Coverage for agentos/core/context.py: 56%

1""" 

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

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass 

8from typing import Any 

9 

10 

11@dataclass 

12class ToolCall: 

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

14 

15 name: str 

16 arguments: dict[str, Any] 

17 

18 

19@dataclass 

20class 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 

34class 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 

44class AgentContext: 

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

46 

47 messages: list[Message] 

48 tools: list[dict] | None = None 

49 model_type: str = "openai" 

50 

51 

52class ContextManager: 

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

54 

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

56 self.system_prompt = system_prompt 

57 self.max_history = max_history 

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

59 self.session_id: str = "" 

60 self._current_task: str = "" 

61 self._step_count: int = 0 

62 self._plan: str = "" 

63 

64 async def init_session(self, session_id: str, task: str): 

65 self.session_id = session_id 

66 self._current_task = task 

67 self._step_count = 0 

68 self._plan = "" 

69 self._messages = [] 

70 if self.system_prompt: 

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

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

73 

74 def build_context( 

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

76 ) -> AgentContext: 

77 self._step_count += 1 

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

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

80 

81 def update_plan(self, new_plan: str): 

82 self._plan = new_plan 

83 self._messages.append(Message(role="system", content=f"[计划更新] {new_plan}")) 

84 

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

86 for r in results: 

87 self._messages.append( 

88 Message( 

89 role="tool", 

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

91 tool_call_id=r.call_id, 

92 ) 

93 ) 

94 

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

96 self._messages.append(Message(role="assistant", content=content, tool_calls=tool_calls)) 

97 

98 def add_user_message(self, content: str): 

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

100 

101 def estimate_context_usage(self) -> float: 

102 """估算上下文使用比例 (0.0-1.0)。按 max_history 消息数 vs 最大 Token 估算。""" 

103 if len(self._messages) == 0: 

104 return 0.0 

105 estimated_tokens = sum(len(m.content) for m in self._messages) // 4 

106 max_tokens = self.max_history * 200 # 每消息 ~200 token 粗估 

107 return min(estimated_tokens / max_tokens if max_tokens > 0 else 0.0, 1.0) 

108 

109 @property 

110 def current_task(self) -> str: 

111 return self._current_task 

112 

113 @property 

114 def step_count(self) -> int: 

115 return self._step_count 

116 

117 @property 

118 def plan(self) -> str: 

119 return self._plan 

120 

121 @property 

122 def message_count(self) -> int: 

123 return len(self._messages) 

124 

125 @property 

126 def estimated_tokens(self) -> int: 

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

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

129 return total_chars // 4