Added tool calling base

1 files changed, 69 insertions, 17 deletions

diff --git a/agent.py b/agent.py
index 51661fd..816f976 100644
--- a/agent.py
+++ b/agent.py
@@ -1,45 +1,97 @@
-from collections.abc import Iterator
+import json
 
 from context import ContextBuilder
 from identity import AgentIdentity, UserProfile
 from llm import LLM
+from tools import ToolRegistry
+
 
 class Agent:
     def __init__(
         self,
         llm: LLM,
         agent_identity: AgentIdentity,
-        user_profile: UserProfile
+        user_profile: UserProfile,
+        tools: ToolRegistry,
     ) -> None:
         self.llm = llm
+        self.tools = tools
 
         self.context = ContextBuilder(
             agent_identity,
             user_profile,
         )
 
-        self.messages: list[dict[str, str]] = [
+        self.messages: list[dict[str, object]] = [
             {
                 "role": "system",
                 "content": self.context.build_system_prompt(),
             }
         ]
 
-    def chat(self, user_input: str) -> Iterator[str]:
-        self.messages.append({
-            "role": "user",
-            "content": user_input,
-        })
+    def chat(self, user_input: str) -> str:
+        self.messages.append(
+            {
+                "role": "user",
+                "content": user_input,
+            }
+        )
+
+        for _ in range(5):
+            response = self.llm.chat(
+                self.messages,
+                tools=self.tools.schemas(),
+            )
+
+            if response.tool_calls:
+                print(f"\n[tool call: {response.tool_calls}]\n")
+
+            if not response.tool_calls:
+                answer = response.content or ""
+
+                self.messages.append(
+                    {
+                        "role": "assistant",
+                        "content": answer,
+                    }
+                )
+
+                return answer
 
-        answer_parts: list[str] = []
+            self.messages.append(
+                {
+                    "role": "assistant",
+                    "content": response.content,
+                    "tool_calls": [
+                        {
+                            "id": call.id,
+                            "type": "function",
+                            "function": {
+                                "name": call.name,
+                                "arguments": json.dumps(call.arguments),
+                            },
+                        }
+                        for call in response.tool_calls
+                    ],
+                }
+            )
 
-        for chunk in self.llm.chat(self.messages):
-            yield chunk
-            answer_parts.append(chunk)
+            for call in response.tool_calls:
+                tool = self.tools.get(call.name)
 
-        answer = "".join(answer_parts)
+                if tool is None:
+                    result = f"Unknown tool: {call.name}"
+                else:
+                    try:
+                        result = tool.execute(call.arguments).content
+                    except Exception as error:
+                        result = f"Tool error: {error}"
 
-        self.messages.append({
-            "role": "assistant",
-            "content": answer,
-        })
+                self.messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": call.id,
+                        "content": result,
+                    }
+                )
+        raise RuntimeError("Tool call limit exceeded")