import logging import time from typing import Any, Dict, List, Optional, Tuple, Union import google.adk.agents from google.adk.agents import callback_context from google.adk import models from google.adk.tools import base_tool from google.adk.tools import tool_context import opik from opik import context_storage from opik.api_objects import opik_client, span, trace from opik.types import DistributedTraceHeadersDict from opik.decorator import span_creation_handler, arguments_helpers from . import ( helpers as adk_helpers, callback_context_info_extractors, patchers, ) from .patchers import ( litellm_wrappers, llm_response_wrapper, ) from .patchers.adk_otel_tracer import llm_span_helpers from .graph import mermaid_graph_builder LOGGER = logging.getLogger(__name__) SpanOrTraceData = Union[span.SpanData, trace.TraceData] class OpikTracer: """ Opik tracer for google-adk. """ def __init__( self, name: Optional[str] = None, tags: Optional[List[str]] = None, metadata: Optional[Dict[str, Any]] = None, project_name: Optional[str] = None, distributed_headers: Optional[DistributedTraceHeadersDict] = None, ): """ Initialize OpikTracer. Arguments: name: The default name for root span or trace created by the tracer. tags: The default tags for all the traces and spans created by the tracer. metadata: The default metadata for all the traces and spans created by the tracer. project_name: The name of the project for tracing. distributed_headers: The distributed trace headers. """ self.name = name self.tags = tags self.metadata = metadata or {} self.metadata["created_from"] = "google-adk" self.project_name = project_name self._distributed_headers = distributed_headers self._init_internal_attributes() def _init_internal_attributes(self) -> None: self._last_model_output: Optional[Dict[str, Any]] = None self._opik_client = opik_client.get_client_cached() # Track time-to-first-token: map span_id -> (request_start_time, first_token_time) self._ttft_tracking: Dict[str, Tuple[float, Optional[float]]] = {} patchers.patch_adk( self._opik_client, distributed_headers=self._distributed_headers ) def _has_response_content(self, llm_response: models.LlmResponse) -> bool: """ Check if the LlmResponse contains actual content (text or function calls). Arguments: llm_response: The LLM response to check. Returns: True if the response contains text content or function calls, False otherwise. """ try: # Check the LlmResponse object directly for content structure if llm_response.content is not None and llm_response.content.parts: for part in llm_response.content.parts: # Check for text content if part.text and part.text.strip(): return True # Check for function call content (tool calls) if part.function_call: return True return False except Exception as e: LOGGER.debug( f"Error checking LlmResponse.content.parts for TTFT: {e}", exc_info=True, ) return False def _safe_ttft_tracking( self, span_id: Optional[str], pop: bool = False ) -> Tuple[Optional[float], Optional[float]]: """ Safely retrieve time-to-first-token tracking data for a span. Arguments: span_id: The span ID to look up in tracking. pop: If True, remove the entry after fetching. If False, keep it. Returns: Tuple of (request_start_time, first_token_time). Returns (None, None) if span_id is None or not found in tracking. """ if span_id is None or span_id not in self._ttft_tracking: return (None, None) if pop: return self._ttft_tracking.pop(span_id) return self._ttft_tracking[span_id] def flush(self) -> None: self._opik_client.flush() def before_agent_callback( self, callback_context: callback_context.CallbackContext, *args: Any, **kwargs: Any, ) -> None: try: current_trace = context_storage.get_trace_data() current_span = context_storage.top_span_data() thread_id, session_metadata = ( callback_context_info_extractors.try_get_session_info(callback_context) ) agent_metadata = self.metadata.copy() agent_metadata["adk_invocation_id"] = callback_context.invocation_id agent_metadata.update(session_metadata) _try_add_agent_graph_to_metadata(agent_metadata, callback_context) if callback_context.user_content is not None: user_input = adk_helpers.convert_adk_base_model_to_dict( callback_context.user_content ) else: user_input = None name = self.name or callback_context.agent_name if current_span is not None: current_span.update( name=name, metadata={**agent_metadata}, input=user_input, tags=self.tags, project_name=self.project_name, ) elif current_trace is not None: current_trace.update( name=name, metadata={**agent_metadata}, input=user_input, tags=self.tags, thread_id=thread_id, project_name=self.project_name, ) else: LOGGER.warning( f"No current span or trace found in context for agent: {callback_context.agent_name}" ) except Exception as e: LOGGER.error(f"Failed during before_agent_callback(): {e}", exc_info=True) def after_agent_callback( self, callback_context: callback_context.CallbackContext, *args: Any, **kwargs: Any, ) -> None: try: output = self._last_model_output current_span = context_storage.top_span_data() current_trace = context_storage.get_trace_data() if current_span is not None: current_span.update( output=output, project_name=self.project_name, ) elif current_trace is not None: current_trace.update( output=output, project_name=self.project_name, ) self._last_model_output = None else: LOGGER.warning( "No current span or trace found in context for agent output update" ) except Exception as e: LOGGER.error(f"Failed during after_agent_callback(): {e}", exc_info=True) def before_model_callback( self, callback_context: callback_context.CallbackContext, llm_request: models.LlmRequest, *args: Any, **kwargs: Any, ) -> None: try: input = adk_helpers.convert_adk_base_model_to_dict(llm_request) provider, model = litellm_wrappers.parse_provider_and_model( llm_request.model ) if provider is None: provider = adk_helpers.get_adk_provider() # ADK runs `before_model_callback` before running `start_as_current_span` function for the LLM call, # which makes it impossible to update the Opik span from this method. # So we create a span manually here. This flow is handled inside ADKTracerWrapper. result = span_creation_handler.create_span_respecting_context( start_span_arguments=arguments_helpers.StartSpanParameters( name=model, project_name=self.project_name, metadata={ **self.metadata, llm_span_helpers.SPAN_STATUS: llm_span_helpers.LLMSpanStatus.STARTED, }, type="llm", model=model, provider=provider, input=input, ), distributed_trace_headers=None, ) context_storage.add_span_data(result.span_data) # Track request start time for time-to-first-token calculation request_start_time = time.time() self._ttft_tracking[result.span_data.id] = (request_start_time, None) except Exception as e: LOGGER.error(f"Failed during before_model_callback(): {e}", exc_info=True) def after_model_callback( self, callback_context: callback_context.CallbackContext, llm_response: models.LlmResponse, *args: Any, **kwargs: Any, ) -> None: try: is_partial = llm_response.partial is True except Exception: LOGGER.debug("Error checking for partial chunks", exc_info=True) is_partial = False span_id: Optional[str] = None exception_occurred = False try: model = None usage = None output = None if adk_helpers.has_empty_text_part_content(llm_response): # Clean up TTFT tracking if it exists before early return current_span = context_storage.top_span_data() if current_span is not None and current_span.id is not None: self._safe_ttft_tracking(current_span.id, pop=True) return current_span = context_storage.top_span_data() if current_span is None: LOGGER.warning( "No current span found in context for model output update" ) return # Store span_id early for cleanup on all exit paths span_id = current_span.id # Track time-to-first-token: detect first token arrival # We check for first token on EVERY callback (including partial chunks) # to catch the first moment content appears request_start_time, first_token_time = self._safe_ttft_tracking( span_id, pop=False ) if ( first_token_time is None and request_start_time is not None and span_id is not None ): # Check if this response contains actual content (first token) # Content can be text or function calls (tool calls) if self._has_response_content(llm_response): # First token detected - record the time first_token_time = time.time() self._ttft_tracking[span_id] = ( request_start_time, first_token_time, ) # Ignore partial chunks for final processing, ADK will call this method with the full response at the end # Note: We intentionally keep the TTFT tracking entry for partial chunks since ADK will call # this method again with the final non-partial response, where we'll properly clean it up if is_partial: return try: output = adk_helpers.convert_adk_base_model_to_dict(llm_response) usage_data = llm_response_wrapper.pop_llm_usage_data( output, current_span.provider ) if usage_data is not None: model = usage_data.model usage = usage_data.opik_usage except Exception as e: LOGGER.debug( f"Error converting LlmResponse to dict or extracting usage data, reason: {e}", exc_info=True, ) # Calculate time-to-first-token and add to metadata metadata_update = {} request_start_time, first_token_time = self._safe_ttft_tracking( span_id, pop=True ) if first_token_time is not None and request_start_time is not None: time_to_first_token = first_token_time - request_start_time metadata_update["time_to_first_token"] = time_to_first_token # Merge with existing metadata if current_span.metadata is None: current_span.metadata = {} current_span.metadata.update(metadata_update) current_span.metadata[llm_span_helpers.SPAN_STATUS] = ( llm_span_helpers.LLMSpanStatus.READY_FOR_FINALIZATION.value ) current_span.update( output=output, name=model or current_span.model, type="llm", model=model, usage=usage, metadata=current_span.metadata, project_name=self.project_name, ) context_storage.pop_span_data(ensure_id=current_span.id) current_span.init_end_time() # We close this span manually because otherwise ADK will close it too late, # and it will also add tool spans inside of it, which we want to avoid. if opik.is_tracing_active(): self._opik_client.span(**current_span.as_parameters) self._last_model_output = output except Exception as e: exception_occurred = True LOGGER.error(f"Failed during after_model_callback(): {e}", exc_info=True) finally: # Clean up TTFT tracking entry on all exit paths to prevent memory leak # Skip cleanup for partial chunks (normal return) since ADK will call again with final response # For final responses, entry is already popped at line 325, so this is a no-op # For errors (exception_occurred=True) or early returns, this ensures cleanup happens if span_id is not None and (exception_occurred or not is_partial): self._ttft_tracking.pop(span_id, None) def before_tool_callback( self, tool: base_tool.BaseTool, args: Dict[str, Any], tool_context: tool_context.ToolContext, *other_args: Any, **kwargs: Any, ) -> None: try: current_span = context_storage.top_span_data() tool_metadata = { "function_call_id": tool_context.function_call_id, **self.metadata, } # Update existing span with tool information if current_span is not None: current_span.update( name=tool.name, type="tool", input=args, metadata={**tool_metadata}, project_name=self.project_name, ) else: LOGGER.warning( f"No current span found in context for tool: {tool.name}" ) _log_tool_context_warning(context=tool_context) except Exception as e: LOGGER.error(f"Failed during before_tool_callback(): {e}", exc_info=True) def after_tool_callback( self, tool: base_tool.BaseTool, args: Dict[str, Any], tool_context: tool_context.ToolContext, tool_response: Any, *other_args: Any, **kwargs: Any, ) -> None: try: # Debug logging for callback invocation current_span = context_storage.top_span_data() output = ( tool_response if isinstance(tool_response, dict) else {"output": tool_response} ) # Update existing span with tool output if current_span is not None: current_span.update( output=output, project_name=self.project_name, ) else: LOGGER.warning( f"No current span found in context for tool output update: {tool.name}" ) _log_tool_context_warning(context=tool_context) except Exception as e: LOGGER.error(f"Failed during after_tool_callback(): {e}", exc_info=True) def __getstate__(self) -> Dict[str, Any]: state = self.__dict__.copy() state.pop("_opik_client", None) # Don't serialize TTFT tracking as it's runtime state state.pop("_ttft_tracking", None) return state def __setstate__(self, state: Dict[str, Any]) -> None: self.__dict__.update(state) self._init_internal_attributes() def _try_add_agent_graph_to_metadata( metadata: Dict[str, Any], callback_context: callback_context.CallbackContext ) -> None: current_agent: Optional[google.adk.agents.BaseAgent] = ( callback_context_info_extractors.try_get_current_agent_instance( callback_context ) ) if current_agent is None: return try: metadata["_opik_graph_definition"] = { "format": "mermaid", "data": mermaid_graph_builder.build_mermaid_graph_definition( current_agent.root_agent ), } except Exception: LOGGER.error("Failed to build mermaid graph for agent.", exc_info=True) def _log_tool_context_warning(context: tool_context.ToolContext) -> None: if context is not None: warning = f"Function call id: {context.function_call_id}, agent name: {context.agent_name}" if context.actions is not None: warning += f", is escalate: {context.actions.escalate}, transfer to: {context.actions.transfer_to_agent}" LOGGER.warning(warning)