Coverage for src / agent / trace_logger.py: 100%

1# Copyright 2025-2026 Microsoft Corporation 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Trace-level logging for LLM request/response capture. 

16 

17Provides structured JSON logging of LLM interactions with token usage, 

18timing, and optional message content for offline analysis and optimization. 

19""" 

20 

21import json 

22import logging 

23from datetime import datetime 

24from pathlib import Path 

25from typing import Any 

26 

27logger = logging.getLogger(__name__) 

28 

29 

30class TraceLogger: 

31 """Logger for capturing detailed LLM request/response traces.""" 

32 

33 def __init__(self, trace_file: Path, include_messages: bool = False): 

34 """Initialize trace logger. 

35 

36 Args: 

37 trace_file: Path to trace log file 

38 include_messages: Whether to include full message content in traces 

39 """ 

40 self.trace_file = trace_file 

41 self.include_messages = include_messages 

42 self._ensure_trace_file() 

43 

44 def _ensure_trace_file(self) -> None: 

45 """Ensure trace log file and directory exist.""" 

46 self.trace_file.parent.mkdir(parents=True, exist_ok=True) 

47 if not self.trace_file.exists(): 

48 self.trace_file.touch() 

49 logger.debug(f"Created trace log file: {self.trace_file}") 

50 

51 def log_interaction( 

52 self, 

53 *, 

54 request_id: str, 

55 messages: list[dict[str, Any]] | None = None, 

56 response_content: str | None = None, 

57 model: str | None = None, 

58 input_tokens: int | None = None, 

59 output_tokens: int | None = None, 

60 total_tokens: int | None = None, 

61 latency_ms: float | None = None, 

62 provider: str | None = None, 

63 error: str | None = None, 

64 ) -> None: 

65 """Log a complete LLM interaction with request and response data. 

66 

67 Args: 

68 request_id: Unique identifier for this request 

69 messages: Request messages (logged only if include_messages=True) 

70 response_content: LLM response text 

71 model: Model name used 

72 input_tokens: Number of input tokens 

73 output_tokens: Number of output tokens 

74 total_tokens: Total tokens used 

75 latency_ms: Response latency in milliseconds 

76 provider: LLM provider name 

77 error: Error message if request failed 

78 """ 

79 trace_entry: dict[str, Any] = { 

80 "timestamp": datetime.now().isoformat(), 

81 "request_id": request_id, 

82 "model": model, 

83 "provider": provider, 

84 } 

85 

86 # Add message data if enabled 

87 if self.include_messages and messages: 

88 trace_entry["message_count"] = len(messages) 

89 trace_entry["messages"] = messages 

90 elif messages: 

91 # Include count but not content 

92 trace_entry["message_count"] = len(messages) 

93 

94 # Add response data 

95 if response_content is not None: 

96 if self.include_messages: 

97 trace_entry["response"] = response_content 

98 else: 

99 # Include length but not content 

100 trace_entry["response_length"] = len(response_content) 

101 

102 # Add token usage 

103 if input_tokens is not None or output_tokens is not None or total_tokens is not None: 

104 trace_entry["tokens"] = {} 

105 if input_tokens is not None: 

106 trace_entry["tokens"]["input"] = input_tokens 

107 if output_tokens is not None: 

108 trace_entry["tokens"]["output"] = output_tokens 

109 if total_tokens is not None: 

110 trace_entry["tokens"]["total"] = total_tokens 

111 

112 # Add timing 

113 if latency_ms is not None: 

114 trace_entry["latency_ms"] = round(latency_ms, 2) 

115 

116 # Add error if present 

117 if error: 

118 trace_entry["error"] = error 

119 

120 # Write to trace log file (append mode) 

121 try: 

122 with open(self.trace_file, "a") as f: 

123 json.dump(trace_entry, f) 

124 f.write("\n") # Write each entry on its own line 

125 except Exception as e: 

126 logger.error(f"Failed to write trace log: {e}") 

127 

128 def log_request( 

129 self, 

130 *, 

131 request_id: str, 

132 messages: list[dict[str, Any]], 

133 model: str | None = None, 

134 provider: str | None = None, 

135 system_instructions: str | None = None, 

136 tools_summary: dict[str, Any] | None = None, 

137 ) -> None: 

138 """Log LLM request data. 

139 

140 Args: 

141 request_id: Unique identifier for this request 

142 messages: Request messages 

143 model: Model name 

144 provider: Provider name 

145 system_instructions: System prompt sent to LLM (if include_messages=True) 

146 tools_summary: Tool definitions summary (if include_messages=True) 

147 """ 

148 trace_entry: dict[str, Any] = { 

149 "timestamp": datetime.now().isoformat(), 

150 "request_id": request_id, 

151 "type": "request", 

152 "model": model, 

153 "provider": provider, 

154 } 

155 

156 if self.include_messages: 

157 trace_entry["message_count"] = len(messages) 

158 trace_entry["messages"] = messages 

159 

160 # Add system instructions if provided 

161 if system_instructions: 

162 trace_entry["system_instructions"] = system_instructions 

163 trace_entry["system_instructions_length"] = len(system_instructions) 

164 trace_entry["system_instructions_tokens_est"] = len(system_instructions) // 4 

165 

166 # Add tools summary if provided 

167 if tools_summary: 

168 trace_entry["tools"] = tools_summary 

169 else: 

170 trace_entry["message_count"] = len(messages) 

171 

172 try: 

173 with open(self.trace_file, "a") as f: 

174 json.dump(trace_entry, f, default=str, ensure_ascii=False) 

175 f.write("\n") 

176 except Exception as e: 

177 logger.error(f"Failed to write trace log: {e}") 

178 

179 def log_response( 

180 self, 

181 *, 

182 request_id: str, 

183 response_content: str, 

184 model: str | None = None, 

185 input_tokens: int | None = None, 

186 output_tokens: int | None = None, 

187 total_tokens: int | None = None, 

188 latency_ms: float | None = None, 

189 error: str | None = None, 

190 ) -> None: 

191 """Log LLM response data with token usage. 

192 

193 Args: 

194 request_id: Unique identifier matching the request 

195 response_content: LLM response text 

196 model: Model name 

197 input_tokens: Number of input tokens 

198 output_tokens: Number of output tokens 

199 total_tokens: Total tokens used 

200 latency_ms: Response latency in milliseconds 

201 error: Error message if request failed 

202 """ 

203 trace_entry: dict[str, Any] = { 

204 "timestamp": datetime.now().isoformat(), 

205 "request_id": request_id, 

206 "type": "response", 

207 "model": model, 

208 } 

209 

210 if self.include_messages: 

211 trace_entry["response"] = response_content 

212 else: 

213 trace_entry["response_length"] = len(response_content) 

214 

215 if input_tokens is not None or output_tokens is not None or total_tokens is not None: 

216 trace_entry["tokens"] = {} 

217 if input_tokens is not None: 

218 trace_entry["tokens"]["input"] = input_tokens 

219 if output_tokens is not None: 

220 trace_entry["tokens"]["output"] = output_tokens 

221 if total_tokens is not None: 

222 trace_entry["tokens"]["total"] = total_tokens 

223 

224 if latency_ms is not None: 

225 trace_entry["latency_ms"] = round(latency_ms, 2) 

226 

227 if error: 

228 trace_entry["error"] = error 

229 

230 try: 

231 with open(self.trace_file, "a") as f: 

232 json.dump(trace_entry, f, default=str, ensure_ascii=False) 

233 f.write("\n") 

234 except Exception as e: 

235 logger.error(f"Failed to write trace log: {e}")