opik 1.6.4__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/integrations/adk/opik_tracer.py

@@ -0,0 +1,370 @@
+import logging
+from typing import Any, Dict, List, Optional, 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()
+
+        patchers.patch_adk(self._opik_client)
+
+    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
+            # Debug logging for callback invocation
+            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.
+            _, span_data = 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(span_data)
+        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:
+            # Ignore partial chunks, ADK will call this method with the full response at the end
+            if llm_response.partial is True:
+                return
+
+        except Exception:
+            LOGGER.debug("Error checking for partial chunks", exc_info=True)
+
+        try:
+            model = None
+            usage = None
+            output = None
+
+            if adk_helpers.has_empty_text_part_content(llm_response):
+                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
+
+            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,
+                )
+
+            current_span.update(
+                output=output,
+                name=model or current_span.model,
+                type="llm",
+                model=model,
+                usage=usage,
+                project_name=self.project_name,
+            )
+            if current_span.metadata is not None:
+                current_span.metadata[llm_span_helpers.SPAN_STATUS] = (
+                    llm_span_helpers.LLMSpanStatus.READY_FOR_FINALIZATION
+                )
+
+            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:
+            LOGGER.error(f"Failed during after_model_callback(): {e}", exc_info=True)
+
+    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)
+        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)

opik/integrations/adk/patchers/__init__.py

@@ -0,0 +1,4 @@
+from .patchers import patch_adk
+
+
+__all__ = ["patch_adk"]

opik/integrations/adk/patchers/adk_otel_tracer/llm_span_helpers.py

@@ -0,0 +1,30 @@
+import enum
+from opik.api_objects import span
+
+SPAN_STATUS = "_OPIK_SPAN_STATUS"
+
+
+class LLMSpanStatus(str, enum.Enum):
+    STARTED = "started"
+    READY_FOR_FINALIZATION = "ready_for_finalization"
+
+
+def is_externally_created_llm_span_ready_for_immediate_finalization(
+    span_data: span.SpanData,
+) -> bool:
+    return (
+        span_data.type == "llm"
+        and span_data.metadata is not None
+        and span_data.metadata.get(SPAN_STATUS, None)
+        == LLMSpanStatus.READY_FOR_FINALIZATION
+    )
+
+
+def is_externally_created_llm_span_that_just_started(
+    span_data: span.SpanData,
+) -> bool:
+    return (
+        span_data.type == "llm"
+        and span_data.metadata is not None
+        and span_data.metadata.get(SPAN_STATUS, None) == LLMSpanStatus.STARTED
+    )

opik/integrations/adk/patchers/adk_otel_tracer/opik_adk_otel_tracer.py

@@ -0,0 +1,201 @@
+import logging
+from typing import Iterator, Optional, Tuple
+
+import opentelemetry.trace
+import opik
+import opik.context_storage
+from opik.api_objects import trace, span
+from opik.decorator import (
+    span_creation_handler,
+    arguments_helpers,
+    error_info_collector,
+)
+from opik.api_objects import opik_client
+
+from . import llm_span_helpers
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class OpikADKOtelTracer(opentelemetry.trace.NoOpTracer):
+    """
+    A patched OpenTelemetry tracer for ADK integration.
+
+    This tracer intercepts OpenTelemetry span creation calls from ADK and converts them
+    into Opik spans and traces, providing seamless integration between ADK's telemetry
+    system and Opik's tracking capabilities.
+
+    Key Features:
+    - Converts OpenTelemetry spans to Opik spans and traces
+    - Handles trace lifecycle management automatically
+    - Skips internal ADK spans that shouldn't have Opik spans created for them
+    - Maintains context storage for nested span relationships
+    - Provides special logic for handling LLM spans and their lifecycle. In particular:
+        * LLM spans are created by OpikTracer.before_model_callback
+        * No other Opik spans are created inside LLM spans (ADK creates tool spans inside LLM spans)
+    - Supports the flows in which Opik spans are being created, updated or finalized outside of this class.
+
+    The tracer operates as a no-op from OpenTelemetry's perspective (returning INVALID_SPAN)
+    while creating corresponding Opik tracking objects in the background.
+
+    Args:
+        opik_client: The Opik client instance used for creating spans and traces
+
+    Attributes:
+        _ADK_INTERNAL_SPAN_NAME_SKIP_LIST: List of span names that should be skipped
+            during tracking.
+    """
+
+    _ADK_INTERNAL_SPAN_NAME_SKIP_LIST = [
+        "invocation",  # This is the span that is created by ADK when the invocation is started.
+    ]
+
+    def __init__(self, opik_client: opik_client.Opik):
+        self.opik_client = opik_client
+
+    def start_span(
+        self,
+        name: str,
+        context: Optional[opentelemetry.trace.Context] = None,
+        kind: opentelemetry.trace.SpanKind = opentelemetry.trace.SpanKind.INTERNAL,
+        attributes: opentelemetry.trace.types.Attributes = None,
+        links: opentelemetry.trace._Links = None,
+        start_time: Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+    ) -> "opentelemetry.trace.Span":
+        return opentelemetry.trace.INVALID_SPAN
+
+    @opentelemetry.util._decorator._agnosticcontextmanager
+    def start_as_current_span(
+        self,
+        name: str,
+        context: Optional[opentelemetry.trace.Context] = None,
+        kind: opentelemetry.trace.SpanKind = opentelemetry.trace.SpanKind.INTERNAL,
+        attributes: opentelemetry.trace.types.Attributes = None,
+        links: opentelemetry.trace._Links = None,
+        start_time: Optional[int] = None,
+        record_exception: bool = True,
+        set_status_on_exception: bool = True,
+        end_on_exit: bool = True,
+    ) -> Iterator["opentelemetry.trace.Span"]:
+        if name in self._ADK_INTERNAL_SPAN_NAME_SKIP_LIST:
+            yield opentelemetry.trace.INVALID_SPAN
+            return
+
+        trace_to_close_in_finally_block = None
+        span_to_close_in_finally_block = None
+
+        current_trace_data = opik.context_storage.get_trace_data()
+        current_span_data = opik.context_storage.top_span_data()
+
+        if (
+            current_span_data is not None
+            and llm_span_helpers.is_externally_created_llm_span_ready_for_immediate_finalization(
+                current_span_data
+            )
+        ):
+            # LLM span has already been updated from the OpikTracer.after_model_call.
+            # ADK will create tool spans inside the llm spans (which we consider wrong practice),
+            # so we manually finalize it here to avoid incorrect span nesting.
+            opik.context_storage.pop_span_data(ensure_id=current_span_data.id)
+            current_span_data.init_end_time()
+            if opik.is_tracing_active():
+                self.opik_client.span(**current_span_data.as_parameters)
+            current_span_data = opik.context_storage.top_span_data()
+
+        try:
+            trace_to_close_in_finally_block, span_to_close_in_finally_block = (
+                _prepare_trace_and_span_to_be_finalized(
+                    name=name,
+                    current_trace_data=current_trace_data,
+                    current_span_data=current_span_data,
+                )
+            )
+
+            yield opentelemetry.trace.INVALID_SPAN
+        except Exception as exception:
+            # The expected exception here is the exception that happened during the agent
+            # execution and was re-raised from the `opentelemetry.util._decorator._agnosticcontextmanagergenerator`s
+            # `__exit__` method via the `gen.throw(...)` statement.
+            #
+            # More context: https://docs.python.org/3.12/reference/expressions.html#examples
+            error_info = error_info_collector.collect(exception)
+
+            if trace_to_close_in_finally_block is not None:
+                trace_to_close_in_finally_block.update(error_info=error_info)
+            if span_to_close_in_finally_block is not None:
+                span_to_close_in_finally_block.update(error_info=error_info)
+            raise
+        finally:
+            if trace_to_close_in_finally_block is not None:
+                self._ensure_trace_is_finalized(trace_to_close_in_finally_block.id)
+
+            if span_to_close_in_finally_block is not None:
+                self._ensure_span_is_finalized(span_to_close_in_finally_block.id)
+
+            if (
+                trace_to_close_in_finally_block is None
+                and span_to_close_in_finally_block is None
+            ):
+                LOGGER.warning(
+                    "No span or trace to finalize in ADK tracer. This is unexpected."
+                )
+
+    def _ensure_trace_is_finalized(self, trace_id: str) -> None:
+        trace_data = opik.context_storage.pop_trace_data(ensure_id=trace_id)
+        if trace_data is not None:
+            trace_data.init_end_time()
+            if opik.is_tracing_active():
+                self.opik_client.trace(**trace_data.as_parameters)
+
+    def _ensure_span_is_finalized(self, span_id: str) -> None:
+        opik.context_storage.trim_span_data_stack_to_certain_span(span_id)
+
+        span_data = opik.context_storage.pop_span_data(ensure_id=span_id)
+        if span_data is not None:
+            span_data.init_end_time()
+            if opik.is_tracing_active():
+                self.opik_client.span(**span_data.as_parameters)
+
+
+def _prepare_trace_and_span_to_be_finalized(
+    name: str,
+    current_trace_data: Optional[trace.TraceData],
+    current_span_data: Optional[span.SpanData],
+) -> Tuple[Optional[trace.TraceData], Optional[span.SpanData]]:
+    """
+    Prepares a trace and a span to be finalized in the finally block.
+    """
+
+    trace_to_close_in_finally_block = None
+    span_to_close_in_finally_block = None
+
+    if current_trace_data is None:
+        trace_to_close_in_finally_block = trace.TraceData(name=name)
+        opik.context_storage.set_trace_data(trace_to_close_in_finally_block)
+    elif (
+        current_span_data is not None
+        and llm_span_helpers.is_externally_created_llm_span_that_just_started(
+            current_span_data
+        )
+    ):
+        # LLM span has just been created and put into context storage from the OpikTracer.before_model_call
+        # Not need to create a new one, just remember it to close it in finally block
+        span_to_close_in_finally_block = current_span_data
+    else:
+        start_span_arguments = arguments_helpers.StartSpanParameters(
+            name=name,
+            type="general",
+        )
+
+        _, span_to_close_in_finally_block = (
+            span_creation_handler.create_span_respecting_context(
+                start_span_arguments=start_span_arguments,
+                distributed_trace_headers=None,
+            )
+        )
+        opik.context_storage.add_span_data(span_to_close_in_finally_block)
+
+    return trace_to_close_in_finally_block, span_to_close_in_finally_block

opik/integrations/adk/patchers/litellm_wrappers.py

@@ -0,0 +1,91 @@
+import logging
+from functools import wraps
+from typing import Any, Callable, Optional, Tuple, Union, TYPE_CHECKING
+
+from google.adk import models as adk_models
+
+import opik
+
+if TYPE_CHECKING:
+    import litellm
+
+LOGGER = logging.getLogger(__name__)
+
+
+def parse_provider_and_model(
+    model: str,
+) -> Tuple[Optional[Union[opik.LLMProvider, str]], str]:
+    parts = model.split("/", 1)
+    if len(parts) != 2:
+        return None, parts[0]
+
+    provider = parts[0]
+    try:
+        provider = opik.LLMProvider(provider)
+    except ValueError:
+        pass
+    return provider, parts[1]
+
+
+def generate_content_response_decorator(func: Callable) -> Callable:
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> adk_models.LlmResponse:
+        """
+        This wrapper puts token usage data into custom metadata to use it later
+        """
+        LOGGER.debug("generate_content_response_decorator called")
+        result = func(*args, **kwargs)
+        model_response = args[0]
+        model_response_dict = model_response.to_dict()
+
+        provider_and_model = model_response_dict.get("provider_and_model", None)
+        LOGGER.debug(
+            "generate_content_response_decorator: provider_and_model=%s",
+            provider_and_model,
+        )
+
+        if provider_and_model is None:
+            return result
+
+        if result.custom_metadata is None:
+            LOGGER.debug(
+                "generate_content_response_decorator: result.custom_metadata is None, creating new custom metadata"
+            )
+            result.custom_metadata = {}
+
+        provider, model = parse_provider_and_model(provider_and_model)
+        LOGGER.debug(
+            "generate_content_response_decorator: provider=%s, model=%s",
+            provider,
+            model,
+        )
+
+        result.custom_metadata["opik_usage"] = model_response_dict["usage"]
+        result.custom_metadata["provider"] = provider
+        result.custom_metadata["model_version"] = model_response_dict.get(
+            "model", model
+        )
+        LOGGER.debug(
+            "generate_content_response_decorator finished: result.custom_metadata=%s",
+            result.custom_metadata,
+        )
+        return result
+
+    return wrapper
+
+
+def litellm_client_acompletion_decorator(func: Callable) -> Callable:
+    @wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> "litellm.types.utils.ModelResponse":
+        """
+        this adds more precise provider/model name and it's version
+        """
+        result = await func(*args, **kwargs)
+        result.provider_and_model = kwargs.get("model", None)
+        LOGGER.debug(
+            "litellm_client_acompletion_decorator called: provider_and_model=%s",
+            result.provider_and_model,
+        )
+        return result
+
+    return wrapper