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

1""" 

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

3""" 

4 

5from __future__ import annotations 

6 

7from dataclasses import dataclass, field 

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 messages: list[Message] 

47 tools: list[dict] | None = None 

48 model_type: str = "openai" 

49 

50 

51class 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 self.session_id: str = "" 

59 self._current_task: str = "" 

60 self._step_count: int = 0 

61 self._plan: str = "" 

62 

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

64 self.session_id = session_id 

65 self._current_task = task 

66 self._step_count = 0 

67 self._plan = "" 

68 self._messages = [] 

69 if self.system_prompt: 

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

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

72 

73 def build_context( 

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

75 ) -> AgentContext: 

76 self._step_count += 1 

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

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

79 

80 def update_plan(self, new_plan: str): 

81 self._plan = new_plan 

82 self._messages.append( 

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

84 ) 

85 

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

87 for r in results: 

88 self._messages.append( 

89 Message( 

90 role="tool", 

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

92 tool_call_id=r.call_id, 

93 ) 

94 ) 

95 

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

97 self._messages.append( 

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

99 ) 

100 

101 def add_user_message(self, content: str): 

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

103 

104 def estimate_context_usage(self) -> float: 

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

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

107 return 0.0 

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

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

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

111 

112 @property 

113 def current_task(self) -> str: 

114 return self._current_task 

115 

116 @property 

117 def step_count(self) -> int: 

118 return self._step_count 

119 

120 @property 

121 def plan(self) -> str: 

122 return self._plan 

123 

124 @property 

125 def message_count(self) -> int: 

126 return len(self._messages) 

127 

128 @property 

129 def estimated_tokens(self) -> int: 

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

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

132 return total_chars // 4