Coverage for pydantic_ai_slim/pydantic_ai/models/instrumented.py: 96.51%

1from __future__ import annotations (empty)

2 

3import json (empty)

4from collections.abc import AsyncIterator, Iterator, Mapping (empty)

5from contextlib import asynccontextmanager, contextmanager (empty)

6from dataclasses import dataclass, field (empty)

7from typing import Any, Callable, Literal (empty)

8from urllib.parse import urlparse (empty)

9 

10from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider (empty)

11from opentelemetry.trace import Span, Tracer, TracerProvider, get_tracer_provider (empty)

12from opentelemetry.util.types import AttributeValue (empty)

13from pydantic import TypeAdapter (empty)

14 

15from ..messages import ( (empty)

16 ModelMessage, 

17 ModelRequest, 

18 ModelResponse, 

19) 

20from ..settings import ModelSettings (empty)

21from ..usage import Usage (empty)

22from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse (empty)

23from .wrapper import WrapperModel (empty)

24 

25MODEL_SETTING_ATTRIBUTES: tuple[ (empty)

26 Literal[ 

27 'max_tokens', 

28 'top_p', 

29 'seed', 

30 'temperature', 

31 'presence_penalty', 

32 'frequency_penalty', 

33 ], 

34 ..., 

35] = ( 

36 'max_tokens', 

37 'top_p', 

38 'seed', 

39 'temperature', 

40 'presence_penalty', 

41 'frequency_penalty', 

42) 

43 

44ANY_ADAPTER = TypeAdapter[Any](Any) (empty)

45 

46 

47@dataclass(init=False) (empty)

48class InstrumentationSettings: (empty)

49 """Options for instrumenting models and agents with OpenTelemetry. 

50 

51 Used in: 

52 

53 - `Agent(instrument=...)` 

54 - [`Agent.instrument_all()`][pydantic_ai.agent.Agent.instrument_all] 

55 - [`InstrumentedModel`][pydantic_ai.models.instrumented.InstrumentedModel] 

56 

57 See the [Debugging and Monitoring guide](https://ai.pydantic.dev/logfire/) for more info. 

58 """ 

59 

60 tracer: Tracer = field(repr=False) (empty)

61 event_logger: EventLogger = field(repr=False) (empty)

62 event_mode: Literal['attributes', 'logs'] = 'attributes' (empty)

63 

64 def __init__( (empty)

65 self, 

66 *, 

67 event_mode: Literal['attributes', 'logs'] = 'attributes', 

68 tracer_provider: TracerProvider | None = None, 

69 event_logger_provider: EventLoggerProvider | None = None, 

70 ): 

71 """Create instrumentation options. 

72 

73 Args: 

74 event_mode: The mode for emitting events. If `'attributes'`, events are attached to the span as attributes. 

75 If `'logs'`, events are emitted as OpenTelemetry log-based events. 

76 tracer_provider: The OpenTelemetry tracer provider to use. 

77 If not provided, the global tracer provider is used. 

78 Calling `logfire.configure()` sets the global tracer provider, so most users don't need this. 

79 event_logger_provider: The OpenTelemetry event logger provider to use. 

80 If not provided, the global event logger provider is used. 

81 Calling `logfire.configure()` sets the global event logger provider, so most users don't need this. 

82 This is only used if `event_mode='logs'`. 

83 """ 

84 from pydantic_ai import __version__ (empty)

85 

86 tracer_provider = tracer_provider or get_tracer_provider() (empty)

87 event_logger_provider = event_logger_provider or get_event_logger_provider() (empty)

88 self.tracer = tracer_provider.get_tracer('pydantic-ai', __version__) (empty)

89 self.event_logger = event_logger_provider.get_event_logger('pydantic-ai', __version__) (empty)

90 self.event_mode = event_mode (empty)

91 

92 

93GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system' (empty)

94GEN_AI_REQUEST_MODEL_ATTRIBUTE = 'gen_ai.request.model' (empty)

95 

96 

97@dataclass (empty)

98class InstrumentedModel(WrapperModel): (empty)

99 """Model which wraps another model so that requests are instrumented with OpenTelemetry. 

100 

101 See the [Debugging and Monitoring guide](https://ai.pydantic.dev/logfire/) for more info. 

102 """ 

103 

104 settings: InstrumentationSettings (empty)

105 """Configuration for instrumenting requests.""" (empty)

106 

107 def __init__( (empty)

108 self, 

109 wrapped: Model | KnownModelName, 

110 options: InstrumentationSettings | None = None, 

111 ) -> None: 

112 super().__init__(wrapped) (empty)

113 self.settings = options or InstrumentationSettings() (empty)

114 

115 async def request( (empty)

116 self, 

117 messages: list[ModelMessage], 

118 model_settings: ModelSettings | None, 

119 model_request_parameters: ModelRequestParameters, 

120 ) -> tuple[ModelResponse, Usage]: 

121 with self._instrument(messages, model_settings, model_request_parameters) as finish: (empty)

122 response, usage = await super().request(messages, model_settings, model_request_parameters) (empty)

123 finish(response, usage) (empty)

124 return response, usage (empty)

125 

126 @asynccontextmanager (empty)

127 async def request_stream( (empty)

128 self, 

129 messages: list[ModelMessage], 

130 model_settings: ModelSettings | None, 

131 model_request_parameters: ModelRequestParameters, 

132 ) -> AsyncIterator[StreamedResponse]: 

133 with self._instrument(messages, model_settings, model_request_parameters) as finish: (empty)

134 response_stream: StreamedResponse | None = None (empty)

135 try: (empty)

136 async with super().request_stream( (empty)

137 messages, model_settings, model_request_parameters 

138 ) as response_stream: 

139 yield response_stream (empty)

140 finally: 

141 if response_stream: 141 ↛ exitline 141 didn't jump to the function exit(empty)

142 finish(response_stream.get(), response_stream.usage()) (empty)

143 

144 @contextmanager (empty)

145 def _instrument( (empty)

146 self, 

147 messages: list[ModelMessage], 

148 model_settings: ModelSettings | None, 

149 model_request_parameters: ModelRequestParameters, 

150 ) -> Iterator[Callable[[ModelResponse, Usage], None]]: 

151 operation = 'chat' (empty)

152 span_name = f'{operation} {self.model_name}' (empty)

153 # TODO Missing attributes: 

154 # - error.type: unclear if we should do something here or just always rely on span exceptions 

155 # - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these 

156 attributes: dict[str, AttributeValue] = { (empty)

157 'gen_ai.operation.name': operation, 

158 **self.model_attributes(self.wrapped), 

159 'model_request_parameters': json.dumps(InstrumentedModel.serialize_any(model_request_parameters)), 

160 'logfire.json_schema': json.dumps( 

161 { 

162 'type': 'object', 

163 'properties': {'model_request_parameters': {'type': 'object'}}, 

164 } 

165 ), 

166 } 

167 

168 if model_settings: (empty)

169 for key in MODEL_SETTING_ATTRIBUTES: (empty)

170 if isinstance(value := model_settings.get(key), (float, int)): (empty)

171 attributes[f'gen_ai.request.{key}'] = value (empty)

172 

173 with self.settings.tracer.start_as_current_span(span_name, attributes=attributes) as span: (empty)

174 

175 def finish(response: ModelResponse, usage: Usage): (empty)

176 if not span.is_recording(): (empty)

177 return (empty)

178 

179 events = self.messages_to_otel_events(messages) (empty)

180 for event in self.messages_to_otel_events([response]): (empty)

181 events.append( (empty)

182 Event( 

183 'gen_ai.choice', 

184 body={ 

185 # TODO finish_reason 

186 'index': 0, 

187 'message': event.body, 

188 }, 

189 ) 

190 ) 

191 new_attributes: dict[str, AttributeValue] = usage.opentelemetry_attributes() # type: ignore (empty)

192 attributes.update(getattr(span, 'attributes', {})) (empty)

193 request_model = attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] (empty)

194 new_attributes['gen_ai.response.model'] = response.model_name or request_model (empty)

195 span.set_attributes(new_attributes) (empty)

196 span.update_name(f'{operation} {request_model}') (empty)

197 for event in events: (empty)

198 event.attributes = { (empty)

199 GEN_AI_SYSTEM_ATTRIBUTE: attributes[GEN_AI_SYSTEM_ATTRIBUTE], 

200 **(event.attributes or {}), 

201 } 

202 self._emit_events(span, events) (empty)

203 

204 yield finish (empty)

205 

206 def _emit_events(self, span: Span, events: list[Event]) -> None: (empty)

207 if self.settings.event_mode == 'logs': (empty)

208 for event in events: (empty)

209 self.settings.event_logger.emit(event) (empty)

210 else: 

211 attr_name = 'events' (empty)

212 span.set_attributes( (empty)

213 { 

214 attr_name: json.dumps([self.event_to_dict(event) for event in events]), 

215 'logfire.json_schema': json.dumps( 

216 { 

217 'type': 'object', 

218 'properties': { 

219 attr_name: {'type': 'array'}, 

220 'model_request_parameters': {'type': 'object'}, 

221 }, 

222 } 

223 ), 

224 } 

225 ) 

226 

227 @staticmethod (empty)

228 def model_attributes(model: Model): (empty)

229 attributes: dict[str, AttributeValue] = { (empty)

230 GEN_AI_SYSTEM_ATTRIBUTE: model.system, 

231 GEN_AI_REQUEST_MODEL_ATTRIBUTE: model.model_name, 

232 } 

233 if base_url := model.base_url: (empty)

234 try: (empty)

235 parsed = urlparse(base_url) (empty)

236 except Exception: # pragma: no cover 

237 pass 

238 else: 

239 if parsed.hostname: 239 ↛ 241line 239 didn't jump to line 241 because the condition on line 239 was always true(empty)

240 attributes['server.address'] = parsed.hostname (empty)

241 if parsed.port: 241 ↛ 244line 241 didn't jump to line 244 because the condition on line 241 was always true(empty)

242 attributes['server.port'] = parsed.port (empty)

243 

244 return attributes (empty)

245 

246 @staticmethod (empty)

247 def event_to_dict(event: Event) -> dict[str, Any]: (empty)

248 if not event.body: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true(empty)

249 body = {} 

250 elif isinstance(event.body, Mapping): (empty)

251 body = event.body # type: ignore (empty)

252 else: 

253 body = {'body': event.body} (empty)

254 return {**body, **(event.attributes or {})} (empty)

255 

256 @staticmethod (empty)

257 def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]: (empty)

258 result: list[Event] = [] (empty)

259 for message_index, message in enumerate(messages): (empty)

260 message_events: list[Event] = [] (empty)

261 if isinstance(message, ModelRequest): (empty)

262 for part in message.parts: (empty)

263 if hasattr(part, 'otel_event'): (empty)

264 message_events.append(part.otel_event()) (empty)

265 elif isinstance(message, ModelResponse): 265 ↛ 267line 265 didn't jump to line 267 because the condition on line 265 was always true(empty)

266 message_events = message.otel_events() (empty)

267 for event in message_events: (empty)

268 event.attributes = { (empty)

269 'gen_ai.message.index': message_index, 

270 **(event.attributes or {}), 

271 } 

272 result.extend(message_events) (empty)

273 for event in result: (empty)

274 event.body = InstrumentedModel.serialize_any(event.body) (empty)

275 return result (empty)

276 

277 @staticmethod (empty)

278 def serialize_any(value: Any) -> str: (empty)

279 try: (empty)

280 return ANY_ADAPTER.dump_python(value, mode='json') (empty)

281 except Exception: (empty)

282 try: (empty)

283 return str(value) (empty)

284 except Exception as e: (empty)

285 return f'Unable to serialize: {e}' (empty)