diff options
| -rw-r--r-- | agent.py | 86 | ||||
| -rw-r--r-- | client.py | 20 | ||||
| -rw-r--r-- | context.py | 1 | ||||
| -rw-r--r-- | identity.py | 4 | ||||
| -rw-r--r-- | llm.py | 89 | ||||
| -rw-r--r-- | tools.py | 149 |
6 files changed, 302 insertions, 47 deletions
@@ -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") @@ -1,16 +1,24 @@ from agent import Agent from identity import load_agent_identity, load_user_profile from llm import LLM +from tools import CalculatorTool, HelloTool, ToolRegistry -llm = LLM("http://localhost:8080/v1/chat/completions",) +llm = LLM( + "http://localhost:8080/v1/chat/completions", +) agent_identity = load_agent_identity("agent.toml") user_profile = load_user_profile("user.toml") +tools = ToolRegistry() +tools.register(HelloTool()) +tools.register(CalculatorTool()) + agent = Agent( llm, agent_identity, user_profile, + tools, ) while True: @@ -18,10 +26,6 @@ while True: if user_input.lower() in {"exit", "quit"}: break - - print("Agent: ", end="", flush=True) - - for chunk in agent.chat(user_input): - print(chunk, end="", flush=True) - - print() + + answer = agent.chat(user_input) + print(f"Agent: {answer}") @@ -1,5 +1,6 @@ from identity import AgentIdentity, UserProfile + class ContextBuilder: def __init__( self, diff --git a/identity.py b/identity.py index 9b4faa0..cbcd6ad 100644 --- a/identity.py +++ b/identity.py @@ -1,6 +1,7 @@ from dataclasses import dataclass import tomllib + @dataclass class AgentIdentity: name: str @@ -8,11 +9,13 @@ class AgentIdentity: values: dict[str, bool] quirks: dict[str, str] + @dataclass class UserProfile: name: str preferences: dict[str, str] + def load_agent_identity(path: str) -> AgentIdentity: with open(path, "rb") as file: data = tomllib.load(file) @@ -24,6 +27,7 @@ def load_agent_identity(path: str) -> AgentIdentity: quirks=data["quirks"], ) + def load_user_profile(path: str) -> UserProfile: with open(path, "rb") as file: data = tomllib.load(file) @@ -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) diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..64238e7 --- /dev/null +++ b/tools.py @@ -0,0 +1,149 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass + + +@dataclass +class ToolResult: + content: str + + +class Tool(ABC): + @property + @abstractmethod + def name(self) -> str: ... + + @property + @abstractmethod + def description(self) -> str: ... + + @property + @abstractmethod + def schema(self) -> dict[str, object]: ... + + @abstractmethod + def execute(self, arguments: dict[str, object]) -> ToolResult: ... + + +class ToolRegistry: + def __init__(self) -> None: + self.tools: dict[str, Tool] = {} + + def register(self, tool: Tool) -> None: + self.tools[tool.name] = tool + + def get(self, name: str) -> Tool | None: + return self.tools.get(name) + + def all(self) -> list[Tool]: + return list(self.tools.values()) + + def schemas(self) -> list[dict[str, object]]: + return [ + { + "type": "function", + "function": { + "name": tool.name, + "description": tool.description, + "parameters": tool.schema, + }, + } + for tool in self.tools.values() + ] + + +class HelloTool(Tool): + @property + def name(self) -> str: + return "hello" + + @property + def description(self) -> str: + return "Say hello to someone." + + @property + def schema(self) -> dict[str, object]: + return { + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "The person to greet.", + }, + }, + "required": ["name"], + } + + def execute(self, arguments: dict[str, object]) -> ToolResult: + name = arguments.get("name", "stranger") + return ToolResult(content=f"Hello, {name}!") + +class CalculatorTool(Tool): + @property + def name(self) -> str: + return "calculator" + + @property + def description(self) -> str: + return "Perform basic arithemtic" + + @property + def schema(self) -> dict[str, object]: + return { + "type": "object", + "properties": { + "a": { + "type": "number", + "description": "The first number.", + }, + "b": { + "type": "number", + "description": "The second number.", + }, + "operation": { + "type": "string", + "enum": ["add", "subtract", "multiply", "divide"], + "description": "The arithemtic operation to perform", + }, + }, + "required": ["a", "b", "operation"], + } + + def execute(self, arguments: dict[str, object]) -> ToolResult: + a = arguments["a"] + b = arguments["b"] + operation = arguments["operation"] + + if not isinstance(a, (int, float)): + raise TypeError("a must be a number") + + if not isinstance(b, (int, float)): + raise TypeError("b must be a number") + + if not isinstance(operation, str): + raise TypeError("operation must be a string") + + if operation == "add": + result = a + b + elif operation == "subtract": + result = a - b + elif operation == "multiply": + result = a * b + elif operation == "divide": + if b == 0: + raise ValueError("cannot divide by zero") + result = a / b + else: + raise ValueError(f"unknown operation: {operation}") + + return ToolResult(content=str(result)) + +if __name__ == "__main__": + registry = ToolRegistry() + + registry.register(HelloTool()) + print(registry.schemas()) + tool = registry.get("hello") + + if tool is not None: + result = tool.execute({"name": "Joseph"}) + print(result.content) |