"""Auto-instrumentation for Google ADK.
The original implementation crammed three independent monkey-patches into
a single ``patch()`` body, which made partial failures hard to recover
from (a successful first patch left ``_patched=True`` even when later
patches blew up, and ``unpatch()`` would then try to undo work that
never happened). This module now exposes three small patch helpers and
``patch()`` simply orchestrates them with isolated state.
"""
from __future__ import annotations
import contextlib
import logging
import os
from typing import Any, Callable
from taimoe.platform.observability import current_trace, trace
from ._helpers import safe_dump
logger = logging.getLogger(__name__)
# Per-patch state. Each block has its own pair of (patched flag, originals)
# so a failure in one doesn't leak into the others.
_gemini_patched = False
_original_generate: Callable[..., Any] | None = None
_llm_agent_patched = False
_original_canonical_model: Any = None
_tool_to_mldev_patched = False
_original_tool_to_mldev: Callable[..., Any] | None = None
# --- Public surface --------------------------------------------------------
[docs]
def patch() -> bool:
"""Apply every supported monkey-patch. Returns True if any succeeded."""
results = [
_patch_gemini_generate_content(),
_patch_llm_agent_canonical_model(),
_patch_tool_to_mldev(),
]
return any(results)
[docs]
def unpatch() -> None:
"""Reverse whatever ``patch()`` managed to apply."""
_unpatch_gemini_generate_content()
_unpatch_llm_agent_canonical_model()
_unpatch_tool_to_mldev()
# --- Patch 1: Gemini.generate_content_async --------------------------------
#
# Wraps every Gemini call with: gateway re-routing (per-call cost; see
# WARNING below) + an observability span recording usage and grounding
# citations.
#
# WARNING — gateway hijacking cost. ``patched_generate`` rewrites
# ``self.base_url`` and ``os.environ["GEMINI_API_KEY"]`` on *every* call.
# This is intentional but expensive: ADK's ``Gemini`` is a frozen Pydantic
# model whose ``base_url`` we can only bypass via ``object.__setattr__``,
# and there's no public init point we can hook to do it once. Moving to
# a once-at-init wrapper is on the roadmap; until then, the cost of doing
# this per-call is two attr writes + one ``os.environ`` set per generate
# (small but non-zero).
#
# IMPORTANT — ``os.environ`` is a process-wide global. In a multi-tenant
# host that runs Gemini calls for more than one Taimoe token concurrently,
# tokens can race. We accept that limitation today because every Taimoe
# runtime owns its own process.
def _patch_gemini_generate_content() -> bool:
"""Wrap ``Gemini.generate_content_async`` with gateway hijack + tracing."""
global _gemini_patched, _original_generate
if _gemini_patched:
return True
try:
import google.adk.models # type: ignore
except ImportError:
logger.debug("google-adk not installed; skipping Gemini patch.")
return False
target_class = getattr(google.adk.models, "Gemini", None)
if not target_class or not hasattr(target_class, "generate_content_async"):
logger.debug("google.adk.models.Gemini.generate_content_async missing.")
return False
_original_generate = target_class.generate_content_async
target_class.generate_content_async = _build_patched_generate(_original_generate)
_gemini_patched = True
logger.info("Auto-instrumented Google ADK Gemini.generate_content_async.")
return True
def _unpatch_gemini_generate_content() -> None:
global _gemini_patched, _original_generate
if not _gemini_patched:
return
try:
import google.adk.models # type: ignore
target_class = getattr(google.adk.models, "Gemini", None)
if target_class and _original_generate is not None:
target_class.generate_content_async = _original_generate
except Exception as exc: # noqa: BLE001
logger.warning("Failed to unpatch Gemini.generate_content_async: %s", exc)
finally:
_gemini_patched = False
def _build_patched_generate(original_generate: Callable[..., Any]) -> Callable[..., Any]:
"""Construct the async-generator wrapper around ``generate_content_async``.
Pulled out of ``_patch_gemini_generate_content`` purely for readability
— the wrapper is the bulk of the patch.
Note: the wrapper is ``async def`` but does **not** ``await`` the
underlying ``original_generate(...)``; ADK returns an async generator
from that call, which we then ``async for`` over. Don't add ``await``.
"""
async def patched_generate(self: Any, *args: Any, **kwargs: Any) -> Any:
_route_through_gateway(self)
active_trace = current_trace()
llm_request = args[0] if args else kwargs.get("llm_request")
real_model = (
getattr(llm_request, "model", None)
or getattr(self, "model", None)
or "unknown-model"
)
agent_name = (
getattr(self, "_taimoe_agent_name", None)
or _extract_agent_name_from_request(llm_request)
or "gemini-adk"
)
if active_trace is None:
trace_ctx: Any = trace(name=f"adk-{agent_name}", agent_id=agent_name)
else:
trace_ctx = contextlib.nullcontext(active_trace)
with trace_ctx as trc:
with trc.span(
"llm",
name=f"{real_model}.generate_content",
input={"args": safe_dump(args), "kwargs": safe_dump(kwargs)},
attributes={
"model": real_model,
"provider": "google",
"agent_name": agent_name,
},
agent_id=agent_name,
) as span:
try:
gen = original_generate(self, *args, **kwargs)
citations: list[dict[str, str]] = []
seen_keys: set[tuple[int, int]] = set()
stream_event_idx = 0
async for chunk in gen:
_record_usage(span, chunk)
_record_grounding(span, chunk, stream_event_idx, citations, seen_keys)
yield chunk
stream_event_idx += 1
logger.debug(
"[Taimoe] collected %d citations for agent %s",
len(citations), agent_name,
)
if citations:
_emit_synthetic_tool_span(trc, span, citations, agent_name)
except Exception as exc:
span.set_error(exc)
raise
return patched_generate
def _route_through_gateway(gemini_instance: Any) -> None:
"""Redirect a ``Gemini`` instance at the Taimoe Gateway.
Idempotent (no-op if already pointing at the gateway). See the module
docstring for cost / multi-tenant caveats.
"""
try:
from taimoe.platform.init import get_registry
registry = get_registry()
if not (registry.platform_url and registry.platform_token):
return
gateway_url = f"{registry.platform_url.rstrip('/')}/api/v1/gateway/gemini"
if getattr(gemini_instance, "base_url", None) == gateway_url:
return
os.environ["GEMINI_API_KEY"] = registry.platform_token
# ``Gemini`` is a frozen Pydantic model; bypass via __setattr__.
object.__setattr__(gemini_instance, "base_url", gateway_url)
cls = type(gemini_instance)
original_tracking_headers = getattr(cls, "_tracking_headers")
def patched_tracking_headers() -> dict[str, str]:
headers = original_tracking_headers(gemini_instance)
headers["x-api-key"] = registry.platform_token
return headers
object.__setattr__(gemini_instance, "_tracking_headers", patched_tracking_headers)
# Force the API client to rebuild against the new base_url.
object.__setattr__(gemini_instance, "_api_client", None)
except Exception as exc: # noqa: BLE001 — never break a model call
logger.warning("Failed to route Gemini through Taimoe Gateway: %s", exc)
def _record_usage(span: Any, chunk: Any) -> None:
"""Stamp token usage from a stream chunk onto the LLM span."""
usage = getattr(chunk, "usage_metadata", None)
if usage is None:
return
span.set_usage(
prompt_tokens=getattr(usage, "prompt_token_count", None),
completion_tokens=getattr(usage, "candidates_token_count", None),
)
# Gemini 2.5 reports cached + thinking tokens separately;
# downstream cost calculators read these from attributes.
cached = getattr(usage, "cached_content_token_count", None)
thinking = getattr(usage, "thoughts_token_count", None)
if cached:
span.set_attribute("cached_tokens", cached)
if thinking:
span.set_attribute("thinking_tokens", thinking)
def _record_grounding(
span: Any,
chunk: Any,
stream_event_idx: int,
citations: list[dict[str, str]],
seen_keys: set[tuple[int, int]],
) -> None:
"""Stamp grounding citations onto the span and accumulate for tool span."""
grounding_metadata = getattr(chunk, "grounding_metadata", None)
if not grounding_metadata:
candidates = getattr(chunk, "candidates", None)
if candidates and len(candidates) > 0:
grounding_metadata = getattr(candidates[0], "grounding_metadata", None)
if not grounding_metadata:
return
grounding_chunks = getattr(grounding_metadata, "grounding_chunks", None) or []
for i, ch in enumerate(grounding_chunks):
rc = getattr(ch, "retrieved_context", None)
if rc:
doc_name = getattr(rc, "document_name", "") or ""
title = getattr(rc, "title", "") or ""
text = getattr(rc, "text", "") or ""
uri = getattr(rc, "uri", "") or doc_name
span.set_attribute(f"grounding.citations.{i}.title", title or doc_name)
span.set_attribute(f"grounding.citations.{i}.uri", uri)
if text:
span.set_attribute(f"grounding.citations.{i}.text", text[:200])
key = (stream_event_idx, i)
if key not in seen_keys:
seen_keys.add(key)
citations.append({
"title": title or doc_name,
"uri": uri,
"text": text[:200] if text else "",
})
continue
web = getattr(ch, "web", None)
if not web:
continue
title = getattr(web, "title", "") or "來源"
uri = getattr(web, "uri", "") or ""
span.set_attribute(f"grounding.citations.{i}.title", title)
span.set_attribute(f"grounding.citations.{i}.uri", uri)
key = (stream_event_idx, i)
if key not in seen_keys:
seen_keys.add(key)
citations.append({"title": title, "uri": uri, "text": ""})
def _emit_synthetic_tool_span(
trace_ctx: Any, llm_span: Any, citations: list[dict[str, str]], agent_name: str
) -> None:
"""Create a synthetic tool span so the trace graph shows LLM → TOOL."""
tool_span = trace_ctx.span(
"tool",
name="vertex_ai_search",
input={"query": "(grounding via Gemini built-in tool)"},
attributes={
"tool_name": "vertex_ai_search",
"citation_count": len(citations),
},
agent_id=agent_name,
)
# Manually wire parent so the graph renders LLM → TOOL.
tool_span.parent_span_id = llm_span.span_id
with tool_span as ts:
ts.set_output(citations)
def _extract_agent_name_from_request(llm_request: Any) -> str | None:
"""Pull an agent identifier out of ``LlmRequest.config.labels`` if present.
Soft contract: ADK doesn't propagate the calling agent's name to the
model layer by default, but the SDK's ``LlmAgent.canonical_model``
patch (below) and any caller that wants to attribute traces correctly
can set ``config.labels['agent_name']``. Returns None if absent.
"""
if llm_request is None:
return None
try:
config = getattr(llm_request, "config", None)
labels = getattr(config, "labels", None) if config else None
if labels and isinstance(labels, dict):
return labels.get("agent_name") or labels.get("agent")
except Exception as exc: # noqa: BLE001
logger.debug("Failed to extract agent name from request: %s", exc)
return None
# --- Patch 2: LlmAgent.canonical_model -------------------------------------
#
# Stamps the owning agent's name onto the model instance so the Gemini
# patch above can attribute spans correctly. Idempotent via a per-class
# ``_taimoe_patched`` marker.
def _patch_llm_agent_canonical_model() -> bool:
global _llm_agent_patched, _original_canonical_model
if _llm_agent_patched:
return True
try:
import google.adk.agents # type: ignore
except ImportError:
return False
llm_agent_class = getattr(google.adk.agents, "LlmAgent", None)
if llm_agent_class is None:
return False
canonical_model = getattr(llm_agent_class, "canonical_model", None)
if not canonical_model or getattr(llm_agent_class, "_taimoe_patched", False):
return False
_original_canonical_model = canonical_model
original_fget = canonical_model.fget
def patched_fget(self: Any) -> Any:
model_instance = original_fget(self)
try:
agent_name = getattr(self, "name", None)
if agent_name and model_instance is not None:
setattr(model_instance, "_taimoe_agent_name", agent_name)
except Exception as exc: # noqa: BLE001
logger.debug("Failed to stamp agent name on model: %s", exc)
return model_instance
llm_agent_class.canonical_model = property(patched_fget)
llm_agent_class._taimoe_patched = True
_llm_agent_patched = True
return True
def _unpatch_llm_agent_canonical_model() -> None:
global _llm_agent_patched, _original_canonical_model
if not _llm_agent_patched:
return
try:
import google.adk.agents # type: ignore
llm_agent_class = getattr(google.adk.agents, "LlmAgent", None)
if llm_agent_class is not None and _original_canonical_model is not None:
llm_agent_class.canonical_model = _original_canonical_model
try:
del llm_agent_class._taimoe_patched
except AttributeError:
pass
except Exception as exc: # noqa: BLE001
logger.warning("Failed to unpatch LlmAgent.canonical_model: %s", exc)
finally:
_llm_agent_patched = False
_original_canonical_model = None
# --- Patch 3: google.genai.models._Tool_to_mldev ---------------------------
#
# The Taimoe Gateway supports Vertex AI tools (e.g. VertexAiSearchTool)
# even when the client is using the Developer API. The vanilla
# ``_Tool_to_mldev`` raises ``ValueError`` for ``retrieval=`` — we
# temporarily hide it across the conversion and re-attach it on the
# output dict so the Gateway can dispatch the retrieval upstream.
def _patch_tool_to_mldev() -> bool:
global _tool_to_mldev_patched, _original_tool_to_mldev
if _tool_to_mldev_patched:
return True
try:
import google.genai.models # type: ignore
except ImportError:
return False
original = getattr(google.genai.models, "_Tool_to_mldev", None)
if original is None:
return False
_original_tool_to_mldev = original
def patched_tool_to_mldev(
from_object: Any, parent_object: Any = None, root_object: Any = None
) -> Any:
has_retrieval = False
retrieval_val = None
if isinstance(from_object, dict) and "retrieval" in from_object:
has_retrieval = True
retrieval_val = from_object.pop("retrieval")
elif hasattr(from_object, "retrieval") and getattr(from_object, "retrieval") is not None:
has_retrieval = True
retrieval_val = from_object.retrieval
from_object.retrieval = None
try:
to_object = original(from_object, parent_object, root_object)
finally:
if has_retrieval:
if isinstance(from_object, dict):
from_object["retrieval"] = retrieval_val
else:
from_object.retrieval = retrieval_val
if hasattr(retrieval_val, "model_dump"):
to_object["retrieval"] = retrieval_val.model_dump(
exclude_none=True, by_alias=True
)
else:
to_object["retrieval"] = retrieval_val
return to_object
google.genai.models._Tool_to_mldev = patched_tool_to_mldev
_tool_to_mldev_patched = True
return True
def _unpatch_tool_to_mldev() -> None:
global _tool_to_mldev_patched, _original_tool_to_mldev
if not _tool_to_mldev_patched:
return
try:
import google.genai.models # type: ignore
if _original_tool_to_mldev is not None:
google.genai.models._Tool_to_mldev = _original_tool_to_mldev
except Exception as exc: # noqa: BLE001
logger.warning("Failed to unpatch _Tool_to_mldev: %s", exc)
finally:
_tool_to_mldev_patched = False
_original_tool_to_mldev = None