telemetry-dev-google-genai 0.1.0__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.
@@ -0,0 +1,1213 @@
1
+ # Google GenAI accepts both pydantic models and plain dicts; the defensive readers below
2
+ # intentionally traverse dynamic shapes without trusting SDK internals.
3
+ # pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
4
+ # pyright: reportUnknownArgumentType=false, reportAttributeAccessIssue=false
5
+ # pyright: reportUnnecessaryIsInstance=false
6
+
7
+ from __future__ import annotations
8
+
9
+ import contextvars
10
+ import json
11
+ import threading
12
+ import time
13
+ from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence
14
+ from functools import wraps
15
+ from typing import Any, TypeVar, cast
16
+
17
+ import telemetry_dev
18
+ from google.genai import _extra_utils
19
+ from google.genai.models import AsyncModels, Models
20
+
21
+ __version__ = "0.1.0"
22
+
23
+ ProviderResolver = Callable[[object | None], str]
24
+ RequestMapper = Callable[[Any, Any, Any], tuple[str, dict[str, Any]]]
25
+ ResponseMapper = Callable[[Any], dict[str, Any]]
26
+
27
+ _WRAPPED_ATTR = "_telemetry_dev_google_genai_wrapped"
28
+ _ORIGINAL_ATTR = "_telemetry_dev_google_genai_original"
29
+ _ORIGINALS: list[tuple[type[Any], str, Any]] = []
30
+ _installed = False
31
+ _install_lock = threading.Lock()
32
+ _T = TypeVar("_T")
33
+ _AFC_USAGE_STATE: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar(
34
+ "telemetry_dev_google_genai_afc_usage_state", default=None
35
+ )
36
+
37
+ _BUILTIN_TOOL_FIELDS = (
38
+ ("google_search", "googleSearch"),
39
+ ("google_search_retrieval", "googleSearchRetrieval"),
40
+ ("code_execution", "codeExecution"),
41
+ ("url_context", "urlContext"),
42
+ ("computer_use", "computerUse"),
43
+ ("file_search", "fileSearch"),
44
+ ("retrieval", "retrieval"),
45
+ ("google_maps", "googleMaps"),
46
+ ("enterprise_web_search", "enterpriseWebSearch"),
47
+ ("parallel_ai_search", "parallelAiSearch"),
48
+ ("mcp_servers", "mcpServers"),
49
+ )
50
+
51
+
52
+ def _camel_name(name: str) -> str:
53
+ parts = name.split("_")
54
+ return parts[0] + "".join(part.capitalize() for part in parts[1:])
55
+
56
+
57
+ def _field(value: Any, name: str) -> Any:
58
+ camel = _camel_name(name)
59
+ if isinstance(value, Mapping):
60
+ if name in value:
61
+ return value[name]
62
+ return value.get(camel)
63
+ for key in (name, camel):
64
+ if hasattr(value, key):
65
+ return getattr(value, key, None)
66
+ return None
67
+
68
+
69
+ def _sequence_items(value: Any) -> list[Any]:
70
+ if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray):
71
+ return list(cast(Sequence[Any], value))
72
+ return []
73
+
74
+
75
+ def _native(value: Any) -> Any:
76
+ if hasattr(value, "model_dump"):
77
+ try:
78
+ return value.model_dump(mode="json", by_alias=True, exclude_none=True)
79
+ except Exception:
80
+ try:
81
+ return value.model_dump(by_alias=True, exclude_none=True)
82
+ except Exception:
83
+ return str(value)
84
+ if isinstance(value, Mapping):
85
+ return {str(key): _native(item) for key, item in value.items()}
86
+ if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray):
87
+ return [_native(item) for item in value]
88
+ if hasattr(value, "value") and not isinstance(value, str | int | float | bool):
89
+ return _native(value.value)
90
+ return value
91
+
92
+
93
+ def _normalized_content_input(contents: Any) -> list[Any]:
94
+ if contents is None:
95
+ return []
96
+ if isinstance(contents, str):
97
+ return [{"role": "user", "parts": [{"text": contents}]}]
98
+ items = _sequence_items(contents)
99
+ if items:
100
+ if any(
101
+ _string(_field(item, "role")) is not None or _sequence_items(_field(item, "parts"))
102
+ for item in items
103
+ ):
104
+ return [_native(item) for item in items]
105
+ parts: list[Any] = []
106
+ for item in items:
107
+ if isinstance(item, str):
108
+ parts.append({"text": item})
109
+ else:
110
+ parts.append(_native(item))
111
+ return [{"role": "user", "parts": parts}]
112
+ if _string(_field(contents, "role")) is None and not _sequence_items(_field(contents, "parts")):
113
+ return [{"role": "user", "parts": [_native(contents)]}]
114
+ return [_native(contents)]
115
+
116
+
117
+ def _number(value: Any) -> int | float | None:
118
+ if isinstance(value, bool):
119
+ return None
120
+ if isinstance(value, int | float):
121
+ return value
122
+ return None
123
+
124
+
125
+ def _string(value: Any) -> str | None:
126
+ return value if isinstance(value, str) else None
127
+
128
+
129
+ def _enum_value(value: Any) -> str | None:
130
+ if value is None:
131
+ return None
132
+ if hasattr(value, "value") and not isinstance(value, str):
133
+ raw = value.value
134
+ return raw if isinstance(raw, str) else str(raw)
135
+ return _string(value)
136
+
137
+
138
+ def _usage(fields: dict[str, int | float | None]) -> dict[str, int | float] | None:
139
+ usage = {key: value for key, value in fields.items() if value is not None}
140
+ return usage or None
141
+
142
+
143
+ def _sum_usage(
144
+ total: dict[str, int | float] | None, part: dict[str, int | float] | None
145
+ ) -> dict[str, int | float] | None:
146
+ if not total:
147
+ return part
148
+ if not part:
149
+ return total
150
+ merged = dict(total)
151
+ for key, value in part.items():
152
+ merged[key] = merged.get(key, 0) + value
153
+ return merged
154
+
155
+
156
+ def _afc_usage_state(config: Any) -> dict[str, Any] | None:
157
+ try:
158
+ if _extra_utils.should_disable_afc(config):
159
+ return None
160
+ tools = _sequence_items(_field(config, "tools"))
161
+ if not any(
162
+ callable(tool)
163
+ or callable(getattr(tool, "call_tool", None))
164
+ or callable(getattr(tool, "callTool", None))
165
+ for tool in tools
166
+ ):
167
+ return None
168
+ return {}
169
+ except Exception:
170
+ return None
171
+
172
+
173
+ def _record_afc_usage(response: Any) -> Any:
174
+ state = _AFC_USAGE_STATE.get()
175
+ if state is None:
176
+ return response
177
+ try:
178
+ fields = _generate_response_fields(response)
179
+ usage = fields.get("usage")
180
+ if isinstance(usage, dict):
181
+ state["usage"] = _sum_usage(
182
+ cast(dict[str, int | float] | None, state.get("usage")), usage
183
+ )
184
+ attrs = fields.get("attributes")
185
+ if isinstance(attrs, Mapping):
186
+ tool_use = _number(attrs.get("google_genai.usage.tool_use_prompt_tokens"))
187
+ if tool_use is not None:
188
+ state["tool_use_prompt_tokens"] = int(state.get("tool_use_prompt_tokens", 0)) + int(
189
+ tool_use
190
+ )
191
+ except Exception:
192
+ return response
193
+ return response
194
+
195
+
196
+ def _fields_with_afc_usage(fields: dict[str, Any], state: dict[str, Any] | None) -> dict[str, Any]:
197
+ usage = state.get("usage") if state else None
198
+ if isinstance(usage, dict):
199
+ fields["usage"] = usage
200
+ tool_use = _number(state.get("tool_use_prompt_tokens")) if state else None
201
+ if tool_use is not None:
202
+ attrs = fields.get("attributes")
203
+ if not isinstance(attrs, dict):
204
+ attrs = {}
205
+ fields["attributes"] = attrs
206
+ attrs["google_genai.usage.tool_use_prompt_tokens"] = int(tool_use)
207
+ return fields
208
+
209
+
210
+ def _stop_sequences(value: Any) -> list[str] | None:
211
+ if isinstance(value, str):
212
+ return [value]
213
+ strings = [item for item in (_string(part) for part in _sequence_items(value)) if item]
214
+ return strings or None
215
+
216
+
217
+ def _json_attr(value: Any) -> str | None:
218
+ try:
219
+ return json.dumps(_native(value))
220
+ except (TypeError, ValueError):
221
+ return None
222
+
223
+
224
+ def _clean_fields(fields: Mapping[str, Any]) -> dict[str, Any]:
225
+ return {key: value for key, value in fields.items() if value is not None}
226
+
227
+
228
+ def _safe_fields(mapper: ResponseMapper, response: Any) -> dict[str, Any]:
229
+ try:
230
+ return mapper(response)
231
+ except Exception:
232
+ return {}
233
+
234
+
235
+ def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
236
+ ended = False
237
+
238
+ def end(**fields: Any) -> None:
239
+ nonlocal ended
240
+ if ended:
241
+ return
242
+ ended = True
243
+ handle.end(**_clean_fields(fields))
244
+
245
+ return end
246
+
247
+
248
+ def _output_type(config: Any) -> str | None:
249
+ if config is None:
250
+ return None
251
+ mime = _string(_field(config, "response_mime_type"))
252
+ if mime == "application/json":
253
+ return "json"
254
+ if (
255
+ _field(config, "response_schema") is not None
256
+ or _field(config, "response_json_schema") is not None
257
+ ):
258
+ return "json"
259
+ if mime == "text/plain":
260
+ return "text"
261
+ return None
262
+
263
+
264
+ def _tool_definitions(tools: Any) -> str | None:
265
+ items = _sequence_items(tools)
266
+ if not items:
267
+ return None
268
+ definitions: list[Any] = []
269
+ for tool in items:
270
+ if callable(tool):
271
+ tool_name = getattr(tool, "__name__", None)
272
+ if tool_name:
273
+ definitions.append({"type": "function", "name": tool_name})
274
+ continue
275
+ if callable(getattr(tool, "call_tool", None)) or callable(getattr(tool, "callTool", None)):
276
+ tool_name = getattr(tool, "name", None)
277
+ if tool_name:
278
+ definitions.append({"type": "function", "name": tool_name})
279
+ continue
280
+ native = _native(tool)
281
+ tool_dict = native if isinstance(native, Mapping) else {}
282
+ decls = tool_dict.get("function_declarations") or tool_dict.get("functionDeclarations")
283
+ if decls is None:
284
+ decls = _field(tool, "function_declarations")
285
+ for decl in _sequence_items(decls):
286
+ decl_native = _native(decl)
287
+ if not isinstance(decl_native, Mapping):
288
+ continue
289
+ entry: dict[str, Any] = {"type": "function", "name": decl_native.get("name")}
290
+ if decl_native.get("description") is not None:
291
+ entry["description"] = decl_native.get("description")
292
+ parameters = decl_native.get("parameters")
293
+ if parameters is None:
294
+ parameters = decl_native.get("parameters_json_schema")
295
+ if parameters is None:
296
+ parameters = decl_native.get("parametersJsonSchema")
297
+ if parameters is not None:
298
+ entry["parameters"] = parameters
299
+ if entry.get("name"):
300
+ definitions.append(entry)
301
+ for field_name, type_name in _BUILTIN_TOOL_FIELDS:
302
+ if _field(tool, field_name) is not None or (
303
+ isinstance(tool_dict, Mapping) and tool_dict.get(field_name) is not None
304
+ ):
305
+ definitions.append({"type": type_name})
306
+ if not definitions:
307
+ return None
308
+ return _json_attr(definitions)
309
+
310
+
311
+ def _request_attributes(config: Any) -> dict[str, Any]:
312
+ attrs: dict[str, Any] = {}
313
+ candidate_count = _number(_field(config, "candidate_count"))
314
+ if candidate_count is not None:
315
+ attrs["gen_ai.request.choice.count"] = int(candidate_count)
316
+ tools = _tool_definitions(_field(config, "tools"))
317
+ if tools is not None:
318
+ attrs["gen_ai.tool.definitions"] = tools
319
+ for field_name, attr_name in (
320
+ ("tool_config", "google_genai.request.tool_config"),
321
+ ("safety_settings", "google_genai.request.safety_settings"),
322
+ ("thinking_config", "google_genai.request.thinking_config"),
323
+ ("labels", "google_genai.request.labels"),
324
+ ("response_modalities", "google_genai.request.response_modalities"),
325
+ ):
326
+ value = _field(config, field_name)
327
+ if value is None:
328
+ continue
329
+ encoded = _json_attr(value)
330
+ if encoded is not None:
331
+ attrs[attr_name] = encoded
332
+ cached_content = _field(config, "cached_content")
333
+ if cached_content is not None:
334
+ cached = _string(cached_content) or _string(_field(cached_content, "name"))
335
+ if cached is not None:
336
+ attrs["google_genai.request.cached_content"] = cached
337
+ return attrs
338
+
339
+
340
+ def _generate_request_fields(model: Any, contents: Any, config: Any) -> tuple[str, dict[str, Any]]:
341
+ model_name = _string(model) or "unknown"
342
+ fields: dict[str, Any] = {
343
+ "type": "generation",
344
+ "model": _string(model),
345
+ "input": _normalized_content_input(contents),
346
+ }
347
+ system_instruction = _field(config, "system_instruction")
348
+ if system_instruction is not None:
349
+ fields["system_instructions"] = _native(system_instruction)
350
+ for field_name in (
351
+ "temperature",
352
+ "top_p",
353
+ "top_k",
354
+ "seed",
355
+ "presence_penalty",
356
+ "frequency_penalty",
357
+ ):
358
+ value = _number(_field(config, field_name))
359
+ if value is not None:
360
+ fields[field_name] = value
361
+ max_output_tokens = _number(_field(config, "max_output_tokens"))
362
+ if max_output_tokens is not None:
363
+ fields["max_tokens"] = int(max_output_tokens)
364
+ stop_sequences = _stop_sequences(_field(config, "stop_sequences"))
365
+ if stop_sequences is not None:
366
+ fields["stop_sequences"] = stop_sequences
367
+ output_type = _output_type(config)
368
+ if output_type is not None:
369
+ fields["output_type"] = output_type
370
+ attrs = _request_attributes(config)
371
+ if attrs:
372
+ fields["attributes"] = attrs
373
+ return f"chat {model_name}", fields
374
+
375
+
376
+ def _generate_usage(raw: Any, attrs: dict[str, Any]) -> dict[str, int | float] | None:
377
+ tool_use = _number(_field(raw, "tool_use_prompt_token_count"))
378
+ if tool_use is not None:
379
+ attrs["google_genai.usage.tool_use_prompt_tokens"] = int(tool_use)
380
+ return _usage(
381
+ {
382
+ "input_tokens": _number(_field(raw, "prompt_token_count")),
383
+ "output_tokens": _number(_field(raw, "candidates_token_count")),
384
+ "total_tokens": _number(_field(raw, "total_token_count")),
385
+ "cache_read_input_tokens": _number(_field(raw, "cached_content_token_count")),
386
+ "reasoning_output_tokens": _number(_field(raw, "thoughts_token_count")),
387
+ }
388
+ )
389
+
390
+
391
+ def _content_output(content: Any) -> dict[str, Any] | None:
392
+ if content is None:
393
+ return None
394
+ parts = _sequence_items(_field(content, "parts"))
395
+ if not parts:
396
+ return None
397
+ role = _string(_field(content, "role")) or "model"
398
+ return {"role": role, "parts": [_native(part) for part in parts]}
399
+
400
+
401
+ def _candidate_index(candidate: Any, fallback: int) -> int:
402
+ index = _number(_field(candidate, "index"))
403
+ return int(index) if index is not None else fallback
404
+
405
+
406
+ def _record_candidate_metadata(candidate: Any, index: int, attrs: dict[str, Any]) -> None:
407
+ ratings = _field(candidate, "safety_ratings")
408
+ if ratings:
409
+ existing = attrs.get("google_genai.response.safety_ratings")
410
+ entries: list[Any]
411
+ if isinstance(existing, str):
412
+ try:
413
+ parsed = json.loads(existing)
414
+ entries = parsed if isinstance(parsed, list) else []
415
+ except json.JSONDecodeError:
416
+ entries = []
417
+ else:
418
+ entries = []
419
+ entries.append({"candidateIndex": index, "ratings": _native(ratings)})
420
+ encoded = _json_attr(entries)
421
+ if encoded is not None:
422
+ attrs["google_genai.response.safety_ratings"] = encoded
423
+ if index == 0:
424
+ grounding = _field(candidate, "grounding_metadata")
425
+ if grounding is not None:
426
+ encoded = _json_attr(grounding)
427
+ if encoded is not None:
428
+ attrs["google_genai.response.grounding_metadata"] = encoded
429
+ url_context = _field(candidate, "url_context_metadata")
430
+ if url_context is not None:
431
+ encoded = _json_attr(url_context)
432
+ if encoded is not None:
433
+ attrs["google_genai.response.url_context_metadata"] = encoded
434
+
435
+
436
+ def _generate_response_fields(response: Any) -> dict[str, Any]:
437
+ fields: dict[str, Any] = {}
438
+ attrs: dict[str, Any] = {}
439
+ response_model = _string(_field(response, "model_version"))
440
+ if response_model:
441
+ fields["response_model"] = response_model
442
+ response_id = _string(_field(response, "response_id"))
443
+ if response_id:
444
+ fields["response_id"] = response_id
445
+ candidates = _sequence_items(_field(response, "candidates"))
446
+ finish_reasons: list[str] = []
447
+ output: list[dict[str, Any]] = []
448
+ for index, candidate in enumerate(candidates):
449
+ candidate_index = _candidate_index(candidate, index)
450
+ finish = _enum_value(_field(candidate, "finish_reason"))
451
+ if finish is not None:
452
+ finish_reasons.append(finish)
453
+ content_output = _content_output(_field(candidate, "content"))
454
+ if content_output is not None:
455
+ output.append(content_output)
456
+ _record_candidate_metadata(candidate, candidate_index, attrs)
457
+ if finish_reasons:
458
+ fields["finish_reason"] = finish_reasons[0]
459
+ if len(finish_reasons) > 1:
460
+ attrs["gen_ai.response.finish_reasons"] = finish_reasons
461
+ if output:
462
+ fields["output"] = output
463
+ usage = _generate_usage(_field(response, "usage_metadata"), attrs)
464
+ if usage:
465
+ fields["usage"] = usage
466
+ feedback = _field(response, "prompt_feedback")
467
+ block_reason = _enum_value(_field(feedback, "block_reason"))
468
+ if block_reason:
469
+ attrs["google_genai.response.block_reason"] = block_reason
470
+ block_message = _string(_field(feedback, "block_reason_message"))
471
+ if block_message:
472
+ attrs["google_genai.response.block_reason_message"] = block_message
473
+ prompt_ratings = _sequence_items(_field(feedback, "safety_ratings"))
474
+ if prompt_ratings:
475
+ encoded = _json_attr([_native(rating) for rating in prompt_ratings])
476
+ if encoded is not None:
477
+ attrs["google_genai.response.prompt_safety_ratings"] = encoded
478
+ history = _field(response, "automatic_function_calling_history")
479
+ if history:
480
+ native_history = _native(history)
481
+ if native_history:
482
+ fields["input"] = native_history
483
+ attrs["google_genai.automatic_function_calling"] = True
484
+ if attrs:
485
+ fields["attributes"] = attrs
486
+ return fields
487
+
488
+
489
+ def _embed_request_fields(model: Any, contents: Any, config: Any) -> tuple[str, dict[str, Any]]:
490
+ model_name = _string(model) or "unknown"
491
+ fields: dict[str, Any] = {
492
+ "type": "embedding",
493
+ "model": _string(model),
494
+ "input": _native(contents),
495
+ }
496
+ attrs: dict[str, Any] = {}
497
+ task_type = _string(_field(config, "task_type"))
498
+ if task_type:
499
+ attrs["google_genai.request.task_type"] = task_type
500
+ output_dimensionality = _number(_field(config, "output_dimensionality"))
501
+ if output_dimensionality is not None:
502
+ attrs["google_genai.request.output_dimensionality"] = int(output_dimensionality)
503
+ if attrs:
504
+ fields["attributes"] = attrs
505
+ return f"embeddings {model_name}", fields
506
+
507
+
508
+ def _embed_response_fields(response: Any) -> dict[str, Any]:
509
+ fields: dict[str, Any] = {}
510
+ attrs: dict[str, Any] = {}
511
+ embeddings = _sequence_items(_field(response, "embeddings"))
512
+ if embeddings:
513
+ attrs["google_genai.response.embedding_count"] = len(embeddings)
514
+ values = _field(embeddings[0], "values")
515
+ if values is not None:
516
+ attrs["google_genai.response.embedding_dimensions"] = len(list(values))
517
+ total_tokens = 0
518
+ saw_tokens = False
519
+ for embedding in embeddings:
520
+ token_count = _number(_field(_field(embedding, "statistics"), "token_count"))
521
+ if token_count is not None:
522
+ total_tokens += int(token_count)
523
+ saw_tokens = True
524
+ usage = _usage({"input_tokens": total_tokens if saw_tokens else None})
525
+ if usage:
526
+ fields["usage"] = usage
527
+ billable = _number(_field(_field(response, "metadata"), "billable_character_count"))
528
+ if billable is not None:
529
+ attrs["google_genai.usage.billable_characters"] = int(billable)
530
+ if attrs:
531
+ fields["attributes"] = attrs
532
+ return fields
533
+
534
+
535
+ class _CandidateState:
536
+ def __init__(self) -> None:
537
+ self.role = "model"
538
+ self.parts: list[Any] = []
539
+
540
+
541
+ class _StreamState:
542
+ def __init__(self) -> None:
543
+ self.candidates: dict[int, _CandidateState] = {}
544
+ self.finish_reasons: dict[int, str] = {}
545
+ self.usage: dict[str, int | float] | None = None
546
+ self.response_id: str | None = None
547
+ self.response_model: str | None = None
548
+ self.attrs: dict[str, Any] = {}
549
+ self.committed_usage: dict[str, int | float] | None = None
550
+ self.prior_contents: list[dict[str, Any]] = []
551
+ self.afc_history: Any = None
552
+
553
+
554
+ def _append_part(parts: list[Any], part: Any) -> None:
555
+ native = _native(part)
556
+ if not isinstance(native, dict):
557
+ parts.append(native)
558
+ return
559
+ text = native.get("text")
560
+ if text is not None:
561
+ thought = native.get("thought", False)
562
+ if parts:
563
+ previous = parts[-1]
564
+ if isinstance(previous, dict) and previous.get("text") is not None:
565
+ if previous.get("thought", False) == thought:
566
+ previous["text"] = f"{previous.get('text', '')}{text}"
567
+ return
568
+ parts.append(native)
569
+
570
+
571
+ def _fold_turn(state: _StreamState) -> None:
572
+ # An AFC turn ended: archive its output and fold its usage so the next
573
+ # internal model call does not overwrite what this turn produced.
574
+ for index in sorted(state.candidates):
575
+ candidate_state = state.candidates[index]
576
+ if candidate_state.parts:
577
+ state.prior_contents.append(
578
+ {"role": candidate_state.role, "parts": candidate_state.parts}
579
+ )
580
+ state.candidates.clear()
581
+ state.finish_reasons.clear()
582
+ state.committed_usage = _sum_usage(state.committed_usage, state.usage)
583
+ state.usage = None
584
+
585
+
586
+ def _record_chunk(chunk: Any, state: _StreamState) -> None:
587
+ try:
588
+ _record_chunk_inner(chunk, state)
589
+ except Exception:
590
+ return
591
+
592
+
593
+ def _record_chunk_inner(chunk: Any, state: _StreamState) -> None:
594
+ response_id = _string(_field(chunk, "response_id"))
595
+ if response_id:
596
+ if state.response_id and response_id != state.response_id:
597
+ _fold_turn(state)
598
+ state.response_id = response_id
599
+ response_model = _string(_field(chunk, "model_version"))
600
+ if response_model:
601
+ state.response_model = response_model
602
+ for index, candidate in enumerate(_sequence_items(_field(chunk, "candidates"))):
603
+ candidate_index = _candidate_index(candidate, index)
604
+ candidate_state = state.candidates.setdefault(candidate_index, _CandidateState())
605
+ content = _field(candidate, "content")
606
+ role = _string(_field(content, "role"))
607
+ if role:
608
+ candidate_state.role = role
609
+ for part in _sequence_items(_field(content, "parts")):
610
+ _append_part(candidate_state.parts, part)
611
+ finish = _enum_value(_field(candidate, "finish_reason"))
612
+ if finish is not None:
613
+ state.finish_reasons[candidate_index] = finish
614
+ _record_candidate_metadata(candidate, candidate_index, state.attrs)
615
+ usage = _generate_usage(_field(chunk, "usage_metadata"), state.attrs)
616
+ if usage:
617
+ state.usage = usage
618
+ feedback = _field(chunk, "prompt_feedback")
619
+ block_reason = _enum_value(_field(feedback, "block_reason"))
620
+ if block_reason:
621
+ state.attrs["google_genai.response.block_reason"] = block_reason
622
+ block_message = _string(_field(feedback, "block_reason_message"))
623
+ if block_message:
624
+ state.attrs["google_genai.response.block_reason_message"] = block_message
625
+ prompt_ratings = _sequence_items(_field(feedback, "safety_ratings"))
626
+ if prompt_ratings:
627
+ encoded = _json_attr([_native(rating) for rating in prompt_ratings])
628
+ if encoded is not None:
629
+ state.attrs["google_genai.response.prompt_safety_ratings"] = encoded
630
+ history = _field(chunk, "automatic_function_calling_history")
631
+ if history:
632
+ native_history = _native(history)
633
+ if native_history:
634
+ state.afc_history = native_history
635
+
636
+
637
+ def _stream_fields(state: _StreamState) -> dict[str, Any]:
638
+ try:
639
+ return _stream_fields_inner(state)
640
+ except Exception:
641
+ return {}
642
+
643
+
644
+ def _stream_fields_inner(state: _StreamState) -> dict[str, Any]:
645
+ fields: dict[str, Any] = {}
646
+ attrs = dict(state.attrs)
647
+ if state.response_id:
648
+ fields["response_id"] = state.response_id
649
+ if state.response_model:
650
+ fields["response_model"] = state.response_model
651
+ output: list[dict[str, Any]] = []
652
+ finish_reasons: list[str] = []
653
+ for index in sorted(state.candidates):
654
+ candidate_state = state.candidates[index]
655
+ if candidate_state.parts:
656
+ output.append({"role": candidate_state.role, "parts": candidate_state.parts})
657
+ if index in state.finish_reasons:
658
+ finish_reasons.append(state.finish_reasons[index])
659
+ if state.afc_history:
660
+ # AFC turns are fully represented in the history (function calls and
661
+ # responses included), mirroring the non-streaming AFC mapping.
662
+ fields["input"] = state.afc_history
663
+ attrs["google_genai.automatic_function_calling"] = True
664
+ elif state.prior_contents:
665
+ output = [*state.prior_contents, *output]
666
+ if finish_reasons:
667
+ fields["finish_reason"] = finish_reasons[0]
668
+ if len(finish_reasons) > 1:
669
+ attrs["gen_ai.response.finish_reasons"] = finish_reasons
670
+ if output:
671
+ fields["output"] = output
672
+ usage = _sum_usage(state.committed_usage, state.usage)
673
+ if usage:
674
+ fields["usage"] = usage
675
+ if attrs:
676
+ fields["attributes"] = attrs
677
+ return fields
678
+
679
+
680
+ def _provider_for_client(client: object | None) -> str:
681
+ if client is not None and getattr(client, "vertexai", False):
682
+ return "gcp.vertex_ai"
683
+ return "gcp.gemini"
684
+
685
+
686
+ def _provider_for_resource(resource: object | None) -> str:
687
+ api_client = getattr(resource, "_api_client", None) if resource is not None else None
688
+ if api_client is not None and getattr(api_client, "vertexai", False):
689
+ return "gcp.vertex_ai"
690
+ return "gcp.gemini"
691
+
692
+
693
+ def _start_generate_span(
694
+ model: Any, contents: Any, config: Any, provider: str
695
+ ) -> tuple[telemetry_dev.SpanHandle, Callable[..., None], float]:
696
+ try:
697
+ name, fields = _generate_request_fields(model, contents, config)
698
+ except Exception:
699
+ name, fields = f"chat {_string(model) or 'unknown'}", {"type": "generation"}
700
+ handle = telemetry_dev.start_span(name, provider=provider, **_clean_fields(fields))
701
+ return handle, _end_once(handle), time.perf_counter()
702
+
703
+
704
+ def _start_embed_span(
705
+ model: Any, contents: Any, config: Any, provider: str
706
+ ) -> tuple[telemetry_dev.SpanHandle, Callable[..., None], float]:
707
+ try:
708
+ name, fields = _embed_request_fields(model, contents, config)
709
+ except Exception:
710
+ name, fields = f"embeddings {_string(model) or 'unknown'}", {"type": "embedding"}
711
+ handle = telemetry_dev.start_span(name, provider=provider, **_clean_fields(fields))
712
+ return handle, _end_once(handle), time.perf_counter()
713
+
714
+
715
+ def _first_chunk_update(chunk: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
716
+ try:
717
+ update: dict[str, Any] = {
718
+ "time_to_first_chunk_ms": (time.perf_counter() - started_at) * 1000,
719
+ }
720
+ response_id = _string(_field(chunk, "response_id"))
721
+ if response_id:
722
+ update["response_id"] = response_id
723
+ response_model = _string(_field(chunk, "model_version"))
724
+ if response_model:
725
+ update["response_model"] = response_model
726
+ handle.update(**_clean_fields(update))
727
+ except Exception:
728
+ return
729
+
730
+
731
+ class _ObservedStream:
732
+ def __init__(
733
+ self,
734
+ inner: Iterator[Any],
735
+ handle: telemetry_dev.SpanHandle,
736
+ end: Callable[..., None],
737
+ started_at: float,
738
+ afc_usage_state: dict[str, Any] | None = None,
739
+ ) -> None:
740
+ self._inner = inner
741
+ self._handle = handle
742
+ self._end = end
743
+ self._started_at = started_at
744
+ self._state = _StreamState()
745
+ self._saw_first = False
746
+ self._afc_usage_state = afc_usage_state
747
+
748
+ def _fields(self) -> dict[str, Any]:
749
+ return _fields_with_afc_usage(_stream_fields(self._state), self._afc_usage_state)
750
+
751
+ def __iter__(self) -> Iterator[Any]:
752
+ try:
753
+ while True:
754
+ try:
755
+ yield self.__next__()
756
+ except StopIteration:
757
+ return
758
+ except BaseException as exc:
759
+ if not isinstance(exc, GeneratorExit):
760
+ self._end(**_clean_fields(self._fields()), error=exc)
761
+ raise
762
+ finally:
763
+ self.close()
764
+
765
+ def __next__(self) -> Any:
766
+ return self._advance(self._inner.__next__)
767
+
768
+ def send(self, value: Any) -> Any:
769
+ def pull() -> Any:
770
+ return self._inner.send(value)
771
+
772
+ return self._advance(pull)
773
+
774
+ def throw(self, *args: Any) -> Any:
775
+ def pull() -> Any:
776
+ return self._inner.throw(*args)
777
+
778
+ return self._advance(pull)
779
+
780
+ def _advance(self, pull: Callable[[], Any]) -> Any:
781
+ token = (
782
+ _AFC_USAGE_STATE.set(self._afc_usage_state)
783
+ if self._afc_usage_state is not None
784
+ else None
785
+ )
786
+ try:
787
+ chunk = pull()
788
+ except StopIteration:
789
+ self._end(**_clean_fields(self._fields()))
790
+ raise
791
+ except BaseException as exc:
792
+ if not isinstance(exc, GeneratorExit):
793
+ self._end(**_clean_fields(self._fields()), error=exc)
794
+ raise
795
+ finally:
796
+ if token is not None:
797
+ _AFC_USAGE_STATE.reset(token)
798
+ if not self._saw_first:
799
+ self._saw_first = True
800
+ _first_chunk_update(chunk, self._handle, self._started_at)
801
+ _record_chunk(chunk, self._state)
802
+ return chunk
803
+
804
+ def close(self) -> None:
805
+ self._end(**_clean_fields(self._fields()))
806
+ close = getattr(self._inner, "close", None)
807
+ if close is not None:
808
+ close()
809
+
810
+
811
+ class _ObservedAsyncStream:
812
+ def __init__(
813
+ self,
814
+ inner: AsyncIterator[Any],
815
+ handle: telemetry_dev.SpanHandle,
816
+ end: Callable[..., None],
817
+ started_at: float,
818
+ afc_usage_state: dict[str, Any] | None = None,
819
+ ) -> None:
820
+ self._inner = inner
821
+ self._handle = handle
822
+ self._end = end
823
+ self._started_at = started_at
824
+ self._state = _StreamState()
825
+ self._saw_first = False
826
+ self._afc_usage_state = afc_usage_state
827
+
828
+ def _fields(self) -> dict[str, Any]:
829
+ return _fields_with_afc_usage(_stream_fields(self._state), self._afc_usage_state)
830
+
831
+ def __aiter__(self) -> AsyncIterator[Any]:
832
+ return self._iterate()
833
+
834
+ async def _iterate(self) -> AsyncIterator[Any]:
835
+ try:
836
+ while True:
837
+ try:
838
+ yield await self.__anext__()
839
+ except StopAsyncIteration:
840
+ return
841
+ except BaseException as exc:
842
+ if not isinstance(exc, GeneratorExit):
843
+ self._end(**_clean_fields(self._fields()), error=exc)
844
+ raise
845
+ finally:
846
+ await self.aclose()
847
+
848
+ def __anext__(self) -> Any:
849
+ return self._advance(self._inner.__anext__)
850
+
851
+ def asend(self, value: Any) -> Any:
852
+ def pull() -> Awaitable[Any]:
853
+ return self._inner.asend(value)
854
+
855
+ return self._advance(pull)
856
+
857
+ def athrow(self, *args: Any) -> Any:
858
+ def pull() -> Awaitable[Any]:
859
+ return self._inner.athrow(*args)
860
+
861
+ return self._advance(pull)
862
+
863
+ async def _advance(self, pull: Callable[[], Awaitable[Any]]) -> Any:
864
+ token = (
865
+ _AFC_USAGE_STATE.set(self._afc_usage_state)
866
+ if self._afc_usage_state is not None
867
+ else None
868
+ )
869
+ try:
870
+ chunk = await pull()
871
+ except StopAsyncIteration:
872
+ self._end(**_clean_fields(self._fields()))
873
+ raise
874
+ except BaseException as exc:
875
+ if not isinstance(exc, GeneratorExit):
876
+ self._end(**_clean_fields(self._fields()), error=exc)
877
+ raise
878
+ finally:
879
+ if token is not None:
880
+ _AFC_USAGE_STATE.reset(token)
881
+ if not self._saw_first:
882
+ self._saw_first = True
883
+ _first_chunk_update(chunk, self._handle, self._started_at)
884
+ _record_chunk(chunk, self._state)
885
+ return chunk
886
+
887
+ async def aclose(self) -> None:
888
+ self._end(**_clean_fields(self._fields()))
889
+ close = getattr(self._inner, "aclose", None)
890
+ if close is not None:
891
+ await close()
892
+
893
+
894
+ def _wrap_sync_collect_generate(
895
+ original: Callable[..., Any], _provider_resolver: ProviderResolver
896
+ ) -> Callable[..., Any]:
897
+ @wraps(original)
898
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
899
+ return _record_afc_usage(original(*args, **kwargs))
900
+
901
+ setattr(wrapper, _WRAPPED_ATTR, True)
902
+ setattr(wrapper, _ORIGINAL_ATTR, original)
903
+ return wrapper
904
+
905
+
906
+ def _wrap_sync_collect_stream(
907
+ original: Callable[..., Any], _provider_resolver: ProviderResolver
908
+ ) -> Callable[..., Any]:
909
+ @wraps(original)
910
+ def wrapper(*args: Any, **kwargs: Any) -> Iterator[Any]:
911
+ for chunk in original(*args, **kwargs):
912
+ yield _record_afc_usage(chunk)
913
+
914
+ setattr(wrapper, _WRAPPED_ATTR, True)
915
+ setattr(wrapper, _ORIGINAL_ATTR, original)
916
+ return wrapper
917
+
918
+
919
+ def _wrap_async_collect_generate(
920
+ original: Callable[..., Any], _provider_resolver: ProviderResolver
921
+ ) -> Callable[..., Any]:
922
+ @wraps(original)
923
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
924
+ return _record_afc_usage(await original(*args, **kwargs))
925
+
926
+ setattr(wrapper, _WRAPPED_ATTR, True)
927
+ setattr(wrapper, _ORIGINAL_ATTR, original)
928
+ return wrapper
929
+
930
+
931
+ def _wrap_async_collect_stream(
932
+ original: Callable[..., Any], _provider_resolver: ProviderResolver
933
+ ) -> Callable[..., Any]:
934
+ @wraps(original)
935
+ async def wrapper(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
936
+ inner = await original(*args, **kwargs)
937
+
938
+ async def observed() -> AsyncIterator[Any]:
939
+ async for chunk in inner:
940
+ yield _record_afc_usage(chunk)
941
+
942
+ return observed()
943
+
944
+ setattr(wrapper, _WRAPPED_ATTR, True)
945
+ setattr(wrapper, _ORIGINAL_ATTR, original)
946
+ return wrapper
947
+
948
+
949
+ def _wrap_sync_generate(
950
+ original: Callable[..., Any], provider_resolver: ProviderResolver
951
+ ) -> Callable[..., Any]:
952
+ @wraps(original)
953
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
954
+ resource = args[0] if args else None
955
+ _handle, end, _started_at = _start_generate_span(
956
+ kwargs.get("model"),
957
+ kwargs.get("contents"),
958
+ kwargs.get("config"),
959
+ provider_resolver(resource),
960
+ )
961
+ afc_state = _afc_usage_state(kwargs.get("config"))
962
+ token = _AFC_USAGE_STATE.set(afc_state) if afc_state is not None else None
963
+ try:
964
+ result = original(*args, **kwargs)
965
+ except BaseException as exc:
966
+ end(error=exc)
967
+ raise
968
+ finally:
969
+ if token is not None:
970
+ _AFC_USAGE_STATE.reset(token)
971
+ end(**_fields_with_afc_usage(_safe_fields(_generate_response_fields, result), afc_state))
972
+ return result
973
+
974
+ setattr(wrapper, _WRAPPED_ATTR, True)
975
+ setattr(wrapper, _ORIGINAL_ATTR, original)
976
+ return wrapper
977
+
978
+
979
+ def _wrap_sync_stream(
980
+ original: Callable[..., Any], provider_resolver: ProviderResolver
981
+ ) -> Callable[..., Any]:
982
+ @wraps(original)
983
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
984
+ resource = args[0] if args else None
985
+ handle, end, started_at = _start_generate_span(
986
+ kwargs.get("model"),
987
+ kwargs.get("contents"),
988
+ kwargs.get("config"),
989
+ provider_resolver(resource),
990
+ )
991
+ afc_state = _afc_usage_state(kwargs.get("config"))
992
+ try:
993
+ result = original(*args, **kwargs)
994
+ except BaseException as exc:
995
+ end(error=exc)
996
+ raise
997
+ return _ObservedStream(result, handle, end, started_at, afc_state)
998
+
999
+ setattr(wrapper, _WRAPPED_ATTR, True)
1000
+ setattr(wrapper, _ORIGINAL_ATTR, original)
1001
+ return wrapper
1002
+
1003
+
1004
+ def _wrap_sync_embed(
1005
+ original: Callable[..., Any], provider_resolver: ProviderResolver
1006
+ ) -> Callable[..., Any]:
1007
+ @wraps(original)
1008
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
1009
+ resource = args[0] if args else None
1010
+ _handle, end, _started_at = _start_embed_span(
1011
+ kwargs.get("model"),
1012
+ kwargs.get("contents"),
1013
+ kwargs.get("config"),
1014
+ provider_resolver(resource),
1015
+ )
1016
+ try:
1017
+ result = original(*args, **kwargs)
1018
+ except BaseException as exc:
1019
+ end(error=exc)
1020
+ raise
1021
+ end(**_safe_fields(_embed_response_fields, result))
1022
+ return result
1023
+
1024
+ setattr(wrapper, _WRAPPED_ATTR, True)
1025
+ setattr(wrapper, _ORIGINAL_ATTR, original)
1026
+ return wrapper
1027
+
1028
+
1029
+ def _wrap_async_generate(
1030
+ original: Callable[..., Any], provider_resolver: ProviderResolver
1031
+ ) -> Callable[..., Any]:
1032
+ @wraps(original)
1033
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
1034
+ resource = args[0] if args else None
1035
+ _handle, end, _started_at = _start_generate_span(
1036
+ kwargs.get("model"),
1037
+ kwargs.get("contents"),
1038
+ kwargs.get("config"),
1039
+ provider_resolver(resource),
1040
+ )
1041
+ afc_state = _afc_usage_state(kwargs.get("config"))
1042
+ token = _AFC_USAGE_STATE.set(afc_state) if afc_state is not None else None
1043
+ try:
1044
+ result = await original(*args, **kwargs)
1045
+ except BaseException as exc:
1046
+ end(error=exc)
1047
+ raise
1048
+ finally:
1049
+ if token is not None:
1050
+ _AFC_USAGE_STATE.reset(token)
1051
+ end(**_fields_with_afc_usage(_safe_fields(_generate_response_fields, result), afc_state))
1052
+ return result
1053
+
1054
+ setattr(wrapper, _WRAPPED_ATTR, True)
1055
+ setattr(wrapper, _ORIGINAL_ATTR, original)
1056
+ return wrapper
1057
+
1058
+
1059
+ def _wrap_async_stream(
1060
+ original: Callable[..., Any], provider_resolver: ProviderResolver
1061
+ ) -> Callable[..., Any]:
1062
+ @wraps(original)
1063
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
1064
+ resource = args[0] if args else None
1065
+ handle, end, started_at = _start_generate_span(
1066
+ kwargs.get("model"),
1067
+ kwargs.get("contents"),
1068
+ kwargs.get("config"),
1069
+ provider_resolver(resource),
1070
+ )
1071
+ afc_state = _afc_usage_state(kwargs.get("config"))
1072
+ try:
1073
+ result = await original(*args, **kwargs)
1074
+ except BaseException as exc:
1075
+ end(error=exc)
1076
+ raise
1077
+ return _ObservedAsyncStream(result, handle, end, started_at, afc_state)
1078
+
1079
+ setattr(wrapper, _WRAPPED_ATTR, True)
1080
+ setattr(wrapper, _ORIGINAL_ATTR, original)
1081
+ return wrapper
1082
+
1083
+
1084
+ def _wrap_async_embed(
1085
+ original: Callable[..., Any], provider_resolver: ProviderResolver
1086
+ ) -> Callable[..., Any]:
1087
+ @wraps(original)
1088
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
1089
+ resource = args[0] if args else None
1090
+ _handle, end, _started_at = _start_embed_span(
1091
+ kwargs.get("model"),
1092
+ kwargs.get("contents"),
1093
+ kwargs.get("config"),
1094
+ provider_resolver(resource),
1095
+ )
1096
+ try:
1097
+ result = await original(*args, **kwargs)
1098
+ except BaseException as exc:
1099
+ end(error=exc)
1100
+ raise
1101
+ end(**_safe_fields(_embed_response_fields, result))
1102
+ return result
1103
+
1104
+ setattr(wrapper, _WRAPPED_ATTR, True)
1105
+ setattr(wrapper, _ORIGINAL_ATTR, original)
1106
+ return wrapper
1107
+
1108
+
1109
+ def _patch_instance(
1110
+ resource: object,
1111
+ method: str,
1112
+ wrapper_factory: Callable[[Callable[..., Any], ProviderResolver], Callable[..., Any]],
1113
+ provider_name: str,
1114
+ ) -> None:
1115
+ current = getattr(resource, method)
1116
+ if getattr(current, _WRAPPED_ATTR, False):
1117
+ if method in vars(resource):
1118
+ return
1119
+ original = getattr(current, _ORIGINAL_ATTR, None)
1120
+ if original is None:
1121
+ return
1122
+ current = original.__get__(resource, type(resource))
1123
+ wrapped = wrapper_factory(current, lambda _: provider_name)
1124
+ setattr(resource, method, wrapped)
1125
+
1126
+
1127
+ def _patch_class(
1128
+ cls: type[Any],
1129
+ method: str,
1130
+ wrapper_factory: Callable[[Callable[..., Any], ProviderResolver], Callable[..., Any]],
1131
+ ) -> None:
1132
+ original = getattr(cls, method)
1133
+ if getattr(original, _WRAPPED_ATTR, False):
1134
+ return
1135
+ _ORIGINALS.append((cls, method, original))
1136
+ setattr(cls, method, wrapper_factory(original, _provider_for_resource))
1137
+
1138
+
1139
+ def _patch_models(resource: object, provider_name: str) -> None:
1140
+ _patch_instance(resource, "_generate_content", _wrap_sync_collect_generate, provider_name)
1141
+ _patch_instance(resource, "_generate_content_stream", _wrap_sync_collect_stream, provider_name)
1142
+ _patch_instance(resource, "generate_content", _wrap_sync_generate, provider_name)
1143
+ _patch_instance(resource, "generate_content_stream", _wrap_sync_stream, provider_name)
1144
+ _patch_instance(resource, "embed_content", _wrap_sync_embed, provider_name)
1145
+
1146
+
1147
+ def _patch_async_models(resource: object, provider_name: str) -> None:
1148
+ _patch_instance(resource, "_generate_content", _wrap_async_collect_generate, provider_name)
1149
+ _patch_instance(resource, "_generate_content_stream", _wrap_async_collect_stream, provider_name)
1150
+ _patch_instance(resource, "generate_content", _wrap_async_generate, provider_name)
1151
+ _patch_instance(resource, "generate_content_stream", _wrap_async_stream, provider_name)
1152
+ _patch_instance(resource, "embed_content", _wrap_async_embed, provider_name)
1153
+
1154
+
1155
+ def wrap_google_genai(client: _T) -> _T:
1156
+ if getattr(client, _WRAPPED_ATTR, False):
1157
+ return client
1158
+ provider_name = _provider_for_client(client)
1159
+ models = getattr(client, "models", None)
1160
+ if models is not None:
1161
+ _patch_models(models, provider_name)
1162
+ aio = getattr(client, "aio", None)
1163
+ async_models = getattr(aio, "models", None) if aio is not None else None
1164
+ if async_models is not None:
1165
+ _patch_async_models(async_models, provider_name)
1166
+ setattr(client, _WRAPPED_ATTR, True)
1167
+ return client
1168
+
1169
+
1170
+ def instrument_google_genai() -> None:
1171
+ global _installed
1172
+ with _install_lock:
1173
+ if _installed:
1174
+ return
1175
+ for method, factory in (
1176
+ ("_generate_content", _wrap_sync_collect_generate),
1177
+ ("_generate_content_stream", _wrap_sync_collect_stream),
1178
+ ("generate_content", _wrap_sync_generate),
1179
+ ("generate_content_stream", _wrap_sync_stream),
1180
+ ("embed_content", _wrap_sync_embed),
1181
+ ):
1182
+ _patch_class(Models, method, factory)
1183
+ for method, factory in (
1184
+ ("_generate_content", _wrap_async_collect_generate),
1185
+ ("_generate_content_stream", _wrap_async_collect_stream),
1186
+ ("generate_content", _wrap_async_generate),
1187
+ ("generate_content_stream", _wrap_async_stream),
1188
+ ("embed_content", _wrap_async_embed),
1189
+ ):
1190
+ _patch_class(AsyncModels, method, factory)
1191
+ _installed = True
1192
+
1193
+
1194
+ def uninstrument_google_genai() -> None:
1195
+ global _installed
1196
+ with _install_lock:
1197
+ while _ORIGINALS:
1198
+ cls, method, original = _ORIGINALS.pop()
1199
+ current = getattr(cls, method, None)
1200
+ if (
1201
+ getattr(current, _WRAPPED_ATTR, False)
1202
+ and getattr(current, _ORIGINAL_ATTR, None) is original
1203
+ ):
1204
+ setattr(cls, method, original)
1205
+ _installed = False
1206
+
1207
+
1208
+ __all__ = [
1209
+ "__version__",
1210
+ "instrument_google_genai",
1211
+ "uninstrument_google_genai",
1212
+ "wrap_google_genai",
1213
+ ]
File without changes
@@ -0,0 +1,134 @@
1
+ Metadata-Version: 2.4
2
+ Name: telemetry-dev-google-genai
3
+ Version: 0.1.0
4
+ Summary: Google GenAI (Gemini) integration for telemetry.dev Python SDK
5
+ Keywords: telemetry,opentelemetry,gemini,google-genai,llm,genai,tracing
6
+ Author: telemetry.dev
7
+ License-Expression: MIT
8
+ Classifier: Development Status :: 4 - Beta
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: Programming Language :: Python :: 3.10
12
+ Classifier: Programming Language :: Python :: 3.11
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Programming Language :: Python :: 3.13
15
+ Classifier: Typing :: Typed
16
+ Requires-Dist: telemetry-dev>=0.1.0
17
+ Requires-Dist: google-genai>=2,<3
18
+ Requires-Python: >=3.10
19
+ Project-URL: Homepage, https://telemetry.dev
20
+ Project-URL: Repository, https://github.com/telemetry-dev/telemetry.dev
21
+ Description-Content-Type: text/markdown
22
+
23
+ # telemetry-dev-google-genai
24
+
25
+ Google GenAI (Gemini) SDK instrumentation for telemetry.dev. It wraps the official `google-genai` Python SDK and emits telemetry.dev generation and embedding spans through `telemetry-dev`.
26
+
27
+ ## Install
28
+
29
+ ```sh
30
+ pip install telemetry-dev-google-genai
31
+ ```
32
+
33
+ Initialize the core SDK first:
34
+
35
+ ```py
36
+ import telemetry_dev
37
+
38
+ telemetry_dev.init(
39
+ api_key="td_live_...",
40
+ base_url="http://localhost:4318",
41
+ service_name="my-service",
42
+ )
43
+ ```
44
+
45
+ ## Per-client wrapping
46
+
47
+ ```py
48
+ from google import genai
49
+ from telemetry_dev_google_genai import wrap_google_genai
50
+
51
+ client = wrap_google_genai(genai.Client(api_key="..."))
52
+
53
+ client.models.generate_content(
54
+ model="gemini-2.5-flash",
55
+ contents=[{"role": "user", "parts": [{"text": "Tell me a joke about OpenTelemetry"}]}],
56
+ )
57
+ ```
58
+
59
+ Use this when you want explicit control over which clients are instrumented. `wrap_google_genai` patches both sync `client.models` and async `client.aio.models`.
60
+
61
+ ## Global instrumentation
62
+
63
+ ```py
64
+ from google import genai
65
+ from telemetry_dev_google_genai import instrument_google_genai, uninstrument_google_genai
66
+
67
+ instrument_google_genai()
68
+ client = genai.Client(api_key="...")
69
+
70
+ try:
71
+ client.models.generate_content(
72
+ model="gemini-2.5-flash",
73
+ contents="Tell me a joke about OpenTelemetry",
74
+ )
75
+ finally:
76
+ uninstrument_google_genai()
77
+ ```
78
+
79
+ Use this as the app-wide one-liner at startup when all GenAI clients should be instrumented.
80
+
81
+ ## Instrumented surfaces
82
+
83
+ Sync and async variants are covered:
84
+
85
+ - `client.models.generate_content(...)` / `client.aio.models.generate_content(...)`
86
+ - `client.models.generate_content_stream(...)` / `client.aio.models.generate_content_stream(...)`
87
+ - `client.models.embed_content(...)` / `client.aio.models.embed_content(...)`
88
+
89
+ `client.chats.create(...).send_message(...)` and `send_message_stream(...)` are covered automatically because they call the wrapped `models` methods.
90
+
91
+ ## What gets captured
92
+
93
+ | Gemini signal | telemetry.dev field / attribute |
94
+ |---|---|
95
+ | `model` | `model` |
96
+ | `contents` | `input` |
97
+ | `config.system_instruction` | `system_instructions` |
98
+ | sampling params (`temperature`, `top_p`, `top_k`, `seed`, penalties, `max_output_tokens`, `stop_sequences`) | same-named span fields |
99
+ | JSON / schema output config | `output_type` (`json` or `text`) |
100
+ | `config.candidate_count` | `gen_ai.request.choice.count` |
101
+ | `config.tools` | `gen_ai.tool.definitions` |
102
+ | `config.tool_config`, `safety_settings`, `thinking_config`, `labels`, `cached_content`, `response_modalities` | `google_genai.request.*` attributes |
103
+ | response IDs, model version, finish reasons, output messages, usage tokens | mapped span fields |
104
+ | block/safety/grounding/url-context metadata | `google_genai.response.*` attributes |
105
+ | AFC history on the final response | replaces span `input`; sets `google_genai.automatic_function_calling=true` |
106
+
107
+ The integration maps native Gemini request/response shapes directly. It never mutates caller requests.
108
+
109
+ ## Streaming
110
+
111
+ Gemini streams already include cumulative `usage_metadata` on chunks, so no request injection is needed. Stream spans record time-to-first-chunk, aggregate text parts (merging consecutive text with the same `thought` flag), last-seen usage/finish reasons, and end once when the stream completes, errors, or is closed.
112
+
113
+ ## Automatic function calling (AFC)
114
+
115
+ When Python callables are passed in `tools`, the SDK may run an internal AFC loop across multiple transport calls. The integration emits one span for the public `generate_content` call and, when present, sets span input to `automatic_function_calling_history`.
116
+
117
+ ## Embeddings
118
+
119
+ Embedding calls emit `gen_ai.operation.name = "embeddings"`, request model/input, embedding count/dimension attributes, optional token usage, and optional billable character counts. Embedding vectors are not captured as output.
120
+
121
+ ## Provider values
122
+
123
+ - Gemini Developer API clients record provider `gcp.gemini`.
124
+ - Vertex AI clients (`vertexai=True`) record provider `gcp.vertex_ai`.
125
+
126
+ ## Fail-open guarantee
127
+
128
+ Mapping code is defensive; wrapped calls return the SDK response unchanged and re-raise exceptions untouched. Telemetry bugs never break callers.
129
+
130
+ ## Limitations
131
+
132
+ - Not instrumented: `count_tokens`, `compute_tokens`, `generate_images`, `generate_videos`, `live`, `caches`, `files`, `tunings`, `batches`.
133
+ - Deliberately not captured: logprobs, citation metadata, per-modality token detail arrays, `create_time`, `sdk_http_response`.
134
+ - Unconsumed streams end their spans only when the stream is exhausted, errors, or is closed.
@@ -0,0 +1,5 @@
1
+ telemetry_dev_google_genai/__init__.py,sha256=hLlMg5EgpJ2syD3TkW840Tooz8wg5OSCMzEVW7xE7qc,43687
2
+ telemetry_dev_google_genai/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
3
+ telemetry_dev_google_genai-0.1.0.dist-info/WHEEL,sha256=bu7Cckf7DKpj8ztfOQHBiWtXM4xigujiTkrhvV6U2aU,81
4
+ telemetry_dev_google_genai-0.1.0.dist-info/METADATA,sha256=Msg1oPgHiBTlWYZehVGlSk-91jLZAU2IIHONCyu-SiA,5441
5
+ telemetry_dev_google_genai-0.1.0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: uv 0.11.27
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any