judgeval 0.16.6__py3-none-any.whl → 0.16.8__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of judgeval might be problematic. Click here for more details.

Files changed (43) hide show
  1. judgeval/api/api_types.py +1 -2
  2. judgeval/data/judgment_types.py +1 -2
  3. judgeval/tracer/__init__.py +7 -52
  4. judgeval/tracer/llm/config.py +12 -44
  5. judgeval/tracer/llm/constants.py +0 -1
  6. judgeval/tracer/llm/llm_anthropic/config.py +3 -17
  7. judgeval/tracer/llm/llm_anthropic/messages.py +440 -0
  8. judgeval/tracer/llm/llm_anthropic/messages_stream.py +322 -0
  9. judgeval/tracer/llm/llm_anthropic/wrapper.py +40 -621
  10. judgeval/tracer/llm/llm_google/__init__.py +3 -0
  11. judgeval/tracer/llm/llm_google/config.py +3 -21
  12. judgeval/tracer/llm/llm_google/generate_content.py +125 -0
  13. judgeval/tracer/llm/llm_google/wrapper.py +19 -454
  14. judgeval/tracer/llm/llm_openai/beta_chat_completions.py +192 -0
  15. judgeval/tracer/llm/llm_openai/chat_completions.py +437 -0
  16. judgeval/tracer/llm/llm_openai/config.py +3 -29
  17. judgeval/tracer/llm/llm_openai/responses.py +444 -0
  18. judgeval/tracer/llm/llm_openai/wrapper.py +43 -641
  19. judgeval/tracer/llm/llm_together/__init__.py +3 -0
  20. judgeval/tracer/llm/llm_together/chat_completions.py +398 -0
  21. judgeval/tracer/llm/llm_together/config.py +3 -20
  22. judgeval/tracer/llm/llm_together/wrapper.py +34 -485
  23. judgeval/tracer/llm/providers.py +4 -48
  24. judgeval/utils/decorators/dont_throw.py +30 -14
  25. judgeval/utils/wrappers/README.md +3 -0
  26. judgeval/utils/wrappers/__init__.py +15 -0
  27. judgeval/utils/wrappers/immutable_wrap_async.py +74 -0
  28. judgeval/utils/wrappers/immutable_wrap_async_iterator.py +84 -0
  29. judgeval/utils/wrappers/immutable_wrap_sync.py +66 -0
  30. judgeval/utils/wrappers/immutable_wrap_sync_iterator.py +84 -0
  31. judgeval/utils/wrappers/mutable_wrap_async.py +67 -0
  32. judgeval/utils/wrappers/mutable_wrap_sync.py +67 -0
  33. judgeval/utils/wrappers/utils.py +35 -0
  34. judgeval/version.py +1 -1
  35. {judgeval-0.16.6.dist-info → judgeval-0.16.8.dist-info}/METADATA +1 -1
  36. {judgeval-0.16.6.dist-info → judgeval-0.16.8.dist-info}/RECORD +40 -27
  37. judgeval/tracer/llm/llm_groq/config.py +0 -23
  38. judgeval/tracer/llm/llm_groq/wrapper.py +0 -498
  39. judgeval/tracer/local_eval_queue.py +0 -199
  40. /judgeval/{tracer/llm/llm_groq/__init__.py → utils/wrappers/py.typed} +0 -0
  41. {judgeval-0.16.6.dist-info → judgeval-0.16.8.dist-info}/WHEEL +0 -0
  42. {judgeval-0.16.6.dist-info → judgeval-0.16.8.dist-info}/entry_points.txt +0 -0
  43. {judgeval-0.16.6.dist-info → judgeval-0.16.8.dist-info}/licenses/LICENSE.md +0 -0
@@ -0,0 +1,192 @@
1
+ from __future__ import annotations
2
+ from typing import (
3
+ TYPE_CHECKING,
4
+ Any,
5
+ Awaitable,
6
+ Callable,
7
+ Dict,
8
+ ParamSpec,
9
+ TypeVar,
10
+ )
11
+
12
+ from judgeval.tracer.keys import AttributeKeys
13
+ from judgeval.tracer.utils import set_span_attribute
14
+ from judgeval.utils.serialize import safe_serialize
15
+ from judgeval.utils.wrappers import (
16
+ immutable_wrap_sync,
17
+ immutable_wrap_async,
18
+ )
19
+
20
+ if TYPE_CHECKING:
21
+ from judgeval.tracer import Tracer
22
+ from openai import OpenAI, AsyncOpenAI
23
+ from openai.types.chat.parsed_chat_completion import ParsedChatCompletion
24
+
25
+ P = ParamSpec("P")
26
+ T = TypeVar("T")
27
+
28
+
29
+ def wrap_beta_chat_completions_parse_sync(tracer: Tracer, client: OpenAI) -> None:
30
+ original_func = client.beta.chat.completions.parse
31
+ wrapped = _wrap_beta_non_streaming_sync(tracer, original_func)
32
+ setattr(client.beta.chat.completions, "parse", wrapped)
33
+
34
+
35
+ def _wrap_beta_non_streaming_sync(
36
+ tracer: Tracer, original_func: Callable[P, ParsedChatCompletion[T]]
37
+ ) -> Callable[P, ParsedChatCompletion[T]]:
38
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
39
+ ctx["span"] = tracer.get_tracer().start_span(
40
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
41
+ )
42
+ tracer.add_agent_attributes_to_span(ctx["span"])
43
+ set_span_attribute(
44
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
45
+ )
46
+ ctx["model_name"] = kwargs.get("model", "")
47
+ set_span_attribute(
48
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
49
+ )
50
+
51
+ def post_hook(ctx: Dict[str, Any], result: ParsedChatCompletion[T]) -> None:
52
+ span = ctx.get("span")
53
+ if not span:
54
+ return
55
+
56
+ set_span_attribute(
57
+ span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
58
+ )
59
+
60
+ usage_data = result.usage
61
+ if usage_data:
62
+ prompt_tokens = usage_data.prompt_tokens or 0
63
+ completion_tokens = usage_data.completion_tokens or 0
64
+ cache_read = 0
65
+ prompt_tokens_details = usage_data.prompt_tokens_details
66
+ if prompt_tokens_details:
67
+ cache_read = prompt_tokens_details.cached_tokens or 0
68
+
69
+ set_span_attribute(
70
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
71
+ )
72
+ set_span_attribute(
73
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
74
+ )
75
+ set_span_attribute(
76
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
77
+ )
78
+ set_span_attribute(
79
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
80
+ )
81
+ set_span_attribute(
82
+ span,
83
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
84
+ safe_serialize(usage_data),
85
+ )
86
+
87
+ set_span_attribute(
88
+ span,
89
+ AttributeKeys.GEN_AI_RESPONSE_MODEL,
90
+ result.model or ctx["model_name"],
91
+ )
92
+
93
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
94
+ span = ctx.get("span")
95
+ if span:
96
+ span.record_exception(error)
97
+
98
+ def finally_hook(ctx: Dict[str, Any]) -> None:
99
+ span = ctx.get("span")
100
+ if span:
101
+ span.end()
102
+
103
+ return immutable_wrap_sync(
104
+ original_func,
105
+ pre_hook=pre_hook,
106
+ post_hook=post_hook,
107
+ error_hook=error_hook,
108
+ finally_hook=finally_hook,
109
+ )
110
+
111
+
112
+ def wrap_beta_chat_completions_parse_async(tracer: Tracer, client: AsyncOpenAI) -> None:
113
+ original_func = client.beta.chat.completions.parse
114
+ wrapped = _wrap_beta_non_streaming_async(tracer, original_func)
115
+ setattr(client.beta.chat.completions, "parse", wrapped)
116
+
117
+
118
+ def _wrap_beta_non_streaming_async(
119
+ tracer: Tracer, original_func: Callable[P, Awaitable[ParsedChatCompletion[T]]]
120
+ ) -> Callable[P, Awaitable[ParsedChatCompletion[T]]]:
121
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
122
+ ctx["span"] = tracer.get_tracer().start_span(
123
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
124
+ )
125
+ tracer.add_agent_attributes_to_span(ctx["span"])
126
+ set_span_attribute(
127
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
128
+ )
129
+ ctx["model_name"] = kwargs.get("model", "")
130
+ set_span_attribute(
131
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
132
+ )
133
+
134
+ def post_hook(ctx: Dict[str, Any], result: ParsedChatCompletion[T]) -> None:
135
+ span = ctx.get("span")
136
+ if not span:
137
+ return
138
+
139
+ set_span_attribute(
140
+ span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
141
+ )
142
+
143
+ usage_data = result.usage
144
+ if usage_data:
145
+ prompt_tokens = usage_data.prompt_tokens or 0
146
+ completion_tokens = usage_data.completion_tokens or 0
147
+ cache_read = 0
148
+ prompt_tokens_details = usage_data.prompt_tokens_details
149
+ if prompt_tokens_details:
150
+ cache_read = prompt_tokens_details.cached_tokens or 0
151
+
152
+ set_span_attribute(
153
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
154
+ )
155
+ set_span_attribute(
156
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
157
+ )
158
+ set_span_attribute(
159
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
160
+ )
161
+ set_span_attribute(
162
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
163
+ )
164
+ set_span_attribute(
165
+ span,
166
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
167
+ safe_serialize(usage_data),
168
+ )
169
+
170
+ set_span_attribute(
171
+ span,
172
+ AttributeKeys.GEN_AI_RESPONSE_MODEL,
173
+ result.model or ctx["model_name"],
174
+ )
175
+
176
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
177
+ span = ctx.get("span")
178
+ if span:
179
+ span.record_exception(error)
180
+
181
+ def finally_hook(ctx: Dict[str, Any]) -> None:
182
+ span = ctx.get("span")
183
+ if span:
184
+ span.end()
185
+
186
+ return immutable_wrap_async(
187
+ original_func,
188
+ pre_hook=pre_hook,
189
+ post_hook=post_hook,
190
+ error_hook=error_hook,
191
+ finally_hook=finally_hook,
192
+ )
@@ -0,0 +1,437 @@
1
+ from __future__ import annotations
2
+ from typing import (
3
+ TYPE_CHECKING,
4
+ Any,
5
+ Awaitable,
6
+ Callable,
7
+ Dict,
8
+ Iterator,
9
+ AsyncIterator,
10
+ Generator,
11
+ AsyncGenerator,
12
+ ParamSpec,
13
+ TypeVar,
14
+ )
15
+ from packaging import version
16
+
17
+ from judgeval.tracer.keys import AttributeKeys
18
+ from judgeval.tracer.utils import set_span_attribute
19
+ from judgeval.utils.serialize import safe_serialize
20
+ from judgeval.utils.wrappers import (
21
+ immutable_wrap_async,
22
+ immutable_wrap_sync,
23
+ mutable_wrap_sync,
24
+ mutable_wrap_async,
25
+ immutable_wrap_sync_iterator,
26
+ immutable_wrap_async_iterator,
27
+ )
28
+
29
+ if TYPE_CHECKING:
30
+ from judgeval.tracer import Tracer
31
+ from openai import OpenAI, AsyncOpenAI
32
+ from openai.types.chat import ChatCompletion, ChatCompletionChunk
33
+
34
+ P = ParamSpec("P")
35
+ T = TypeVar("T")
36
+
37
+
38
+ def _supports_stream_options() -> bool:
39
+ try:
40
+ import openai
41
+
42
+ return version.parse(openai.__version__) >= version.parse("1.26.0")
43
+ except Exception:
44
+ return False
45
+
46
+
47
+ def wrap_chat_completions_create_sync(tracer: Tracer, client: OpenAI) -> None:
48
+ original_func = client.chat.completions.create
49
+
50
+ def dispatcher(*args: Any, **kwargs: Any) -> Any:
51
+ if kwargs.get("stream", False):
52
+ return _wrap_streaming_sync(tracer, original_func)(*args, **kwargs)
53
+ return _wrap_non_streaming_sync(tracer, original_func)(*args, **kwargs)
54
+
55
+ setattr(client.chat.completions, "create", dispatcher)
56
+
57
+
58
+ def _wrap_non_streaming_sync(
59
+ tracer: Tracer, original_func: Callable[..., ChatCompletion]
60
+ ) -> Callable[..., ChatCompletion]:
61
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
62
+ ctx["span"] = tracer.get_tracer().start_span(
63
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
64
+ )
65
+ tracer.add_agent_attributes_to_span(ctx["span"])
66
+ set_span_attribute(
67
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
68
+ )
69
+ ctx["model_name"] = kwargs.get("model", "")
70
+ set_span_attribute(
71
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
72
+ )
73
+
74
+ def post_hook(ctx: Dict[str, Any], result: ChatCompletion) -> None:
75
+ span = ctx.get("span")
76
+ if not span:
77
+ return
78
+
79
+ set_span_attribute(
80
+ span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
81
+ )
82
+
83
+ usage_data = result.usage
84
+ if usage_data:
85
+ prompt_tokens = usage_data.prompt_tokens or 0
86
+ completion_tokens = usage_data.completion_tokens or 0
87
+ cache_read = 0
88
+ prompt_tokens_details = usage_data.prompt_tokens_details
89
+ if prompt_tokens_details:
90
+ cache_read = prompt_tokens_details.cached_tokens or 0
91
+
92
+ set_span_attribute(
93
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
94
+ )
95
+ set_span_attribute(
96
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
97
+ )
98
+ set_span_attribute(
99
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
100
+ )
101
+ set_span_attribute(
102
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
103
+ )
104
+ set_span_attribute(
105
+ span,
106
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
107
+ safe_serialize(usage_data),
108
+ )
109
+
110
+ set_span_attribute(
111
+ span,
112
+ AttributeKeys.GEN_AI_RESPONSE_MODEL,
113
+ result.model or ctx["model_name"],
114
+ )
115
+
116
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
117
+ span = ctx.get("span")
118
+ if span:
119
+ span.record_exception(error)
120
+
121
+ def finally_hook(ctx: Dict[str, Any]) -> None:
122
+ span = ctx.get("span")
123
+ if span:
124
+ span.end()
125
+
126
+ return immutable_wrap_sync(
127
+ original_func,
128
+ pre_hook=pre_hook,
129
+ post_hook=post_hook,
130
+ error_hook=error_hook,
131
+ finally_hook=finally_hook,
132
+ )
133
+
134
+
135
+ def _wrap_streaming_sync(
136
+ tracer: Tracer, original_func: Callable[..., Iterator[ChatCompletionChunk]]
137
+ ) -> Callable[..., Iterator[ChatCompletionChunk]]:
138
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
139
+ ctx["span"] = tracer.get_tracer().start_span(
140
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
141
+ )
142
+ tracer.add_agent_attributes_to_span(ctx["span"])
143
+ set_span_attribute(
144
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
145
+ )
146
+ ctx["model_name"] = kwargs.get("model", "")
147
+ set_span_attribute(
148
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
149
+ )
150
+ ctx["accumulated_content"] = ""
151
+
152
+ def mutate_kwargs_hook(ctx: Dict[str, Any], kwargs: Any) -> Any:
153
+ if "stream_options" not in kwargs and _supports_stream_options():
154
+ modified_kwargs = dict(kwargs)
155
+ modified_kwargs["stream_options"] = {"include_usage": True}
156
+ return modified_kwargs
157
+ return kwargs
158
+
159
+ def mutate_hook(
160
+ ctx: Dict[str, Any], result: Iterator[ChatCompletionChunk]
161
+ ) -> Iterator[ChatCompletionChunk]:
162
+ def traced_generator() -> Generator[ChatCompletionChunk, None, None]:
163
+ for chunk in result:
164
+ yield chunk
165
+
166
+ def yield_hook(inner_ctx: Dict[str, Any], chunk: ChatCompletionChunk) -> None:
167
+ span = ctx.get("span")
168
+ if not span:
169
+ return
170
+
171
+ if chunk.choices and len(chunk.choices) > 0:
172
+ delta = chunk.choices[0].delta
173
+ if delta and delta.content:
174
+ ctx["accumulated_content"] = (
175
+ ctx.get("accumulated_content", "") + delta.content
176
+ )
177
+
178
+ if hasattr(chunk, "usage") and chunk.usage:
179
+ prompt_tokens = chunk.usage.prompt_tokens or 0
180
+ completion_tokens = chunk.usage.completion_tokens or 0
181
+ cache_read = 0
182
+ if chunk.usage.prompt_tokens_details:
183
+ cache_read = chunk.usage.prompt_tokens_details.cached_tokens or 0
184
+
185
+ set_span_attribute(
186
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
187
+ )
188
+ set_span_attribute(
189
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
190
+ )
191
+ set_span_attribute(
192
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
193
+ )
194
+ set_span_attribute(
195
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
196
+ )
197
+ set_span_attribute(
198
+ span,
199
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
200
+ safe_serialize(chunk.usage),
201
+ )
202
+
203
+ def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
204
+ span = ctx.get("span")
205
+ if span:
206
+ accumulated = ctx.get("accumulated_content", "")
207
+ set_span_attribute(span, AttributeKeys.GEN_AI_COMPLETION, accumulated)
208
+
209
+ def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
210
+ span = ctx.get("span")
211
+ if span:
212
+ span.record_exception(error)
213
+
214
+ def finally_hook_inner(inner_ctx: Dict[str, Any]) -> None:
215
+ span = ctx.get("span")
216
+ if span:
217
+ span.end()
218
+
219
+ wrapped_generator = immutable_wrap_sync_iterator(
220
+ traced_generator,
221
+ yield_hook=yield_hook,
222
+ post_hook=post_hook_inner,
223
+ error_hook=error_hook_inner,
224
+ finally_hook=finally_hook_inner,
225
+ )
226
+
227
+ return wrapped_generator()
228
+
229
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
230
+ span = ctx.get("span")
231
+ if span:
232
+ span.record_exception(error)
233
+
234
+ return mutable_wrap_sync(
235
+ original_func,
236
+ pre_hook=pre_hook,
237
+ mutate_kwargs_hook=mutate_kwargs_hook,
238
+ mutate_hook=mutate_hook,
239
+ error_hook=error_hook,
240
+ )
241
+
242
+
243
+ def wrap_chat_completions_create_async(tracer: Tracer, client: AsyncOpenAI) -> None:
244
+ original_func = client.chat.completions.create
245
+
246
+ async def dispatcher(*args: Any, **kwargs: Any) -> Any:
247
+ if kwargs.get("stream", False):
248
+ return await _wrap_streaming_async(tracer, original_func)(*args, **kwargs)
249
+ return await _wrap_non_streaming_async(tracer, original_func)(*args, **kwargs)
250
+
251
+ setattr(client.chat.completions, "create", dispatcher)
252
+
253
+
254
+ def _wrap_non_streaming_async(
255
+ tracer: Tracer, original_func: Callable[..., Awaitable[ChatCompletion]]
256
+ ) -> Callable[..., Awaitable[ChatCompletion]]:
257
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
258
+ ctx["span"] = tracer.get_tracer().start_span(
259
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
260
+ )
261
+ tracer.add_agent_attributes_to_span(ctx["span"])
262
+ set_span_attribute(
263
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
264
+ )
265
+ ctx["model_name"] = kwargs.get("model", "")
266
+ set_span_attribute(
267
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
268
+ )
269
+
270
+ def post_hook(ctx: Dict[str, Any], result: ChatCompletion) -> None:
271
+ span = ctx.get("span")
272
+ if not span:
273
+ return
274
+
275
+ set_span_attribute(
276
+ span, AttributeKeys.GEN_AI_COMPLETION, safe_serialize(result)
277
+ )
278
+
279
+ usage_data = result.usage
280
+ if usage_data:
281
+ prompt_tokens = usage_data.prompt_tokens or 0
282
+ completion_tokens = usage_data.completion_tokens or 0
283
+ cache_read = 0
284
+ prompt_tokens_details = usage_data.prompt_tokens_details
285
+ if prompt_tokens_details:
286
+ cache_read = prompt_tokens_details.cached_tokens or 0
287
+
288
+ set_span_attribute(
289
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
290
+ )
291
+ set_span_attribute(
292
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
293
+ )
294
+ set_span_attribute(
295
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
296
+ )
297
+ set_span_attribute(
298
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
299
+ )
300
+ set_span_attribute(
301
+ span,
302
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
303
+ safe_serialize(usage_data),
304
+ )
305
+
306
+ set_span_attribute(
307
+ span,
308
+ AttributeKeys.GEN_AI_RESPONSE_MODEL,
309
+ result.model or ctx["model_name"],
310
+ )
311
+
312
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
313
+ span = ctx.get("span")
314
+ if span:
315
+ span.record_exception(error)
316
+
317
+ def finally_hook(ctx: Dict[str, Any]) -> None:
318
+ span = ctx.get("span")
319
+ if span:
320
+ span.end()
321
+
322
+ return immutable_wrap_async(
323
+ original_func,
324
+ pre_hook=pre_hook,
325
+ post_hook=post_hook,
326
+ error_hook=error_hook,
327
+ finally_hook=finally_hook,
328
+ )
329
+
330
+
331
+ def _wrap_streaming_async(
332
+ tracer: Tracer,
333
+ original_func: Callable[..., Awaitable[AsyncIterator[ChatCompletionChunk]]],
334
+ ) -> Callable[..., Awaitable[AsyncIterator[ChatCompletionChunk]]]:
335
+ def pre_hook(ctx: Dict[str, Any], *args: Any, **kwargs: Any) -> None:
336
+ ctx["span"] = tracer.get_tracer().start_span(
337
+ "OPENAI_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
338
+ )
339
+ tracer.add_agent_attributes_to_span(ctx["span"])
340
+ set_span_attribute(
341
+ ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
342
+ )
343
+ ctx["model_name"] = kwargs.get("model", "")
344
+ set_span_attribute(
345
+ ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, ctx["model_name"]
346
+ )
347
+ ctx["accumulated_content"] = ""
348
+
349
+ def mutate_kwargs_hook(ctx: Dict[str, Any], kwargs: Any) -> Any:
350
+ if "stream_options" not in kwargs and _supports_stream_options():
351
+ modified_kwargs = dict(kwargs)
352
+ modified_kwargs["stream_options"] = {"include_usage": True}
353
+ return modified_kwargs
354
+ return kwargs
355
+
356
+ def mutate_hook(
357
+ ctx: Dict[str, Any], result: AsyncIterator[ChatCompletionChunk]
358
+ ) -> AsyncIterator[ChatCompletionChunk]:
359
+ async def traced_generator() -> AsyncGenerator[ChatCompletionChunk, None]:
360
+ async for chunk in result:
361
+ yield chunk
362
+
363
+ def yield_hook(inner_ctx: Dict[str, Any], chunk: ChatCompletionChunk) -> None:
364
+ span = ctx.get("span")
365
+ if not span:
366
+ return
367
+
368
+ if chunk.choices and len(chunk.choices) > 0:
369
+ delta = chunk.choices[0].delta
370
+ if delta and delta.content:
371
+ ctx["accumulated_content"] = (
372
+ ctx.get("accumulated_content", "") + delta.content
373
+ )
374
+
375
+ if hasattr(chunk, "usage") and chunk.usage:
376
+ prompt_tokens = chunk.usage.prompt_tokens or 0
377
+ completion_tokens = chunk.usage.completion_tokens or 0
378
+ cache_read = 0
379
+ if chunk.usage.prompt_tokens_details:
380
+ cache_read = chunk.usage.prompt_tokens_details.cached_tokens or 0
381
+
382
+ set_span_attribute(
383
+ span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
384
+ )
385
+ set_span_attribute(
386
+ span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
387
+ )
388
+ set_span_attribute(
389
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, cache_read
390
+ )
391
+ set_span_attribute(
392
+ span, AttributeKeys.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, 0
393
+ )
394
+ set_span_attribute(
395
+ span,
396
+ AttributeKeys.JUDGMENT_USAGE_METADATA,
397
+ safe_serialize(chunk.usage),
398
+ )
399
+
400
+ def post_hook_inner(inner_ctx: Dict[str, Any]) -> None:
401
+ span = ctx.get("span")
402
+ if span:
403
+ accumulated = ctx.get("accumulated_content", "")
404
+ set_span_attribute(span, AttributeKeys.GEN_AI_COMPLETION, accumulated)
405
+
406
+ def error_hook_inner(inner_ctx: Dict[str, Any], error: Exception) -> None:
407
+ span = ctx.get("span")
408
+ if span:
409
+ span.record_exception(error)
410
+
411
+ def finally_hook_inner(inner_ctx: Dict[str, Any]) -> None:
412
+ span = ctx.get("span")
413
+ if span:
414
+ span.end()
415
+
416
+ wrapped_generator = immutable_wrap_async_iterator(
417
+ traced_generator,
418
+ yield_hook=yield_hook,
419
+ post_hook=post_hook_inner,
420
+ error_hook=error_hook_inner,
421
+ finally_hook=finally_hook_inner,
422
+ )
423
+
424
+ return wrapped_generator()
425
+
426
+ def error_hook(ctx: Dict[str, Any], error: Exception) -> None:
427
+ span = ctx.get("span")
428
+ if span:
429
+ span.record_exception(error)
430
+
431
+ return mutable_wrap_async(
432
+ original_func,
433
+ pre_hook=pre_hook,
434
+ mutate_kwargs_hook=mutate_kwargs_hook,
435
+ mutate_hook=mutate_hook,
436
+ error_hook=error_hook,
437
+ )
@@ -1,32 +1,6 @@
1
1
  from __future__ import annotations
2
+ import importlib.util
2
3
 
3
- HAS_OPENAI = False
4
- openai_OpenAI = None
5
- openai_AsyncOpenAI = None
6
- openai_ChatCompletion = None
7
- openai_Response = None
8
- openai_ParsedChatCompletion = None
4
+ HAS_OPENAI = importlib.util.find_spec("openai") is not None
9
5
 
10
- try:
11
- from openai import OpenAI, AsyncOpenAI
12
- from openai.types.chat.chat_completion import ChatCompletion
13
- from openai.types.responses.response import Response
14
- from openai.types.chat import ParsedChatCompletion
15
-
16
- openai_OpenAI = OpenAI
17
- openai_AsyncOpenAI = AsyncOpenAI
18
- openai_ChatCompletion = ChatCompletion
19
- openai_Response = Response
20
- openai_ParsedChatCompletion = ParsedChatCompletion
21
- HAS_OPENAI = True
22
- except ImportError:
23
- pass
24
-
25
- __all__ = [
26
- "HAS_OPENAI",
27
- "openai_OpenAI",
28
- "openai_AsyncOpenAI",
29
- "openai_ChatCompletion",
30
- "openai_Response",
31
- "openai_ParsedChatCompletion",
32
- ]
6
+ __all__ = ["HAS_OPENAI"]