1、重构Agent代码
2、实现WebSocket连接服务
This commit is contained in:
asypost committed
1 parent
2c1bac9bcb
commit
77ff26094d
16 files changed
+871
-322
No files matched your search
Whitespace-only changes.
@@ -0,0 +1,22 @@
|
||||
from openai.types.chat import ChatCompletionMessageParam
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
|
||||
from mini_agent.tools.tool import Tool
|
||||
|
||||
|
||||
class ConversationContext(BaseModel):
|
||||
conversation_id: str = Field(description="会话ID")
|
||||
messages: list[ChatCompletionMessageParam] = Field(description="消息历史", default_factory=list)
|
||||
|
||||
_discovered_tools: list[Tool] = PrivateAttr(default_factory=list)
|
||||
|
||||
discovered_tool_names: list[str] = Field(description="已发现的工具名称", default_factory=list)
|
||||
|
||||
@property
|
||||
def discovered_tools(self):
|
||||
return self._discovered_tools
|
||||
|
||||
def add_tool(self, tool: Tool):
|
||||
if tool.name not in self.discovered_tool_names:
|
||||
self.discovered_tool_names.append(tool.name)
|
||||
self.discovered_tools.append(tool)
|
||||
@@ -0,0 +1,47 @@
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class MessageRole(StrEnum):
|
||||
USER = "user"
|
||||
ASSISTANT = "assistant"
|
||||
TOOL = "tool"
|
||||
|
||||
|
||||
class MessageType(StrEnum):
|
||||
TEXT = "text"
|
||||
THINK = "think"
|
||||
TOOL_CALL = "tool_call"
|
||||
TOOL_RESULT = "tool_result"
|
||||
|
||||
|
||||
class MessageStatus(StrEnum):
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
id: str = Field(default_factory=lambda: str(uuid4()))
|
||||
conversation_id: str
|
||||
run_id: str | None = None
|
||||
|
||||
role: MessageRole
|
||||
type: MessageType
|
||||
|
||||
content: Any = None
|
||||
status: MessageStatus = MessageStatus.COMPLETED
|
||||
|
||||
tool_call_id: str | None = None
|
||||
|
||||
tool_name: str | None = None
|
||||
|
||||
arguments: dict[str, Any] | None = None
|
||||
|
||||
tool_call_result:Any |None = None
|
||||
|
||||
created_at: datetime = Field(default_factory=lambda: datetime.now())
|
||||
@@ -0,0 +1,247 @@
|
||||
from collections.abc import Iterator
|
||||
import json
|
||||
import os
|
||||
from typing import Any
|
||||
from uuid import uuid7
|
||||
|
||||
from openai import OpenAI
|
||||
from openai.types.chat import ChatCompletionMessage, ChatCompletionMessageToolCallUnion
|
||||
from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from mini_agent.agent.context import ConversationContext
|
||||
from mini_agent.agent.message import Message, MessageRole, MessageStatus, MessageType
|
||||
from mini_agent.tools.tavily_search import TavilySearch
|
||||
from mini_agent.tools.tool_repository import ToolRepository
|
||||
|
||||
|
||||
class ReActAgent:
|
||||
SYSTEM_PROMPT = """
|
||||
你是一个能够调用外部工具的智能助手。
|
||||
|
||||
工作规则：
|
||||
|
||||
1. 首先理解用户的问题，判断是否需要外部工具。
|
||||
2. 如果需要工具，调用工具搜索接口搜索合适的工具。
|
||||
3. 根据搜索结果选择合适的工具。
|
||||
4. 工具执行后，根据返回结果判断下一步行动。
|
||||
5. 当已经获得足够的信息，能够完整回答用户问题时，直接给出最终答案。
|
||||
|
||||
注意事项：
|
||||
|
||||
- 不允许编造工具执行结果。
|
||||
- 工具搜索没有结果时，可以尝试调整关键词重新搜索。
|
||||
- 如果工具无法完成任务，明确告知用户。
|
||||
- 不要为了调用工具而调用工具。
|
||||
- 只有在任务已经完成或明确无法继续时才结束回答。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tool_repository: ToolRepository,
|
||||
model: str | None = None,
|
||||
api_key: str | None = None,
|
||||
base_url: str | None = None,
|
||||
timeout: int | None = None,
|
||||
max_steps: int = 10,
|
||||
) -> None:
|
||||
|
||||
self.model = model or os.getenv("LLM_MODEL_ID")
|
||||
self.base_url = base_url or os.getenv("LLM_BASE_URL")
|
||||
self.api_key = api_key or os.getenv("LLM_API_KEY")
|
||||
self.timeout = timeout if timeout is not None else int(os.getenv("LLM_TIMEOUT", "60"))
|
||||
|
||||
if not all([self.model, self.base_url, self.api_key]):
|
||||
raise ValueError("模型ID、API密钥和服务地址必须被提供或在.env文件中定义。")
|
||||
|
||||
self.client = OpenAI(
|
||||
api_key=self.api_key,
|
||||
base_url=self.base_url,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
|
||||
self.tool_repos = tool_repository
|
||||
self.max_steps = max_steps
|
||||
web_search_tool = TavilySearch()
|
||||
web_search_tool.config(api_key=os.getenv("TAVILY_KEY", ""))
|
||||
self.tool_repos.register(web_search_tool)
|
||||
|
||||
def serialize_json(self, data: Any) -> str:
|
||||
if isinstance(data, BaseModel):
|
||||
return data.model_dump_json()
|
||||
return json.dumps(data, ensure_ascii=False, default=str)
|
||||
|
||||
def _initialize_context(self, context: ConversationContext, question: str) -> None:
|
||||
"""
|
||||
初始化会话context
|
||||
"""
|
||||
if len(context.messages) == 0:
|
||||
context.messages.append({"role": "system", "content": self.SYSTEM_PROMPT})
|
||||
if len(context.discovered_tools) == 0:
|
||||
context.add_tool(self.tool_repos)
|
||||
context.messages.append({"role": "user", "content": question})
|
||||
|
||||
def _build_message_param(self, message: ChatCompletionMessage) -> ChatCompletionMessageParam:
|
||||
"""
|
||||
将ChatCompletionMessage转换成ChatCompletionMessageParam
|
||||
"""
|
||||
message_param: ChatCompletionMessageParam = {
|
||||
"role": message.role,
|
||||
"content": message.content,
|
||||
}
|
||||
|
||||
if message.tool_calls:
|
||||
adapter = TypeAdapter(list[ChatCompletionMessageToolCallUnion])
|
||||
message_param["tool_calls"] = adapter.dump_python(message.tool_calls, exclude_none=True)
|
||||
return message_param
|
||||
|
||||
def _parse_tool_call(
|
||||
self, tool_call: ChatCompletionMessageToolCallUnion
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
解析Tool call,返回(tool_name, arguments)
|
||||
"""
|
||||
if tool_call.type != "function" and tool_call.type != "custom":
|
||||
raise ValueError(f"不支持的工具调用类型{tool_call.type}")
|
||||
if tool_call.type == "function":
|
||||
tool_name = tool_call.function.name
|
||||
raw_arguments = tool_call.function.arguments
|
||||
else:
|
||||
tool_name = tool_call.custom.name
|
||||
raw_arguments = tool_call.custom.input
|
||||
arguments = json.loads(raw_arguments)
|
||||
if not isinstance(arguments, dict):
|
||||
raise ValueError("工具参数必须是 JSON 对象")
|
||||
return (tool_name, arguments)
|
||||
|
||||
def _execute_tool_call(
|
||||
self, tool_name: str, arguments: dict[str, Any], context: ConversationContext
|
||||
) -> Any:
|
||||
"""
|
||||
执行tool call
|
||||
"""
|
||||
observation = self.tool_repos.call_tool(tool_name, arguments)
|
||||
# 更新当前可用工具
|
||||
if tool_name == self.tool_repos.name:
|
||||
available_tools = []
|
||||
for item in observation:
|
||||
if item.name not in context.discovered_tool_names:
|
||||
context.add_tool(item)
|
||||
available_tools.append(self.tool_repos.tool_to_openai(item))
|
||||
observation = available_tools
|
||||
return observation
|
||||
|
||||
def run(self, question: str, context: ConversationContext) -> Iterator[Message]:
|
||||
self._initialize_context(context, question)
|
||||
run_id = str(uuid7())
|
||||
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
role=MessageRole.USER,
|
||||
run_id=run_id,
|
||||
status=MessageStatus.RUNNING,
|
||||
type=MessageType.TEXT,
|
||||
content=question,
|
||||
)
|
||||
for _ in range(self.max_steps):
|
||||
tools = [self.tool_repos.tool_to_openai(tool) for tool in context.discovered_tools]
|
||||
|
||||
assert self.model is not None
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=context.messages,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
|
||||
# 保存模型返回消息
|
||||
message_param = self._build_message_param(message)
|
||||
context.messages.append(message_param)
|
||||
|
||||
# 模型没有调用工具，说明它选择直接回答
|
||||
if not message.tool_calls:
|
||||
if message.content:
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
content=message.content,
|
||||
status=MessageStatus.COMPLETED,
|
||||
type=MessageType.TEXT,
|
||||
)
|
||||
return
|
||||
elif reasoning_content := getattr(message, "reasoning_content", None):
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
type=MessageType.THINK,
|
||||
content=reasoning_content,
|
||||
status=MessageStatus.RUNNING,
|
||||
)
|
||||
else:
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
type=MessageType.TEXT,
|
||||
status=MessageStatus.FAILED,
|
||||
content="模型没有返回有效内容",
|
||||
)
|
||||
return
|
||||
|
||||
if message.tool_calls:
|
||||
for tool_call in message.tool_calls:
|
||||
tool_name = None
|
||||
arguments = {}
|
||||
tool_call_status = MessageStatus.RUNNING
|
||||
try:
|
||||
tool_name, arguments = self._parse_tool_call(tool_call)
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
status=MessageStatus.RUNNING,
|
||||
type=MessageType.TOOL_CALL,
|
||||
content=f"call tool:{tool_name}({arguments})",
|
||||
tool_name=tool_name,
|
||||
arguments=arguments,
|
||||
)
|
||||
|
||||
observation = self._execute_tool_call(tool_name, arguments, context)
|
||||
|
||||
except Exception as e:
|
||||
tool_call_status = MessageStatus.FAILED
|
||||
observation = {"error": str(e)}
|
||||
|
||||
# 统一记录工具执行结果
|
||||
context.messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": self.serialize_json(observation),
|
||||
}
|
||||
)
|
||||
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.TOOL,
|
||||
type=MessageType.TOOL_RESULT,
|
||||
status=tool_call_status,
|
||||
content=f"call {tool_name or 'unknown'} result:{self.serialize_json(observation)}",
|
||||
tool_name=tool_name,
|
||||
arguments=arguments,
|
||||
tool_call_result=observation,
|
||||
)
|
||||
|
||||
yield Message(
|
||||
conversation_id=context.conversation_id,
|
||||
run_id=run_id,
|
||||
role=MessageRole.ASSISTANT,
|
||||
status=MessageStatus.FAILED,
|
||||
type=MessageType.TEXT,
|
||||
content=f"已达到最大执行步数 ({self.max_steps})，任务未完成",
|
||||
)
|
||||
Reference in new issue
Block a user