Added tool calling base

1 files changed, 67 insertions, 22 deletions

diff --git a/llm.py b/llm.py
index 524f173..4dde6bd 100644
--- a/llm.py
+++ b/llm.py
@@ -1,42 +1,87 @@
+from dataclasses import dataclass
 import json
 from collections.abc import Iterator
 
 import requests
 
+
+@dataclass
+class ToolCall:
+    id: str
+    name: str
+    arguments: dict[str, object]
+
+
+@dataclass
+class LLMResponse:
+    content: str | None
+    tool_calls: list[ToolCall]
+
+
 class LLM:
     def __init__(self, url: str) -> None:
         self.url = url
 
-    def chat(self, messages: list[dict[str, str]], max_tokens: int = 256) -> Iterator[str]:
+    def chat(
+        self,
+        messages: list[dict[str, object]],
+        tools: list[dict[str, object]] | None = None,
+        max_tokens: int = 256,
+    ) -> LLMResponse:
+        payload: dict[str, object] = {
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": 0.7,
+        }
+
+        if tools is not None:
+            payload["tools"] = tools
+
         response = requests.post(
             self.url,
-            json = {
-                "messages": messages,
-                "max_tokens": max_tokens,
-                "temperature": 0.7,
-                "stream": True,
-            },
-            stream=True,
+            json=payload,  # type: ignore[arg-type]
         )
         response.raise_for_status()
 
-        for line in response.iter_lines():
-            if not line:
-                continue
-            
-            if not line.startswith(b"data: "):
-                continue
+        data = response.json()
+        message = data["choices"][0]["message"]
+
+        content = message.get("content")
 
-            data = line[6:]
+        tool_calls: list[ToolCall] = []
 
-            if data == b"[DONE]":
-                break
+        for call in message.get("tool_calls", []):
+            function = call["function"]
 
-            chunk = json.loads(data)
+            tool_calls.append(
+                ToolCall(
+                    id=call["id"],
+                    name=function["name"],
+                    arguments=json.loads(function["arguments"]),
+                )
+            )
+
+        return LLMResponse(
+            content=content,
+            tool_calls=tool_calls,
+        )
 
-            delta = chunk["choices"][0]["delta"]
-            content = delta.get("content")
 
-            if content:
-                yield content
+if __name__ == "__main__":
+    from tools import ToolRegistry, HelloTool
+
+    registry = ToolRegistry()
+    registry.register(HelloTool())
+    llm = LLM("http://localhost:8080/v1/chat/completions")
+
+    response = llm.chat(
+        messages=[
+            {
+                "role": "user",
+                "content": "Say hello to Joseph using the hello tool.",
+            },
+        ],
+        tools=registry.schemas(),
+    )
 
+    print(response)