Observability in the Agent Storybook Project: Adding Traces to ADK

# OpenTelemetry GenAI semantic conventions
GEN_AI_SYSTEM = "gen_ai.system"
GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
GEN_AI_USAGE_PROMPT_TOKENS = "gen_ai.usage.prompt_tokens"
GEN_AI_USAGE_COMPLETION_TOKENS = "gen_ai.usage.completion_tokens"


class SpanContextManager:
    def __init__(
        self,
        name: str,
        attributes: dict[str, Any] | None = None,
        tracer: trace.Tracer | None = None,
    ) -> None:
        self.name = name
        self.attributes = attributes or {}
        self.tracer = tracer
        self._cm: Any = None
        self.span: trace.Span | None = None

    def __enter__(self) -> trace.Span:
        tracer = self.tracer or get_tracer()
        self._cm = tracer.start_as_current_span(self.name, attributes=self.attributes)
        self.span = self._cm.__enter__()
        return self.span

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
        if self.span is not None and exc_val is None:
            if self.span.is_recording():
                status = getattr(self.span, "status", None)
                if status is not None and getattr(status, "status_code", None) == trace.StatusCode.UNSET:
                    self.span.set_status(trace.StatusCode.OK)
        if self._cm is not None:
            return self._cm.__exit__(exc_type, exc_val, exc_tb)
        return False

    async def __aenter__(self) -> trace.Span:
        return self.__enter__()

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> Any:
        return self.__exit__(exc_type, exc_val, exc_tb)


def trace_stage(stage_name: str, **attributes: Any) -> SpanContextManager:
    """Wrap a pipeline stage with proper span hierarchy and error recording."""
    return SpanContextManager(stage_name, attributes=attributes)


def trace_agent_call(
    agent_name: str,
    model: str | None = None,
    **attributes: Any,
) -> SpanContextManager:
    """Wrap model/agent calls with standard GenAI semantic conventions."""
    span_name = (
        agent_name
        if ("." in agent_name or agent_name.startswith("agent."))
        else f"agent.{agent_name}"
    )
    system = attributes.pop(GEN_AI_SYSTEM, "gemini")
    operation_name = attributes.pop(GEN_AI_OPERATION_NAME, "generate_content")
    span_attrs: dict[str, Any] = {
        GEN_AI_SYSTEM: system,
        GEN_AI_OPERATION_NAME: operation_name,
        "agent.name": agent_name,
    }
    if model:
        span_attrs[GEN_AI_REQUEST_MODEL] = model
    span_attrs.update(attributes)
    return SpanContextManager(span_name, attributes=span_attrs)


def set_span_token_usage(
    span: trace.Span | None = None,
    prompt_tokens: int | None = None,
    completion_tokens: int | None = None,
) -> None:
    """Attach token usage metrics to a span if available."""
    target_span = span or trace.get_current_span()
    if target_span is not None and target_span.is_recording():
        if prompt_tokens is not None:
            target_span.set_attribute(GEN_AI_USAGE_PROMPT_TOKENS, prompt_tokens)
        if completion_tokens is not None:
            target_span.set_attribute(GEN_AI_USAGE_COMPLETION_TOKENS, completion_tokens)

        with trace_stage("stage.generate_illustrations", session_id=sid, total_spreads=total_spreads):
            log.info(
                "Starting stage.generate_illustrations for session %s (total_spreads=%d, model=%s)",
                sid,
                total_spreads,
                settings.model_image,
            )
            image_sem = asyncio.Semaphore(settings.image_concurrency)
            llm_sem = asyncio.Semaphore(settings.llm_concurrency)
            ref_ready: asyncio.Event = asyncio.Event()
            ref_image: list[bytes] = []

# src/storybook/agents/pipeline.py
from google.adk.models.google_llm import Gemini

_original_generate_content_async = Gemini.generate_content_async


async def _retryable_generate_content_async(
    self: Gemini, llm_request: Any, stream: bool = False
) -> Any:
    """Wrap Gemini.generate_content_async with fine-grained 429 retry.

    Retries transient 429 RESOURCE_EXHAUSTED errors at the individual LLM call
    level with exponential backoff. This prevents an entire LoopAgent (e.g.
    craft_adapt_validate) or multi-step workflow from aborting and restarting
    from scratch when a single validation or generation call hits rate limits.
    """
    max_retries = 5
    for attempt in range(max_retries):
        yielded = False
        try:
            async for item in _original_generate_content_async(self, llm_request, stream=stream):
                yielded = True
                yield item
            return
        except Exception as exc:
            msg = str(exc)
            if not yielded and ("RESOURCE_EXHAUSTED" in msg or "429" in msg) and attempt < max_retries - 1:
                wait = 5 * (2 ** attempt)
                log.warning(
                    "429 RESOURCE_EXHAUSTED on model %s; retrying LLM call in %ds (attempt %d/%d)",
                    getattr(llm_request, "model", "model"),
                    wait,
                    attempt + 1,
                    max_retries,
                )
                await asyncio.sleep(wait)
                continue
            raise


Gemini.generate_content_async = _retryable_generate_content_async

# src/storybook/agents/pipeline.py
with check_image_validation(
    attempt=attempt,
    max_attempts=settings.image_max_retries + 1,
    spread_number=spread_number,
    image_index=image_index,
) as val_span:
    async with llm_sem:
        result = await _run_agent(
            validator_runner,
            val_sid,
            validate_input,
            subject_image=img_bytes,
            reference_image=style_ref,
            prev_spread_image=prev_img,
        )

    # Inspect deterministic ADK session state instead of scraping raw output text!
    vsession = None
    if hasattr(validator_runner, "session_service") and hasattr(validator_runner.session_service, "get_session"):
        session_coro = validator_runner.session_service.get_session(
            app_name=validator_runner.app_name,
            user_id="pipeline",
            session_id=val_sid,
        )
        vsession = await session_coro if inspect.isawaitable(session_coro) else session_coro

    if vsession and "validation.passed" in vsession.state:
        passed = bool(vsession.state["validation.passed"])
        score = float(vsession.state.get("validation.score", 1.0 if passed else 0.0))
        reasons = vsession.state.get("validation.reasons") or ([] if passed else [result.strip()])
        revised_prompt = vsession.state.get("revised_image_prompt")
        if revised_prompt and not passed:
            current_prompt = revised_prompt  # Self-healing retry with revised prompt
    else:
        passed = "approved" in result.lower()
        score = 1.0 if passed else 0.0
        reasons = [] if passed else [result.strip()]

    record_validation_result(val_span, passed=passed, score=score, attempt=attempt, reasons=reasons)