283 lines
9.1 KiB
Python
283 lines
9.1 KiB
Python
"""
|
|
LiteLLM Agent 核心模块
|
|
|
|
基于LiteLLM框架的Agent实现，支持A2A协议
|
|
"""
|
|
import asyncio
|
|
import json
|
|
import uuid
|
|
from typing import AsyncGenerator, Optional, Dict, Any, List
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
|
|
import httpx
|
|
import structlog
|
|
|
|
from config import LiteLLMConfig, AgentConfig, get_config
|
|
|
|
# 配置日志
|
|
logger = structlog.get_logger()
|
|
|
|
|
|
@dataclass
|
|
class Message:
|
|
"""消息数据结构"""
|
|
role: str # user, assistant, system
|
|
content: str
|
|
message_id: str = field(default_factory=lambda: uuid.uuid4().hex)
|
|
timestamp: datetime = field(default_factory=datetime.now)
|
|
metadata: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class Conversation:
|
|
"""对话上下文"""
|
|
conversation_id: str = field(default_factory=lambda: uuid.uuid4().hex)
|
|
messages: List[Message] = field(default_factory=list)
|
|
created_at: datetime = field(default_factory=datetime.now)
|
|
|
|
def add_message(self, role: str, content: str) -> Message:
|
|
"""添加消息到对话"""
|
|
msg = Message(role=role, content=content)
|
|
self.messages.append(msg)
|
|
return msg
|
|
|
|
def to_openai_format(self) -> List[Dict[str, str]]:
|
|
"""转换为OpenAI格式的消息列表"""
|
|
return [{"role": m.role, "content": m.content} for m in self.messages]
|
|
|
|
|
|
class LiteLLMAgent:
|
|
"""
|
|
基于LiteLLM的Agent实现
|
|
|
|
支持功能:
|
|
- 多轮对话
|
|
- 流式响应
|
|
- 工具调用
|
|
- A2A协议兼容
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
litellm_config: Optional[LiteLLMConfig] = None,
|
|
agent_config: Optional[AgentConfig] = None
|
|
):
|
|
"""
|
|
初始化Agent
|
|
|
|
Args:
|
|
api_key: LiteLLM API密钥（可选，优先使用，否则从环境变量获取）
|
|
model: 模型名称（可选，优先使用，否则从环境变量获取）
|
|
litellm_config: LiteLLM配置对象
|
|
agent_config: Agent配置对象
|
|
"""
|
|
if litellm_config:
|
|
self.llm_config = litellm_config
|
|
else:
|
|
self.llm_config = LiteLLMConfig(api_key=api_key, model=model)
|
|
|
|
if agent_config:
|
|
self.agent_config = agent_config
|
|
else:
|
|
self.agent_config = AgentConfig()
|
|
|
|
# 验证配置
|
|
self.llm_config.validate()
|
|
|
|
# HTTP客户端
|
|
self._client: Optional[httpx.AsyncClient] = None
|
|
|
|
# 对话管理
|
|
self.conversations: Dict[str, Conversation] = {}
|
|
|
|
# 工具注册
|
|
self.tools: Dict[str, callable] = {}
|
|
|
|
logger.info(
|
|
"Agent初始化完成",
|
|
agent_name=self.agent_config.name,
|
|
model=self.llm_config.model,
|
|
base_url=self.llm_config.base_url
|
|
)
|
|
|
|
async def _get_client(self) -> httpx.AsyncClient:
|
|
"""获取或创建HTTP客户端"""
|
|
if self._client is None or self._client.is_closed:
|
|
self._client = httpx.AsyncClient(
|
|
timeout=httpx.Timeout(self.llm_config.timeout),
|
|
headers={
|
|
"Authorization": f"Bearer {self.llm_config.api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
)
|
|
return self._client
|
|
|
|
async def close(self):
|
|
"""关闭资源"""
|
|
if self._client and not self._client.is_closed:
|
|
await self._client.aclose()
|
|
|
|
def register_tool(self, name: str, func: callable, description: str = ""):
|
|
"""注册工具函数"""
|
|
self.tools[name] = {
|
|
"function": func,
|
|
"description": description
|
|
}
|
|
logger.info(f"注册工具: {name}")
|
|
|
|
def get_or_create_conversation(self, conversation_id: Optional[str] = None) -> Conversation:
|
|
"""获取或创建对话"""
|
|
if conversation_id and conversation_id in self.conversations:
|
|
return self.conversations[conversation_id]
|
|
|
|
conv = Conversation(conversation_id=conversation_id or uuid.uuid4().hex)
|
|
# 添加系统提示
|
|
conv.add_message("system", self.agent_config.system_prompt)
|
|
self.conversations[conv.conversation_id] = conv
|
|
return conv
|
|
|
|
async def chat(
|
|
self,
|
|
message: str,
|
|
conversation_id: Optional[str] = None,
|
|
stream: bool = False
|
|
) -> str | AsyncGenerator[str, None]:
|
|
"""
|
|
发送消息并获取回复
|
|
|
|
Args:
|
|
message: 用户消息
|
|
conversation_id: 对话ID（用于多轮对话）
|
|
stream: 是否流式响应
|
|
|
|
Returns:
|
|
如果stream=False，返回完整回复字符串
|
|
如果stream=True，返回异步生成器
|
|
"""
|
|
# 获取对话上下文
|
|
conversation = self.get_or_create_conversation(conversation_id)
|
|
conversation.add_message("user", message)
|
|
|
|
if stream:
|
|
return self._stream_chat(conversation)
|
|
else:
|
|
return await self._simple_chat(conversation)
|
|
|
|
async def _simple_chat(self, conversation: Conversation) -> str:
|
|
"""非流式对话"""
|
|
client = await self._get_client()
|
|
|
|
request_body = {
|
|
"model": self.llm_config.model,
|
|
"messages": conversation.to_openai_format(),
|
|
"temperature": self.llm_config.temperature,
|
|
"max_tokens": self.llm_config.max_tokens
|
|
}
|
|
|
|
logger.debug("发送请求", endpoint=self.llm_config.chat_endpoint)
|
|
|
|
try:
|
|
response = await client.post(
|
|
self.llm_config.chat_endpoint,
|
|
json=request_body
|
|
)
|
|
response.raise_for_status()
|
|
|
|
result = response.json()
|
|
assistant_message = result["choices"][0]["message"]["content"]
|
|
|
|
# 保存助手回复到对话
|
|
conversation.add_message("assistant", assistant_message)
|
|
|
|
logger.info("收到回复", length=len(assistant_message))
|
|
return assistant_message
|
|
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error("HTTP错误", status_code=e.response.status_code, detail=e.response.text)
|
|
raise
|
|
except Exception as e:
|
|
logger.error("请求失败", error=str(e))
|
|
raise
|
|
|
|
async def _stream_chat(self, conversation: Conversation) -> AsyncGenerator[str, None]:
|
|
"""流式对话"""
|
|
client = await self._get_client()
|
|
|
|
request_body = {
|
|
"model": self.llm_config.model,
|
|
"messages": conversation.to_openai_format(),
|
|
"temperature": self.llm_config.temperature,
|
|
"max_tokens": self.llm_config.max_tokens,
|
|
"stream": True
|
|
}
|
|
|
|
full_response = ""
|
|
|
|
try:
|
|
async with client.stream(
|
|
"POST",
|
|
self.llm_config.chat_endpoint,
|
|
json=request_body
|
|
) as response:
|
|
response.raise_for_status()
|
|
|
|
async for line in response.aiter_lines():
|
|
if line.startswith("data: "):
|
|
data = line[6:]
|
|
if data == "[DONE]":
|
|
break
|
|
|
|
try:
|
|
chunk = json.loads(data)
|
|
delta = chunk.get("choices", [{}])[0].get("delta", {})
|
|
content = delta.get("content", "")
|
|
if content:
|
|
full_response += content
|
|
yield content
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
# 保存完整回复到对话
|
|
conversation.add_message("assistant", full_response)
|
|
|
|
except Exception as e:
|
|
logger.error("流式请求失败", error=str(e))
|
|
raise
|
|
|
|
async def invoke_tool(self, tool_name: str, **kwargs) -> Any:
|
|
"""调用注册的工具"""
|
|
if tool_name not in self.tools:
|
|
raise ValueError(f"未找到工具: {tool_name}")
|
|
|
|
tool = self.tools[tool_name]
|
|
func = tool["function"]
|
|
|
|
logger.info(f"调用工具: {tool_name}", kwargs=kwargs)
|
|
|
|
if asyncio.iscoroutinefunction(func):
|
|
return await func(**kwargs)
|
|
else:
|
|
return func(**kwargs)
|
|
|
|
|
|
# 示例工具函数
|
|
def tool_get_current_time() -> str:
|
|
"""获取当前时间"""
|
|
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def tool_calculate(expression: str) -> str:
|
|
"""计算数学表达式"""
|
|
try:
|
|
# 安全的数学计算
|
|
allowed_chars = set("0123456789+-*/.() ")
|
|
if not all(c in allowed_chars for c in expression):
|
|
return "错误: 不支持的字符"
|
|
result = eval(expression)
|
|
return str(result)
|
|
except Exception as e:
|
|
return f"计算错误: {str(e)}"
|