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

opik/integrations/bedrock/invoke_model/usage_converters.py

@@ -0,0 +1,87 @@
+"""Shared usage conversion utilities for Bedrock models."""
+
+from typing import Any, Dict
+
+
+def anthropic_to_bedrock_usage(anthropic_usage: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert Anthropic-style usage schema into Bedrock-style usage schema.
+
+    Anthropic usage keys (snake_case):
+      - input_tokens
+      - output_tokens
+      - cache_creation_input_tokens
+      - cache_read_input_tokens
+
+    Bedrock usage keys (camelCase):
+      - inputTokens
+      - outputTokens
+      - cacheWriteInputTokens
+      - cacheReadInputTokens
+      - totalTokens
+    """
+    input_tokens = anthropic_usage.get("input_tokens", 0)
+    output_tokens = anthropic_usage.get("output_tokens", 0)
+    cache_write = anthropic_usage.get("cache_creation_input_tokens", 0)
+    cache_read = anthropic_usage.get("cache_read_input_tokens", 0)
+
+    return {
+        "inputTokens": input_tokens,
+        "outputTokens": output_tokens,
+        "cacheWriteInputTokens": cache_write,
+        "cacheReadInputTokens": cache_read,
+        "totalTokens": input_tokens + output_tokens,
+    }
+
+
+def llama_to_bedrock_usage(llama_usage: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert Llama-style usage schema into Bedrock-style usage schema.
+
+    Llama usage keys:
+      - prompt_token_count
+      - generation_token_count
+    """
+    input_tokens = llama_usage.get("prompt_token_count", 0)
+    output_tokens = llama_usage.get("generation_token_count", 0)
+
+    return {
+        "inputTokens": input_tokens,
+        "outputTokens": output_tokens,
+        "totalTokens": input_tokens + output_tokens,
+    }
+
+
+def openai_to_bedrock_usage(openai_usage: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert OpenAI-style usage schema into Bedrock-style usage schema.
+    Used by Mistral/Pixtral models.
+
+    OpenAI usage keys:
+      - prompt_tokens
+      - completion_tokens
+      - total_tokens
+    """
+    input_tokens = openai_usage.get("prompt_tokens", 0)
+    output_tokens = openai_usage.get("completion_tokens", 0)
+
+    return {
+        "inputTokens": input_tokens,
+        "outputTokens": output_tokens,
+        "totalTokens": input_tokens + output_tokens,
+    }
+
+
+def nova_to_bedrock_usage(nova_usage: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Convert Nova-style usage (already in Bedrock format) - pass through.
+    Nova already uses Bedrock format (inputTokens, outputTokens).
+    """
+    input_tokens = nova_usage.get("inputTokens", 0)
+    output_tokens = nova_usage.get("outputTokens", 0)
+
+    return {
+        "inputTokens": input_tokens,
+        "outputTokens": output_tokens,
+        "totalTokens": input_tokens + output_tokens,
+    }

opik/integrations/bedrock/invoke_model/usage_extraction.py

@@ -0,0 +1,108 @@
+from opik import llm_usage
+from typing import Dict, Any, Optional
+import logging
+import opik._logging
+
+from . import usage_converters
+
+LOGGER = logging.getLogger(__name__)
+
+
+def extract_subprovider_from_model_id(model_id: str) -> str:
+    """
+    Extracts the provider name from a Bedrock modelId.
+
+    Examples:
+        ai21.j2-mid-v1                -> ai21
+        amazon.nova-lite-v1:0         -> amazon
+        anthropic.claude-v2:1         -> anthropic
+        us.meta.llama3-1-70b-instruct -> meta
+    """
+    parts = model_id.split(".")
+
+    if parts[0] in {"us", "eu", "apac"}:
+        return parts[1]
+
+    return parts[0]
+
+
+def try_extract_usage_from_bedrock_response(
+    subprovider: str, response: Dict[str, Any]
+) -> Optional[llm_usage.OpikUsage]:
+    """
+    Since Bedrock's invoke_model response format is not standardized, we need to try different ways to extract the usage.
+
+    This usage may also be not in Bedrock's format, but in the format of the original subprovider.
+    """
+    try:
+        LOGGER.debug("Extracting usage for subprovider: %s", subprovider)
+
+        if subprovider == "anthropic":
+            usage_dict = response["body"]["usage"]
+            bedrock_formatted_usage = usage_converters.anthropic_to_bedrock_usage(
+                usage_dict
+            )
+            opik_usage = llm_usage.OpikUsage.from_bedrock_dict(bedrock_formatted_usage)
+            LOGGER.debug("Anthropic usage extracted: %s", bedrock_formatted_usage)
+            return opik_usage
+
+        elif subprovider == "meta":
+            # Llama models have usage fields directly in body (not in body.usage)
+            body = response.get("body", {})
+            if "prompt_token_count" in body or "generation_token_count" in body:
+                bedrock_formatted_usage = usage_converters.llama_to_bedrock_usage(body)
+                opik_usage = llm_usage.OpikUsage.from_bedrock_dict(
+                    bedrock_formatted_usage
+                )
+                LOGGER.debug("Llama usage extracted: %s", bedrock_formatted_usage)
+                return opik_usage
+
+        elif subprovider == "mistral":
+            # Mistral/Pixtral models use OpenAI-like usage format
+            usage_dict = response["body"].get("usage", {})
+            if usage_dict:
+                bedrock_formatted_usage = usage_converters.openai_to_bedrock_usage(
+                    usage_dict
+                )
+                opik_usage = llm_usage.OpikUsage.from_bedrock_dict(
+                    bedrock_formatted_usage
+                )
+                LOGGER.debug("Mistral usage extracted: %s", bedrock_formatted_usage)
+                return opik_usage
+
+        elif subprovider == "amazon":
+            # Nova models already use Bedrock format
+            usage_dict = response["body"].get("usage", {})
+            if usage_dict:
+                bedrock_formatted_usage = usage_converters.nova_to_bedrock_usage(
+                    usage_dict
+                )
+                opik_usage = llm_usage.OpikUsage.from_bedrock_dict(
+                    bedrock_formatted_usage
+                )
+                LOGGER.debug("Nova usage extracted: %s", bedrock_formatted_usage)
+                return opik_usage
+
+        # Fallback: This is the default case, but it's not guaranteed to find the usage here for all possible subproviders
+        presumably_usage_dict = response["body"].get("usage", {})
+        if presumably_usage_dict:
+            # If it's already in Bedrock's format, we are good (tested with amazon.nova-pro-v1:0, it has bedrock usage format)
+            # If it's not, but it's in some other format that Opik supports, we will at least extract
+            # completion and prompt tokens count so that backend could calculate cost based on them.
+            opik_usage = llm_usage.build_opik_usage_from_unknown_provider(
+                presumably_usage_dict
+            )
+            LOGGER.debug("Fallback usage extracted: %s", presumably_usage_dict)
+            return opik_usage
+
+        LOGGER.debug("No usage found in response body")
+        return None
+
+    except Exception as e:
+        LOGGER.debug("Exception during usage extraction: %s", e)
+        opik._logging.log_once_at_level(
+            logging.WARNING,
+            f"Failed to extract usage from Bedrock's invoke_model response: {response}. It may be because this model response format is currently not supported: please create an issue at https://github.com/opik-ai/opik/issues and we will add support for it.",
+            LOGGER,
+        )
+        return None

opik/integrations/bedrock/opik_tracker.py

@@ -1,6 +1,12 @@
 from typing import Optional, TYPE_CHECKING
 
-from . import chunks_aggregator, converse_decorator, invoke_agent_decorator
+from . import invoke_agent_decorator
+from .converse import chunks_aggregator as converse_chunks_aggregator
+from .converse import converse_decorator
+
+from .invoke_model import invoke_model_decorator
+from .invoke_model import chunks_aggregator as invoke_model_chunks_aggregator
+
 
 if TYPE_CHECKING:
     from mypy_boto3_bedrock_runtime.client import BedrockRuntimeClient
@@ -12,9 +18,15 @@ def track_bedrock(
 ) -> "BedrockRuntimeClient":
     """Adds Opik tracking to an AWS Bedrock client.
 
-    Tracks calls to `converse()` and `converse_stream()` methods
+    Tracks calls to `converse()`, `converse_stream()`, `invoke_model()`, and `invoke_model_with_response_stream()` methods.
     Can be used within other Opik-tracked functions.
 
+    Supported Model subproviders for InvokeModel API (both streaming and non-streaming):
+    - **Anthropic** (Claude)
+    - **Amazon** (Nova)
+    - **Meta** (Llama)
+    - **Mistral** (Pixtral)
+
     Args:
         client: An instance of an AWS Bedrock client (botocore.client.BedrockRuntime or botocore.client.AgentsforBedrockRuntime).
         project_name: The name of the project to log data.
@@ -25,6 +37,7 @@ def track_bedrock(
 
     decorator_for_converse = converse_decorator.BedrockConverseDecorator()
     decorator_for_invoke_agent = invoke_agent_decorator.BedrockInvokeAgentDecorator()
+    decorator_for_invoke_model = invoke_model_decorator.BedrockInvokeModelDecorator()
 
     if hasattr(client, "invoke_agent") and not hasattr(
         client.invoke_agent, "opik_tracked"
@@ -33,7 +46,7 @@ def track_bedrock(
             type="llm",
             name="bedrock_invoke_agent",
             project_name=project_name,
-            generations_aggregator=chunks_aggregator.aggregate_invoke_agent_chunks,
+            generations_aggregator=converse_chunks_aggregator.aggregate_invoke_agent_chunks,
         )
         tracked_invoke_agent = wrapper(client.invoke_agent)
         client.invoke_agent = tracked_invoke_agent
@@ -54,9 +67,34 @@ def track_bedrock(
             type="llm",
             name="bedrock_converse_stream",
             project_name=project_name,
-            generations_aggregator=chunks_aggregator.aggregate_converse_stream_chunks,
+            generations_aggregator=converse_chunks_aggregator.aggregate_converse_stream_chunks,
         )
         tracked_converse_stream = stream_wrapper(client.converse_stream)
         client.converse_stream = tracked_converse_stream
 
+    if hasattr(client, "invoke_model") and not hasattr(
+        client.invoke_model, "opik_tracked"
+    ):
+        wrapper = decorator_for_invoke_model.track(
+            type="llm",
+            name="bedrock_invoke_model",
+            project_name=project_name,
+        )
+        tracked_invoke_model = wrapper(client.invoke_model)
+        client.invoke_model = tracked_invoke_model
+
+    if hasattr(client, "invoke_model_with_response_stream") and not hasattr(
+        client.invoke_model_with_response_stream, "opik_tracked"
+    ):
+        stream_wrapper = decorator_for_invoke_model.track(
+            type="llm",
+            name="bedrock_invoke_model_stream",
+            project_name=project_name,
+            generations_aggregator=invoke_model_chunks_aggregator.aggregate_chunks_to_dataclass,
+        )
+        tracked_invoke_model_stream = stream_wrapper(
+            client.invoke_model_with_response_stream
+        )
+        client.invoke_model_with_response_stream = tracked_invoke_model_stream
+
     return client

opik/integrations/bedrock/types.py

@@ -0,0 +1,19 @@
+from typing import Any, Dict, TypedDict
+
+from botocore import eventstream
+import botocore.response
+
+
+class ConverseStreamOutput(TypedDict):
+    stream: eventstream.EventStream
+    ResponseMetadata: Dict[str, Any]
+
+
+class InvokeModelOutput(TypedDict):
+    body: botocore.response.StreamingBody
+    ResponseMetadata: Dict[str, Any]
+
+
+class InvokeModelWithResponseStreamOutput(TypedDict):
+    body: eventstream.EventStream
+    ResponseMetadata: Dict[str, Any]

opik/integrations/crewai/crewai_decorator.py

@@ -13,10 +13,9 @@ from typing import (
 from typing_extensions import override
 
 from opik.decorator import arguments_helpers, base_track_decorator
-from opik.types import SpanType, LLMProvider
+from opik.types import SpanType
 from opik.api_objects import span
 import opik.jsonable_encoder as jsonable_encoder
-import opik.llm_usage as llm_usage
 import opik.dict_utils as dict_utils
 
 LOGGER = logging.getLogger(__name__)
@@ -144,12 +143,6 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
             input_dict["task"] = task_dict
             name = f"Task: {task.name}"
 
-        elif name == "completion":
-            metadata["object_type"] = "completion"
-            input_dict = {"messages": kwargs.get("messages")}
-            span_type = "llm"
-            name = "llm call"
-
         return input_dict, name, span_type
 
     @override
@@ -161,38 +154,12 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
     ) -> arguments_helpers.EndSpanParameters:
         object_type = None
         metadata = {}
+        output_dict = {}
 
         if current_span_data and current_span_data.metadata:
             metadata = current_span_data.metadata
             object_type = metadata.pop("object_type")
 
-        model, provider, output_dict, usage = self._parse_outputs(object_type, output)
-
-        result = arguments_helpers.EndSpanParameters(
-            output=output_dict,
-            usage=usage,
-            metadata=metadata,
-            model=model,
-            provider=provider,
-        )
-
-        return result
-
-    def _parse_outputs(
-        self,
-        object_type: Optional[str],
-        output: Any,
-    ) -> Tuple[
-        Optional[str],
-        Optional[str],
-        Dict[str, Any],
-        Optional[llm_usage.OpikUsage],
-    ]:
-        model = None
-        provider = None
-        usage = None
-        output_dict = {}
-
         if object_type == "crew":
             output_dict = jsonable_encoder.encode(output)
             output_dict.pop("token_usage", None)
@@ -202,23 +169,13 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
             output_dict = _encode_dict_and_keep_keys(
                 output, TASK_KWARGS_KEYS_TO_LOG_AS_OUTPUT
             )
-        elif object_type == "completion":
-            output_dict = jsonable_encoder.encode(output)
-            if output_dict.get("usage", None) is not None:
-                usage = llm_usage.try_build_opik_usage_or_log_error(
-                    provider=LLMProvider.OPENAI,  # even if it's not openai, we know the format is openai-like
-                    usage=output_dict["usage"],
-                    logger=LOGGER,
-                    error_message="Failed to log token usage from CrewAI LLM call",
-                )
-            else:
-                usage = None
-            model = output_dict.pop("model", None)
-            provider = (
-                "openai" if output_dict.get("object") == "chat.completion" else None
-            )
 
-        return model, provider, output_dict, usage
+        result = arguments_helpers.EndSpanParameters(
+            output=output_dict,
+            metadata=metadata,
+        )
+
+        return result
 
     @override
     def _streams_handler(

opik/integrations/crewai/opik_tracker.py

@@ -1,14 +1,19 @@
+import importlib.metadata
+import logging
 from typing import Optional
 
 import crewai
 
-from . import crewai_decorator
+import opik.semantic_version
 
-__IS_TRACKING_ENABLED = False
+from . import crewai_decorator, patchers
+
+LOGGER = logging.getLogger(__name__)
 
 
 def track_crewai(
     project_name: Optional[str] = None,
+    crew: Optional[crewai.Crew] = None,
 ) -> None:
     """
     Tracks CrewAI activities by enabling tracking decorators for various critical methods.
@@ -21,11 +26,8 @@ def track_crewai(
 
     Parameters:
         project_name: The name of the project to associate with the tracking.
+        crew: The Crew instance to track. Required for CrewAI v1.0.0+ to properly track LLM calls.
     """
-    global __IS_TRACKING_ENABLED
-    if __IS_TRACKING_ENABLED:
-        return
-    __IS_TRACKING_ENABLED = True
 
     decorator_factory = crewai_decorator.CrewAITrackDecorator()
 
@@ -33,12 +35,31 @@ def track_crewai(
         project_name=project_name,
     )
 
-    import litellm
-
     crewai.Crew.kickoff = crewai_wrapper(crewai.Crew.kickoff)
     crewai.Crew.kickoff_for_each = crewai_wrapper(crewai.Crew.kickoff_for_each)
     crewai.Agent.execute_task = crewai_wrapper(crewai.Agent.execute_task)
     crewai.Task.execute_sync = crewai_wrapper(crewai.Task.execute_sync)
-    litellm.completion = crewai_wrapper(litellm.completion)
 
-    return None
+    # Patch LiteLLM functions used by CrewAI
+    patchers.patch_litellm_completion(project_name=project_name)
+
+    # Patch Flow class (v1.0.0+)
+    patchers.patch_flow(project_name=project_name)
+
+    # Patch LLM clients used by CrewAI agents (v1.0.0+)
+    if crew is not None and is_crewai_v1():
+        patchers.patch_llm_client(crew, project_name)
+
+
+def is_crewai_v1() -> bool:
+    """
+    Checks if CrewAI v1.0.0+ is installed.
+
+    Returns:
+        True if CrewAI v1.0.0+ is detected, False otherwise.
+    """
+    try:
+        version_str = importlib.metadata.version("crewai")
+        return opik.semantic_version.SemanticVersion.parse(version_str) >= "1.0.0"  # type: ignore
+    except Exception:
+        return False

opik/integrations/crewai/patchers/__init__.py

@@ -0,0 +1,5 @@
+from .flow import patch_flow
+from .llm_client import patch_llm_client
+from .litellm_completion import patch_litellm_completion
+
+__all__ = ["patch_flow", "patch_llm_client", "patch_litellm_completion"]

opik/integrations/crewai/patchers/flow.py

@@ -0,0 +1,118 @@
+"""
+Patcher for CrewAI Flow class (v1.0.0+).
+
+This module patches the Flow class to automatically track flow methods and execution.
+"""
+
+import functools
+import logging
+from typing import Optional
+
+import opik.decorator.tracker as opik_tracker
+
+LOGGER = logging.getLogger(__name__)
+
+
+def patch_flow(project_name: Optional[str] = None) -> None:
+    """
+    Patches CrewAI Flow class to track flow execution.
+
+    If Flow class is not available (CrewAI < v1.0.0), this function does nothing.
+
+    Args:
+        project_name: The name of the project to associate with tracking.
+    """
+    _patch_flow_init(project_name)
+    _patch_flow_kickoff_async(project_name)
+
+
+def _patch_flow_init(project_name: Optional[str] = None) -> None:
+    """
+    Patches CrewAI Flow.__init__ to automatically track flow methods.
+
+    If Flow class is not available (CrewAI < v1.0.0), this function does nothing.
+    """
+    try:
+        import crewai
+
+        if not hasattr(crewai, "Flow"):
+            LOGGER.debug("CrewAI Flow class not available, skipping Flow patching")
+            return
+
+        if hasattr(_patch_flow_init, "_patched"):
+            return
+
+        original_init = crewai.Flow.__init__
+
+        @functools.wraps(original_init)
+        def _init_wrapper(self, *args, **kwargs) -> None:  # type: ignore
+            original_init(self, *args, **kwargs)
+
+            try:
+                flow_registered_methods = getattr(self, "_methods", {})
+                for method_name, method in list(flow_registered_methods.items()):
+                    if getattr(method, "opik_tracked", False):
+                        continue
+
+                    decorated = opik_tracker.track(
+                        project_name=project_name,
+                        tags=["crewai"],
+                        metadata={"created_from": "crewai"},
+                    )(method)
+
+                    flow_registered_methods[method_name] = decorated
+            except Exception:
+                LOGGER.error(
+                    "An error occurred during Opik instrumentation of CrewAI Flow",
+                    exc_info=True,
+                )
+
+        crewai.Flow.__init__ = _init_wrapper  # type: ignore[assignment]
+
+        setattr(_patch_flow_init, "_patched", True)  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):
+        LOGGER.debug(
+            "CrewAI Flow class not available, skipping Flow patching", exc_info=True
+        )
+
+
+def _patch_flow_kickoff_async(project_name: Optional[str] = None) -> None:
+    """
+    Patches CrewAI Flow.kickoff_async to track flow execution.
+
+    If Flow class is not available (CrewAI < v1.0.0), this function does nothing.
+    """
+    try:
+        import crewai
+
+        if not hasattr(crewai, "Flow"):
+            LOGGER.debug(
+                "CrewAI Flow class not available, skipping Flow.kickoff_async patching"
+            )
+            return
+
+        if hasattr(_patch_flow_kickoff_async, "_patched"):
+            return
+
+        # We only need to patch the async version of the kickoff method because
+        # the sync version calls it internally
+        original_kickoff_async = crewai.Flow.kickoff_async
+
+        @functools.wraps(original_kickoff_async)
+        async def _kickoff_async_wrapper(self, *args, **kwargs):  # type: ignore
+            wrapped = opik_tracker.track(
+                project_name=project_name,
+                tags=["crewai"],
+                name="Flow.kickoff_async",
+                metadata={"created_from": "crewai"},
+            )(original_kickoff_async)
+            return await wrapped(self, *args, **kwargs)
+
+        crewai.Flow.kickoff_async = _kickoff_async_wrapper  # type: ignore[assignment]
+
+        setattr(_patch_flow_kickoff_async, "_patched", True)  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):
+        LOGGER.debug(
+            "CrewAI Flow class not available, skipping Flow.kickoff_async patching",
+            exc_info=True,
+        )

opik/integrations/crewai/patchers/litellm_completion.py

@@ -0,0 +1,30 @@
+"""
+Patcher for LiteLLM completion functions used by CrewAI.
+
+This module patches litellm.completion and litellm.acompletion with Opik tracking.
+CrewAI v0.x uses LiteLLM internally for LLM calls.
+"""
+
+import logging
+from typing import Optional
+
+import litellm
+
+import opik.integrations.litellm
+
+LOGGER = logging.getLogger(__name__)
+
+
+def patch_litellm_completion(project_name: Optional[str] = None) -> None:
+    """
+    Patches LiteLLM completion functions used by CrewAI.
+
+    Args:
+        project_name: The name of the project to associate with tracking.
+    """
+    litellm.completion = opik.integrations.litellm.track_completion(
+        project_name=project_name
+    )(litellm.completion)
+    litellm.acompletion = opik.integrations.litellm.track_completion(
+        project_name=project_name
+    )(litellm.acompletion)