113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from openai import AsyncOpenAI
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self) -> None:
|
|
settings = get_settings()
|
|
self.model = settings.openrouter_model
|
|
self.client = AsyncOpenAI(
|
|
api_key=settings.openrouter_api_key,
|
|
base_url=settings.openrouter_base_url,
|
|
)
|
|
|
|
async def stream_chat(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None = None,
|
|
) -> AsyncIterator[dict[str, Any]]:
|
|
kwargs: dict[str, Any] = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"stream": True,
|
|
"temperature": 0.7,
|
|
}
|
|
if tools:
|
|
kwargs["tools"] = tools
|
|
|
|
stream = await self.client.chat.completions.create(**kwargs)
|
|
|
|
tool_calls: dict[int, dict[str, Any]] = {}
|
|
|
|
async for chunk in stream:
|
|
if not chunk.choices:
|
|
continue
|
|
|
|
choice = chunk.choices[0]
|
|
delta = choice.delta
|
|
|
|
if delta.content:
|
|
yield {"type": "content", "content": delta.content}
|
|
|
|
if delta.tool_calls:
|
|
for tool_call in delta.tool_calls:
|
|
idx = tool_call.index
|
|
if idx not in tool_calls:
|
|
tool_calls[idx] = {
|
|
"id": tool_call.id or "",
|
|
"type": "function",
|
|
"function": {"name": "", "arguments": ""},
|
|
}
|
|
if tool_call.id:
|
|
tool_calls[idx]["id"] = tool_call.id
|
|
if tool_call.function:
|
|
if tool_call.function.name:
|
|
tool_calls[idx]["function"]["name"] = tool_call.function.name
|
|
if tool_call.function.arguments:
|
|
tool_calls[idx]["function"]["arguments"] += tool_call.function.arguments
|
|
|
|
if choice.finish_reason:
|
|
if tool_calls:
|
|
yield {"type": "tool_calls", "tool_calls": list(tool_calls.values())}
|
|
yield {"type": "done", "finish_reason": choice.finish_reason}
|
|
|
|
async def complete(
|
|
self,
|
|
messages: list[dict[str, Any]],
|
|
tools: list[dict[str, Any]] | None = None,
|
|
) -> dict[str, Any]:
|
|
kwargs: dict[str, Any] = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"temperature": 0.7,
|
|
}
|
|
if tools:
|
|
kwargs["tools"] = tools
|
|
|
|
response = await self.client.chat.completions.create(**kwargs)
|
|
message = response.choices[0].message
|
|
|
|
result: dict[str, Any] = {
|
|
"content": message.content or "",
|
|
"tool_calls": [],
|
|
}
|
|
|
|
if message.tool_calls:
|
|
result["tool_calls"] = [
|
|
{
|
|
"id": tc.id,
|
|
"type": "function",
|
|
"function": {
|
|
"name": tc.function.name,
|
|
"arguments": tc.function.arguments,
|
|
},
|
|
}
|
|
for tc in message.tool_calls
|
|
]
|
|
|
|
return result
|
|
|
|
@staticmethod
|
|
def parse_tool_arguments(arguments: str) -> dict[str, Any]:
|
|
if not arguments:
|
|
return {}
|
|
try:
|
|
return json.loads(arguments)
|
|
except json.JSONDecodeError:
|
|
return {}
|