Initial commit: Air Agent Framework v2：全双工对话智能体（打断、插话、主动开口）
This commit is contained in:
WpyQwq
+202
@@ -0,0 +1,202 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user