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 

17 name: str 

18 start_time: float = 0.0 

19 end_time: float = 0.0 

20 duration_ms: float = 0.0 

21 model: str = "" 

22 tokens_in: int = 0 

23 tokens_out: int = 0 

24 tool_calls_count: int = 0 

25 error: str | None = None 

26 

27 

28@dataclass 

29class TokenStats: 

30 """Token 使用统计。""" 

31 

32 total_input: int = 0 

33 total_output: int = 0 

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

35 

36 

37@dataclass 

38class ObservabilityReport: 

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

40 

41 session_id: str 

42 total_duration_ms: float = 0.0 

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

44 tokens: TokenStats = field(default_factory=TokenStats) 

45 model_calls: int = 0 

46 tool_calls_total: int = 0 

47 

48 def summary(self) -> str: 

49 lines = [ 

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

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

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

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

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

55 ] 

56 if self.steps: 

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

58 return "\n".join(lines) 

59 

60 

61class Tracer: 

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

63 

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

65 self.session_id = session_id 

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

67 self.token_stats = TokenStats() 

68 self.start_time = time.time() 

69 

70 @classmethod 

71 def noop(cls) -> Tracer: 

72 return NoopTracer() 

73 

74 @contextmanager 

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

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

77 try: 

78 yield trace 

79 except Exception as e: 

80 trace.error = str(e) 

81 raise 

82 finally: 

83 trace.end_time = time.time() 

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

85 self.steps.append(trace) 

86 

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

88 self.token_stats.total_input += input_tokens 

89 self.token_stats.total_output += output_tokens 

90 if model not in self.token_stats.by_model: 

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

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

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

94 

95 def track_tool_call(self): 

96 if self.steps: 

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

98 

99 def report(self) -> ObservabilityReport: 

100 return ObservabilityReport( 

101 session_id=self.session_id, 

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

103 steps=self.steps, 

104 tokens=self.token_stats, 

105 model_calls=len(self.steps), 

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

107 ) 

108 

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

110 return { 

111 "input": self.token_stats.total_input, 

112 "output": self.token_stats.total_output, 

113 } 

114 

115 

116class NoopTracer(Tracer): 

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

118 

119 def __init__(self): 

120 pass 

121 

122 @contextmanager 

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

124 yield StepTrace(name=name) 

125 

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

127 pass 

128 

129 def track_tool_call(self): 

130 pass 

131 

132 def report(self) -> ObservabilityReport: 

133 return ObservabilityReport(session_id="noop") 

134 

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

136 return {}