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

opik/integrations/adk/patchers/llm_response_wrapper.py

@@ -0,0 +1,105 @@
+import dataclasses
+import logging
+from typing import Any, Callable, Dict, Optional
+
+from google.adk import models as adk_models
+from google.genai import types as genai_types
+
+from .. import helpers as adk_helpers
+
+import opik
+from opik import llm_usage
+from opik.llm_usage import opik_usage
+
+LOGGER = logging.getLogger(__name__)
+
+
+class LlmResponseCreateWrapper:
+    def __init__(
+        self,
+        wrapped_response_create: Callable[
+            [genai_types.GenerateContentResponse], adk_models.LlmResponse
+        ],
+    ) -> None:
+        self.wrapped_response_create = wrapped_response_create
+
+    def __call__(
+        self, generate_content_response: genai_types.GenerateContentResponse
+    ) -> adk_models.LlmResponse:
+        LOGGER.debug("LlmResponseCreateWrapper called")
+        return _wrap_llm_response_create(
+            generate_content_response, self.wrapped_response_create
+        )
+
+
+@dataclasses.dataclass
+class LLMUsageData:
+    opik_usage: opik_usage.OpikUsage
+    model: Optional[str]
+    provider: Optional[str]
+
+
+def pop_llm_usage_data(
+    result_dict: Dict[str, Any], provider: Optional[str]
+) -> Optional[LLMUsageData]:
+    """Extracts Opik usage metadata from ADK output and removes it from the result dict."""
+    opik_usage_metadata = None
+    model = None
+
+    if (custom_metadata := result_dict.get("custom_metadata", None)) is not None:
+        if (opik_usage_metadata := custom_metadata.pop("opik_usage", None)) is not None:
+            model = custom_metadata.pop("model_version", None)
+            if model is not None:
+                model = model.split("/")[-1]
+
+    # in streaming mode ADK returns the usage metadata in the result dict as the last call
+    # to the after_model_callback bypassing our patching (no opik_usage)
+    if opik_usage_metadata is None:
+        if "usage_metadata" in result_dict:
+            opik_usage_metadata = result_dict["usage_metadata"]
+        else:
+            return None
+
+    if provider in [opik.LLMProvider.GOOGLE_AI, opik.LLMProvider.GOOGLE_VERTEXAI]:
+        usage = llm_usage.try_build_opik_usage_or_log_error(
+            provider=opik.LLMProvider(provider),
+            usage=opik_usage_metadata,
+            logger=LOGGER,
+            error_message="Failed to log token usage from ADK Gemini call",
+        )
+    else:
+        usage = llm_usage.build_opik_usage_from_unknown_provider(
+            usage=opik_usage_metadata,
+        )
+
+    if usage is None:
+        return None
+
+    return LLMUsageData(opik_usage=usage, model=model, provider=provider)
+
+
+def _wrap_llm_response_create(
+    generate_content_response: genai_types.GenerateContentResponse,
+    wrapped_response_create: Callable[
+        [genai_types.GenerateContentResponse], adk_models.LlmResponse
+    ],
+) -> adk_models.LlmResponse:
+    usage_metadata = generate_content_response.usage_metadata
+    LOGGER.debug("_wrap_llm_response_create: usage_metadata=%s", usage_metadata)
+
+    response = wrapped_response_create(generate_content_response)
+    if usage_metadata is None:
+        return response
+
+    if response.custom_metadata is None:
+        response.custom_metadata = {}
+
+    response.custom_metadata["opik_usage"] = usage_metadata
+    response.custom_metadata["provider"] = adk_helpers.get_adk_provider()
+    response.custom_metadata["model_version"] = generate_content_response.model_version
+    LOGGER.debug(
+        "_wrap_llm_response_create finished: response.custom_metadata=%s",
+        response.custom_metadata,
+    )
+
+    return response

opik/integrations/adk/patchers/patchers.py

@@ -0,0 +1,64 @@
+import logging
+
+from opik import semantic_version
+from opik.api_objects import opik_client
+from . import llm_response_wrapper
+from . import litellm_wrappers
+from .adk_otel_tracer import opik_adk_otel_tracer
+
+import google.adk.models
+from google.adk.models import lite_llm
+from google.adk import telemetry as adk_telemetry
+from google.adk.agents import base_agent
+
+LOGGER = logging.getLogger(__name__)
+
+
+def patch_adk(opik_client: opik_client.Opik) -> None:
+    _patch_adk_lite_llm()
+    if semantic_version.SemanticVersion.parse(google.adk.__version__) >= "1.3.0":  # type: ignore
+        _patch_adk_opentelemetry_tracers(opik_client)
+
+
+def _patch_adk_lite_llm() -> None:
+    old_function = google.adk.models.LlmResponse.create
+    if not hasattr(old_function, "opik_patched"):
+        create_wrapper = llm_response_wrapper.LlmResponseCreateWrapper(old_function)
+        google.adk.models.LlmResponse.create = create_wrapper
+        google.adk.models.LlmResponse.create.opik_patched = True  # type: ignore
+        LOGGER.debug("Patched LlmResponse.create")
+
+    if hasattr(lite_llm, "LiteLLMClient") and hasattr(
+        lite_llm.LiteLLMClient, "acompletion"
+    ):
+        if not hasattr(lite_llm.LiteLLMClient.acompletion, "opik_patched"):
+            lite_llm.LiteLLMClient.acompletion = (
+                litellm_wrappers.litellm_client_acompletion_decorator(
+                    lite_llm.LiteLLMClient.acompletion
+                )
+            )
+            lite_llm.LiteLLMClient.acompletion.opik_patched = True  # type: ignore
+            LOGGER.debug("Patched LiteLLMClient.acompletion")
+
+    if hasattr(lite_llm, "_model_response_to_generate_content_response"):
+        if not hasattr(
+            lite_llm._model_response_to_generate_content_response, "opik_patched"
+        ):
+            lite_llm._model_response_to_generate_content_response = (
+                litellm_wrappers.generate_content_response_decorator(
+                    lite_llm._model_response_to_generate_content_response
+                )
+            )
+            lite_llm._model_response_to_generate_content_response.opik_patched = True  # type: ignore
+            LOGGER.debug("Patched _model_response_to_generate_content_response")
+
+
+def _patch_adk_opentelemetry_tracers(opik_client: opik_client.Opik) -> None:
+    no_op_opik_tracer = opik_adk_otel_tracer.OpikADKOtelTracer(opik_client)
+
+    adk_telemetry.tracer.start_as_current_span = no_op_opik_tracer.start_as_current_span
+    adk_telemetry.tracer.start_span = no_op_opik_tracer.start_span
+
+    base_agent.tracer.start_as_current_span = no_op_opik_tracer.start_as_current_span
+    base_agent.tracer.start_span = no_op_opik_tracer.start_span
+    LOGGER.debug("Patched ADK tracers")

opik/integrations/adk/recursive_callback_injector.py

@@ -0,0 +1,126 @@
+import types
+from typing import TypeVar, List, Any, Set
+from . import opik_tracer
+import logging
+
+from google.adk.tools import agent_tool
+from google.adk import agents
+
+LOGGER = logging.getLogger(__name__)
+
+ADKAgent = TypeVar("ADKAgent", bound=agents.BaseAgent)
+
+
+class RecursiveCallbackInjector:
+    def __init__(self, tracer: opik_tracer.OpikTracer) -> None:
+        self._opik_tracer = tracer
+        self._seen_instance_ids: Set[int] = set()
+
+    def inject(
+        self,
+        root_agent: agents.BaseAgent,
+    ) -> None:
+        self._process_agent(root_agent)
+
+    def _add_callbacks_to_agent(self, agent: agents.BaseAgent) -> None:
+        callback_fields = {
+            "before_agent_callback": self._opik_tracer.before_agent_callback,
+            "after_agent_callback": self._opik_tracer.after_agent_callback,
+            "before_model_callback": self._opik_tracer.before_model_callback,
+            "after_model_callback": self._opik_tracer.after_model_callback,
+            "before_tool_callback": self._opik_tracer.before_tool_callback,
+            "after_tool_callback": self._opik_tracer.after_tool_callback,
+        }
+
+        for callback_field_name, callback_func in callback_fields.items():
+            if not hasattr(agent, callback_field_name):
+                continue
+
+            current_callback_value = getattr(agent, callback_field_name)
+            if current_callback_value is None:
+                setattr(agent, callback_field_name, callback_func)
+            elif isinstance(
+                current_callback_value, list
+            ) and not self._contains_opik_tracer_callback(
+                callbacks=current_callback_value
+            ):
+                current_callback_value.append(callback_func)
+            elif not self._is_opik_callback_function(
+                current_callback_value
+            ) and callable(current_callback_value):
+                setattr(
+                    agent, callback_field_name, [current_callback_value, callback_func]
+                )
+
+    def _process_agent(
+        self,
+        agent: agents.BaseAgent,
+    ) -> None:
+        if id(agent) in self._seen_instance_ids:
+            return
+
+        self._add_callbacks_to_agent(agent)
+        self._process_sub_agents(agent)
+        self._process_tools(agent)
+
+        self._seen_instance_ids.add(id(agent))
+
+    def _process_sub_agents(
+        self,
+        agent: agents.BaseAgent,
+    ) -> None:
+        if not hasattr(agent, "sub_agents"):
+            return
+
+        for sub_agent in agent.sub_agents:
+            try:
+                self._process_agent(sub_agent)
+            except Exception as e:
+                LOGGER.warning(f"Failed to track subagent: {e}")
+
+    def _process_tools(
+        self,
+        agent: agents.BaseAgent,
+    ) -> None:
+        if not hasattr(agent, "tools"):
+            return
+
+        for tool in agent.tools:
+            if not isinstance(tool, agent_tool.AgentTool):
+                continue
+            try:
+                self._process_agent(tool.agent)
+            except Exception as e:
+                LOGGER.warning(f"Failed to track agent tool: {e}")
+
+    def _is_opik_callback_function(self, obj: Any) -> bool:
+        if not callable(obj):
+            return False
+
+        if isinstance(obj, types.MethodType):
+            return isinstance(obj.__self__, type(self._opik_tracer))
+
+        return False
+
+    def _contains_opik_tracer_callback(self, callbacks: List) -> bool:
+        return any(self._is_opik_callback_function(callback) for callback in callbacks)
+
+
+def track_adk_agent_recursive(
+    root_agent: ADKAgent, tracer: opik_tracer.OpikTracer
+) -> ADKAgent:
+    """
+    Recursively adds opik tracer callbacks to the agent, its subagents, and agent tools.
+
+    Args:
+        root_agent: The root ADK agent to track
+        tracer: The OpikTracer instance to use for tracking
+
+    Returns:
+        The modified root agent with tracking enabled
+    """
+
+    recursive_callback_injector = RecursiveCallbackInjector(tracer)
+    recursive_callback_injector.inject(root_agent)
+
+    return root_agent

opik/integrations/aisuite/aisuite_decorator.py

@@ -10,11 +10,13 @@ from typing import (
     Tuple,
     Union,
 )
+from typing_extensions import override
 
 import aisuite.framework as aisuite_chat_completion
 from openai.types.chat import chat_completion as openai_chat_completion
 
-from opik import dict_utils, llm_usage
+import opik.dict_utils as dict_utils
+import opik.llm_usage as llm_usage
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 from opik.types import LLMProvider
@@ -30,12 +32,13 @@ class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator):
     calls of AISuite's `chat.completion.create`
     """
 
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
         track_options: arguments_helpers.TrackOptions,
-        args: Optional[Tuple],
-        kwargs: Optional[Dict[str, Any]],
+        args: Tuple,
+        kwargs: Dict[str, Any],
     ) -> arguments_helpers.StartSpanParameters:
         assert (
             kwargs is not None
@@ -94,6 +97,7 @@ class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator):
 
         return model, provider
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -150,6 +154,7 @@ class AISuiteTrackDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(
         self,
         output: Any,

opik/integrations/aisuite/opik_tracker.py

@@ -23,6 +23,7 @@ def track_aisuite(
     Returns:
         The modified AISuite client with Opik tracking enabled.
     """
+
     if hasattr(aisuite_client, "opik_tracked"):
         return aisuite_client
 

opik/integrations/anthropic/messages_create_decorator.py

@@ -3,8 +3,10 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import anthropic
 from anthropic.types import Message as AnthropicMessage
+from typing_extensions import override
 
-from opik import dict_utils, llm_usage
+import opik.dict_utils as dict_utils
+import opik.llm_usage as llm_usage
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 
@@ -25,12 +27,13 @@ class AnthropicMessagesCreateDecorator(base_track_decorator.BaseTrackDecorator):
     def __init__(self, provider: str) -> None:
         self.provider: str = provider
 
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
         track_options: arguments_helpers.TrackOptions,
-        args: Optional[Tuple],
-        kwargs: Optional[Dict[str, Any]],
+        args: Tuple,
+        kwargs: Dict[str, Any],
     ) -> arguments_helpers.StartSpanParameters:
         assert (
             kwargs is not None
@@ -57,6 +60,7 @@ class AnthropicMessagesCreateDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Union[str, AnthropicMessage],
@@ -88,6 +92,7 @@ class AnthropicMessagesCreateDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(
         self,
         output: Any,

opik/integrations/anthropic/opik_tracker.py

@@ -47,7 +47,6 @@ def track_anthropic(
     Returns:
         Anthropic client with integrated Opik tracking logic.
     """
-
     if hasattr(anthropic_client, "opik_tracked"):
         return anthropic_client
 

opik/integrations/bedrock/converse/chunks_aggregator.py

@@ -0,0 +1,188 @@
+import logging
+from typing import Any, Dict, List
+
+LOGGER = logging.getLogger(__name__)
+
+
+def _handle_message_start(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract role from messageStart event."""
+    message_start = event.get("messageStart")
+    if isinstance(message_start, dict):
+        role = message_start.get("role")
+        if role:
+            result["output"]["message"]["role"] = role
+
+
+def _handle_content_block_delta(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """
+    Extract content from contentBlockDelta event.
+
+    Handles multiple delta types:
+    - delta.text: Regular text streaming
+    - delta.toolUse: Structured output / tool calls (Issue #3829)
+    """
+    content_block_delta = event.get("contentBlockDelta")
+    if not isinstance(content_block_delta, dict):
+        return
+
+    delta = content_block_delta.get("delta")
+    if not isinstance(delta, dict):
+        return
+
+    content = result["output"]["message"]["content"][0]
+
+    # Handle regular text streaming
+    if "text" in delta:
+        content["text"] += delta["text"]
+        return
+
+    # Handle structured output / tool use (Issue #3829)
+    # Ref: https://github.com/comet-ml/opik/issues/3829
+    if "toolUse" in delta:
+        if "toolUse" not in content:
+            content["toolUse"] = {}
+        content["toolUse"].update(delta["toolUse"])
+        return
+
+    # Log other delta types for future compatibility
+    LOGGER.debug("Unknown delta type in contentBlockDelta: %s", list(delta.keys()))
+
+
+def _handle_message_stop(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract stopReason from messageStop event."""
+    message_stop = event.get("messageStop")
+    if isinstance(message_stop, dict):
+        stop_reason = message_stop.get("stopReason")
+        if stop_reason:
+            result["stopReason"] = stop_reason
+
+
+def _handle_metadata(event: Dict[str, Any], result: Dict[str, Any]) -> None:
+    """Extract usage and metrics from metadata event."""
+    metadata = event.get("metadata")
+    if not isinstance(metadata, dict):
+        return
+
+    # Extract usage information
+    if "usage" in metadata:
+        result["usage"] = metadata["usage"]
+
+    # Extract metrics information
+    if "metrics" in metadata:
+        metrics = metadata["metrics"]
+        if isinstance(metrics, dict) and "latencyMs" in metrics:
+            result["metrics"] = {"latencyMs": metrics["latencyMs"]}
+
+
+def aggregate_converse_stream_chunks(items: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Aggregate streaming chunks from AWS Bedrock converse_stream API into a single response.
+
+    This function handles various event structures from different Bedrock model providers:
+    - Anthropic (Claude): Standard messageStart, contentBlockDelta, messageStop events
+    - Amazon (Nova): Amazon's proprietary event format with potential variations
+    - Meta (Llama): Open-source model events with different tokenization patterns
+    - Mistral (Pixtral): Multimodal model events that may include additional content types
+    - DeepSeek (R1): Reasoning model events with extended thought processes (OPIK-2910 fix)
+
+    Event Structure Variations by Provider:
+    ========================================
+
+    Standard Converse Stream Events (from AWS documentation):
+    - messageStart: Contains role and initial metadata
+    - contentBlockStart: Marks beginning of content block
+    - contentBlockDelta: Contains incremental text in delta.text
+    - contentBlockStop: Marks end of content block
+    - messageStop: Contains stopReason (e.g., "end_turn", "stop_sequence")
+    - metadata: Contains usage stats and performance metrics
+
+    Known Variations:
+    - DeepSeek R1: May have different delta structure or additional reasoning fields (OPIK-2910)
+    - Tool Use/Structured Output: delta.toolUse instead of delta.text (Issue #3829)
+    - Multimodal models: May include non-text content blocks (images, documents)
+    - Different models: Varying field names, nesting levels, or optional fields
+
+    References:
+    - AWS Bedrock Converse API: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_Converse.html
+    - Streaming Events: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStreamOutput.html
+    - ContentBlockDelta: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ContentBlockDelta.html
+    - Tool Use Guide: https://docs.aws.amazon.com/bedrock/latest/userguide/tool-use.html
+    - DeepSeek R1 Issue: https://comet-ml.atlassian.net/browse/OPIK-2910
+    - Tool Use Issue: https://github.com/comet-ml/opik/issues/3829
+
+    Args:
+        items: List of streaming event dictionaries from Bedrock converse_stream
+
+    Returns:
+        Aggregated response dictionary with structure:
+        {
+            "output": {
+                "message": {
+                    "role": "assistant",
+                    "content": [{"text": "aggregated text"}]
+                }
+            },
+            "stopReason": "end_turn",  # Optional
+            "usage": {...},  # Optional
+            "metrics": {"latencyMs": ...}  # Optional
+        }
+    """
+
+    result: Dict[str, Any] = {
+        "output": {"message": {"role": "assistant", "content": [{"text": ""}]}}
+    }
+
+    for event in items:
+        if not isinstance(event, dict):
+            LOGGER.debug("Skipping non-dictionary event: %s", type(event))
+            continue
+
+        try:
+            if "messageStart" in event:
+                _handle_message_start(event, result)
+
+            if "contentBlockDelta" in event:
+                _handle_content_block_delta(event, result)
+
+            if "messageStop" in event:
+                _handle_message_stop(event, result)
+
+            if "metadata" in event:
+                _handle_metadata(event, result)
+
+        except Exception as e:
+            LOGGER.warning(
+                "Unexpected error processing event: %s. Event: %s",
+                str(e),
+                event,
+                exc_info=True,
+            )
+
+    return result
+
+
+def aggregate_invoke_agent_chunks(items: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """
+    Aggregate streaming chunks from AWS Bedrock invoke_agent API.
+
+    Note: The implementation uses a simplified approach as the completion payload
+    only contains chunks without additional metadata (as of implementation date).
+
+    Reference:
+    - https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/invoke_agent.html
+
+    Args:
+        items: List of chunk event dictionaries from invoke_agent stream
+
+    Returns:
+        Aggregated response dictionary with decoded text output
+    """
+    merged_chunks = b""
+
+    for item in items:
+        if isinstance(item, dict) and "chunk" in item:
+            chunk = item["chunk"]
+            if isinstance(chunk, dict) and "bytes" in chunk:
+                merged_chunks += chunk["bytes"]
+
+    return {"output": merged_chunks.decode("utf-8")}

opik/integrations/bedrock/{converse_decorator.py → converse/converse_decorator.py}

@@ -1,11 +1,15 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing_extensions import override
 
-from opik import dict_utils, llm_usage
+import opik
+import opik.dict_utils as dict_utils
+import opik.llm_usage as llm_usage
 from opik.api_objects import span
 from opik.decorator import arguments_helpers, base_track_decorator
 
-from . import helpers, stream_wrappers
+from . import stream_wrappers
+from .. import types
 
 LOGGER = logging.getLogger(__name__)
 
@@ -22,12 +26,13 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
     overrides _generators_handler() method to work correctly with bedrock's streams
     """
 
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
         track_options: arguments_helpers.TrackOptions,
-        args: Optional[Tuple],
-        kwargs: Optional[Dict[str, Any]],
+        args: Tuple,
+        kwargs: Dict[str, Any],
     ) -> arguments_helpers.StartSpanParameters:
         assert (
             kwargs is not None
@@ -47,19 +52,23 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
             tags=tags,
             metadata=metadata,
             project_name=track_options.project_name,
+            model=kwargs.get("modelId", None),
+            provider=opik.LLMProvider.BEDROCK,
         )
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
         capture_output: bool,
         current_span_data: span.SpanData,
     ) -> arguments_helpers.EndSpanParameters:
-        usage = output["usage"]
+        usage = output.get("usage", {})
+
         usage_in_openai_format = llm_usage.try_build_opik_usage_or_log_error(
-            provider="_bedrock",
+            provider=opik.LLMProvider.BEDROCK,
             usage=usage,
             logger=LOGGER,
             error_message="Failed to log token usage from bedrock LLM call",
@@ -76,13 +85,14 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(  # type: ignore
         self,
         output: Any,
         capture_output: bool,
         generations_aggregator: Optional[Callable[[List[Any]], Any]],
     ) -> Union[
-        helpers.ConverseStreamOutput,
+        types.ConverseStreamOutput,
         None,
     ]:
         DECORATED_FUNCTION_IS_NOT_EXPECTED_TO_RETURN_GENERATOR = (
@@ -108,7 +118,7 @@ class BedrockConverseDecorator(base_track_decorator.BaseTrackDecorator):
             )
 
             output["stream"] = wrapped_stream
-            return cast(helpers.ConverseStreamOutput, output)
+            return cast(types.ConverseStreamOutput, output)
 
         STREAM_NOT_FOUND = None