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

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 

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