88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
|
"""
|
|||
|
|
Agent 会话上下文管理：维护消息历史、状态追踪。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import uuid
|
|||
|
|
from typing import Any, Dict, List, Optional
|
|||
|
|
|
|||
|
|
|
|||
|
|
class AgentContext:
|
|||
|
|
"""
|
|||
|
|
Agent 会话上下文：
|
|||
|
|
|
|||
|
|
- 消息历史（messages 列表，OpenAI 格式）
|
|||
|
|
- 会话元信息（session_id, user_id 等）
|
|||
|
|
- 执行追踪（iteration 计数, 工具调用统计）
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
system_prompt: str = "你是一个有用的AI助手。",
|
|||
|
|
user_id: Optional[str] = None,
|
|||
|
|
session_id: Optional[str] = None,
|
|||
|
|
):
|
|||
|
|
self.session_id = session_id or str(uuid.uuid4())
|
|||
|
|
self.user_id = user_id
|
|||
|
|
self._messages: List[Dict[str, Any]] = []
|
|||
|
|
self._system_prompt = system_prompt
|
|||
|
|
# 执行状态
|
|||
|
|
self.iteration = 0
|
|||
|
|
self.tool_calls_made = 0
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def messages(self) -> List[Dict[str, Any]]:
|
|||
|
|
"""获取完整消息列表（含 system prompt）。"""
|
|||
|
|
if self._system_prompt:
|
|||
|
|
# 确保 system prompt 始终在第一条
|
|||
|
|
has_system = (
|
|||
|
|
len(self._messages) > 0
|
|||
|
|
and self._messages[0].get("role") == "system"
|
|||
|
|
)
|
|||
|
|
if not has_system:
|
|||
|
|
return [
|
|||
|
|
{"role": "system", "content": self._system_prompt},
|
|||
|
|
*self._messages,
|
|||
|
|
]
|
|||
|
|
return self._messages
|
|||
|
|
|
|||
|
|
def add_user_message(self, content: str) -> None:
|
|||
|
|
"""添加用户消息。"""
|
|||
|
|
self._messages.append({"role": "user", "content": content})
|
|||
|
|
|
|||
|
|
def add_assistant_message(
|
|||
|
|
self,
|
|||
|
|
content: str,
|
|||
|
|
tool_calls: Optional[List[Dict[str, Any]]] = None,
|
|||
|
|
reasoning_content: Optional[str] = None,
|
|||
|
|
) -> None:
|
|||
|
|
"""添加助手回复。"""
|
|||
|
|
msg: Dict[str, Any] = {"role": "assistant", "content": content or ""}
|
|||
|
|
if tool_calls:
|
|||
|
|
msg["tool_calls"] = tool_calls
|
|||
|
|
if reasoning_content:
|
|||
|
|
msg["reasoning_content"] = reasoning_content
|
|||
|
|
self._messages.append(msg)
|
|||
|
|
|
|||
|
|
def add_tool_result(
|
|||
|
|
self, tool_call_id: str, tool_name: str, result: str
|
|||
|
|
) -> None:
|
|||
|
|
"""添加工具执行结果。"""
|
|||
|
|
self._messages.append({
|
|||
|
|
"role": "tool",
|
|||
|
|
"tool_call_id": tool_call_id,
|
|||
|
|
"content": result,
|
|||
|
|
"name": tool_name,
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
def set_system_prompt(self, prompt: str) -> None:
|
|||
|
|
"""更新 system prompt（仅在未发送过消息时有效）。"""
|
|||
|
|
if not self._messages:
|
|||
|
|
self._system_prompt = prompt
|
|||
|
|
|
|||
|
|
def reset(self) -> None:
|
|||
|
|
"""重置上下文（保留 system prompt 和 session_id）。"""
|
|||
|
|
self._messages = []
|
|||
|
|
self.iteration = 0
|
|||
|
|
self.tool_calls_made = 0
|