1"""

2 上下文管理器 — 构建Agent所需的完整上下文。

3"""

4

5 from __future__ import annotations

6

7 from dataclasses import dataclass, field

8 from typing import Any

9

10

11@dataclass

12 class ToolCall:

13 """模型请求的工具调用。"""

14

15 name: str

16 arguments: dict[str, Any]

17

18

19@dataclass

20 class ToolResult:

21 """工具调用的返回结果。"""

22

23 call_id: str

24 output: str | None = None

25 error: str | None = None

26 exit_code: int | None = None

27

28 @property

29 def is_error(self) -> bool:

30 return self.error is not None

31

32

33@dataclass

34 class Message:

35 """对话中的单条消息。"""

36

37 role: str # system | user | assistant | tool

38 content: str

39 tool_calls: list[ToolCall] | None = None

40 tool_call_id: str | None = None

41

42

43@dataclass

44 class AgentContext:

45 """传给模型的完整上下文。"""

46 messages: list[Message]

47 tools: list[dict] | None = None

48 model_type: str = "openai"

49

50

51 class ContextManager:

52 """管理Agent会话的全部消息历史。"""

53

54 def __init__(self, system_prompt: str = "", max_history: int = 200):

55 self.system_prompt = system_prompt

56 self.max_history = max_history

57 self._messages: list[Message] = []

58

59 def init_session(self, session_id: str, task: str):

60 self.session_id = session_id

61 self._messages = []

62 if self.system_prompt:

63 self._messages.append(Message(role="system", content=self.system_prompt))

64 self._messages.append(Message(role="user", content=task))

65

66 def build_context(

67 self, model_type: str = "openai", tools: list[dict] | None = None

68 ) -> AgentContext:

69 messages = self._messages[-self.max_history:]

70 return AgentContext(messages=messages, tools=tools, model_type=model_type)

71

72 def append_tool_results(self, results: list[ToolResult]):

73 for r in results:

74 self._messages.append(

75 Message(

76 role="tool",

77 content=r.error or r.output or "",

78 tool_call_id=r.call_id,

79 )

80 )

81

82 def add_assistant_message(self, content: str, tool_calls: list[ToolCall] | None = None):

83 self._messages.append(

84 Message(role="assistant", content=content, tool_calls=tool_calls)

85 )

86

87 def add_user_message(self, content: str):

88 self._messages.append(Message(role="user", content=content))

89

90 @property

91 def message_count(self) -> int:

92 return len(self._messages)

93

94 @property

95 def estimated_tokens(self) -> int:

96 """粗略估算当前上下文的总token数（按每4字符≈1token）。"""

97 total_chars = sum(len(m.content) for m in self._messages)

98 return total_chars // 4