diff options
Diffstat (limited to 'llm.py')
| -rw-r--r-- | llm.py | 87 |
1 files changed, 0 insertions, 87 deletions
@@ -1,87 +0,0 @@ -import json -from dataclasses import dataclass - -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, 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.9, - } - - if tools is not None: - payload["tools"] = tools - - response = requests.post( - self.url, - json=payload, # type: ignore[arg-type] - ) - response.raise_for_status() - - data = response.json() - print(f"\n[debug: {data}]\n") - message = data["choices"][0]["message"] - - content = message.get("content") - - tool_calls: list[ToolCall] = [] - - for call in message.get("tool_calls", []): - function = call["function"] - - tool_calls.append( - ToolCall( - id=call["id"], - name=function["name"], - arguments=json.loads(function["arguments"]), - ) - ) - - return LLMResponse( - content=content, - tool_calls=tool_calls, - ) - - -if __name__ == "__main__": - from tools import HelloTool, ToolRegistry - - 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) |
