203 lines
6.4 KiB
Python
203 lines
6.4 KiB
Python
from __future__ import annotations
|
|
import json
|
|
from typing import AsyncIterator, Optional, Callable
|
|
from dataclasses import dataclass, field
|
|
import httpx
|
|
|
|
|
|
@dataclass
|
|
class LLMMessage:
|
|
role: str
|
|
content: str
|
|
tool_calls: list = field(default_factory=list)
|
|
tool_call_id: str = ""
|
|
|
|
|
|
@dataclass
|
|
class ToolDef:
|
|
name: str
|
|
description: str
|
|
parameters: dict
|
|
|
|
|
|
@dataclass
|
|
class LLMResponse:
|
|
content: str
|
|
tool_calls: list = field(default_factory=list)
|
|
finish_reason: str = "stop"
|
|
usage: dict = field(default_factory=dict)
|
|
|
|
|
|
class LLMClient:
|
|
def __init__(self, config):
|
|
self.api_key = config.api_key
|
|
self.base_url = config.base_url.rstrip("/")
|
|
self.model = config.model
|
|
self.max_tokens = config.max_tokens
|
|
self.temperature = config.temperature
|
|
self.timeout = config.timeout
|
|
self._client = httpx.AsyncClient(timeout=config.timeout)
|
|
|
|
def _build_headers(self) -> dict:
|
|
headers = {"Content-Type": "application/json"}
|
|
if self.api_key:
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
return headers
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict],
|
|
tools: Optional[list[ToolDef]] = None,
|
|
stream: bool = False,
|
|
response_format: Optional[dict] = None,
|
|
) -> LLMResponse:
|
|
body = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"max_tokens": self.max_tokens,
|
|
"temperature": self.temperature,
|
|
"stream": stream,
|
|
}
|
|
if tools:
|
|
body["tools"] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": t.name,
|
|
"description": t.description,
|
|
"parameters": t.parameters,
|
|
},
|
|
}
|
|
for t in tools
|
|
]
|
|
if response_format:
|
|
body["response_format"] = response_format
|
|
|
|
if stream:
|
|
return await self._chat_stream(body)
|
|
return await self._chat_sync(body)
|
|
|
|
async def _chat_sync(self, body: dict) -> LLMResponse:
|
|
resp = await self._client.post(
|
|
f"{self.base_url}/chat/completions",
|
|
headers=self._build_headers(),
|
|
json=body,
|
|
)
|
|
resp.raise_for_status()
|
|
data = resp.json()
|
|
choice = data["choices"][0]
|
|
msg = choice["message"]
|
|
return LLMResponse(
|
|
content=msg.get("content", "") or "",
|
|
tool_calls=self._parse_tool_calls(msg.get("tool_calls", [])),
|
|
finish_reason=choice.get("finish_reason", "stop"),
|
|
usage=data.get("usage", {}),
|
|
)
|
|
|
|
async def _chat_stream(self, body: dict) -> LLMResponse:
|
|
content = ""
|
|
tool_calls = {}
|
|
finish_reason = ""
|
|
async with self._client.stream(
|
|
"POST",
|
|
f"{self.base_url}/chat/completions",
|
|
headers=self._build_headers(),
|
|
json=body,
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
async for line in resp.aiter_lines():
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data_str = line[6:].strip()
|
|
if data_str == "[DONE]":
|
|
break
|
|
try:
|
|
data = json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
delta = data["choices"][0].get("delta", {})
|
|
if delta.get("content"):
|
|
content += delta["content"]
|
|
for tc in delta.get("tool_calls", []):
|
|
idx = tc["index"]
|
|
if idx not in tool_calls:
|
|
tool_calls[idx] = {
|
|
"id": tc.get("id", ""),
|
|
"function": {"name": "", "arguments": ""},
|
|
}
|
|
if tc.get("id"):
|
|
tool_calls[idx]["id"] = tc["id"]
|
|
if tc.get("function", {}).get("name"):
|
|
tool_calls[idx]["function"]["name"] += tc["function"]["name"]
|
|
if tc.get("function", {}).get("arguments"):
|
|
tool_calls[idx]["function"]["arguments"] += tc["function"]["arguments"]
|
|
fr = data["choices"][0].get("finish_reason")
|
|
if fr:
|
|
finish_reason = fr
|
|
return LLMResponse(
|
|
content=content,
|
|
tool_calls=list(tool_calls.values()),
|
|
finish_reason=finish_reason,
|
|
)
|
|
|
|
def _parse_tool_calls(self, raw: list) -> list:
|
|
result = []
|
|
for tc in raw:
|
|
result.append({
|
|
"id": tc.get("id", ""),
|
|
"type": "function",
|
|
"function": {
|
|
"name": tc["function"]["name"],
|
|
"arguments": tc["function"]["arguments"],
|
|
},
|
|
})
|
|
return result
|
|
|
|
async def chat_stream_iter(
|
|
self,
|
|
messages: list[dict],
|
|
tools: Optional[list[ToolDef]] = None,
|
|
) -> AsyncIterator[str]:
|
|
body = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
"max_tokens": self.max_tokens,
|
|
"temperature": self.temperature,
|
|
"stream": True,
|
|
}
|
|
if tools:
|
|
body["tools"] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": t.name,
|
|
"description": t.description,
|
|
"parameters": t.parameters,
|
|
},
|
|
}
|
|
for t in tools
|
|
]
|
|
async with self._client.stream(
|
|
"POST",
|
|
f"{self.base_url}/chat/completions",
|
|
headers=self._build_headers(),
|
|
json=body,
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
async for line in resp.aiter_lines():
|
|
if not line.startswith("data: "):
|
|
continue
|
|
data_str = line[6:].strip()
|
|
if data_str == "[DONE]":
|
|
break
|
|
try:
|
|
data = json.loads(data_str)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
delta = data["choices"][0].get("delta", {})
|
|
if delta.get("content"):
|
|
yield delta["content"]
|
|
|
|
async def close(self):
|
|
await self._client.aclose()
|