Coverage for agentos/observability/tracer.py: 58%

1""" 

2全链路追踪 — 每一步可追溯。 

3基因来源: LangSmith + OpenAI Tracing 

4""" 

5 

6from __future__ import annotations 

7 

8import time 

9from contextlib import contextmanager 

10from dataclasses import dataclass, field 

11 

12 

13@dataclass 

14class StepTrace: 

15 """单步追踪记录。""" 

16 name: str 

17 start_time: float = 0.0 

18 end_time: float = 0.0 

19 duration_ms: float = 0.0 

20 model: str = "" 

21 tokens_in: int = 0 

22 tokens_out: int = 0 

23 tool_calls_count: int = 0 

24 error: str | None = None 

25 

26 

27@dataclass 

28class TokenStats: 

29 """Token 使用统计。""" 

30 total_input: int = 0 

31 total_output: int = 0 

32 by_model: dict[str, dict[str, int]] = field(default_factory=dict) 

33 

34 

35@dataclass 

36class ObservabilityReport: 

37 """可观测性报告。""" 

38 session_id: str 

39 total_duration_ms: float = 0.0 

40 steps: list[StepTrace] = field(default_factory=list) 

41 tokens: TokenStats = field(default_factory=TokenStats) 

42 model_calls: int = 0 

43 tool_calls_total: int = 0 

44 

45 def summary(self) -> str: 

46 lines = [ 

47 f"Session: {self.session_id}", 

48 f"Duration: {self.total_duration_ms:.0f}ms", 

49 f"Model calls: {self.model_calls}", 

50 f"Tool calls: {self.tool_calls_total}", 

51 f"Tokens (in/out): {self.tokens.total_input}/{self.tokens.total_output}", 

52 ] 

53 if self.steps: 

54 lines.append(f"Steps: {', '.join(s.name for s in self.steps)}") 

55 return "\n".join(lines) 

56 

57 

58class Tracer: 

59 """全链路追踪器。每步记录耗时、token消耗、工具调用。""" 

60 

61 def __init__(self, session_id: str = ""): 

62 self.session_id = session_id 

63 self.steps: list[StepTrace] = [] 

64 self.token_stats = TokenStats() 

65 self.start_time = time.time() 

66 

67 @classmethod 

68 def noop(cls) -> "Tracer": 

69 return NoopTracer() 

70 

71 @contextmanager 

72 def step(self, name: str, model: str = ""): 

73 trace = StepTrace(name=name, model=model, start_time=time.time()) 

74 try: 

75 yield trace 

76 except Exception as e: 

77 trace.error = str(e) 

78 raise 

79 finally: 

80 trace.end_time = time.time() 

81 trace.duration_ms = (trace.end_time - trace.start_time) * 1000 

82 self.steps.append(trace) 

83 

84 def track_tokens(self, model: str, input_tokens: int, output_tokens: int): 

85 self.token_stats.total_input += input_tokens 

86 self.token_stats.total_output += output_tokens 

87 if model not in self.token_stats.by_model: 

88 self.token_stats.by_model[model] = {"input": 0, "output": 0} 

89 self.token_stats.by_model[model]["input"] += input_tokens 

90 self.token_stats.by_model[model]["output"] += output_tokens 

91 

92 def track_tool_call(self): 

93 if self.steps: 

94 self.steps[-1].tool_calls_count += 1 

95 

96 def report(self) -> ObservabilityReport: 

97 return ObservabilityReport( 

98 session_id=self.session_id, 

99 total_duration_ms=(time.time() - self.start_time) * 1000, 

100 steps=self.steps, 

101 tokens=self.token_stats, 

102 model_calls=len(self.steps), 

103 tool_calls_total=sum(s.tool_calls_count for s in self.steps), 

104 ) 

105 

106 def token_summary(self) -> dict[str, int]: 

107 return { 

108 "input": self.token_stats.total_input, 

109 "output": self.token_stats.total_output, 

110 } 

111 

112 

113class NoopTracer(Tracer): 

114 """空追踪器 — 生产环境中关闭追踪时使用。""" 

115 

116 def __init__(self): 

117 pass 

118 

119 @contextmanager 

120 def step(self, name: str, model: str = ""): 

121 yield StepTrace(name=name) 

122 

123 def track_tokens(self, model: str, input_tokens: int, output_tokens: int): 

124 pass 

125 

126 def track_tool_call(self): 

127 pass 

128 

129 def report(self) -> ObservabilityReport: 

130 return ObservabilityReport(session_id="noop") 

131 

132 def token_summary(self) -> dict[str, int]: 

133 return {}