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

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.
@@ -22,8 +34,10 @@ def track_bedrock(
     Returns:
         The modified bedrock client with Opik tracking enabled.
     """
+
     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"
@@ -32,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
@@ -53,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

@@ -10,11 +10,13 @@ from typing import (
     Tuple,
     Union,
 )
+from typing_extensions import override
 
-from opik import llm_usage
 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.dict_utils as dict_utils
 
 LOGGER = logging.getLogger(__name__)
 
@@ -78,6 +80,7 @@ TASK_KWARGS_KEYS_TO_LOG_AS_OUTPUT = [
 
 
 class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
@@ -123,7 +126,9 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
             metadata["object_type"] = "agent"
             agent = args[0]
             input_dict = {"context": kwargs.get("context")}
-            agent_dict = agent.model_dump(include=AGENT_KWARGS_KEYS_TO_LOG_AS_INPUTS)
+            agent_dict = _encode_dict_and_keep_keys(
+                agent, AGENT_KWARGS_KEYS_TO_LOG_AS_INPUTS
+            )
             input_dict["agent"] = agent_dict
             name = agent.role.strip()
 
@@ -131,18 +136,16 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
         elif name == "execute_sync":
             metadata["object_type"] = "task"
             input_dict = {}
-            task_dict = args[0].model_dump(include=TASK_KWARGS_KEYS_TO_LOG_AS_INPUTS)
+            task = args[0]
+            task_dict = _encode_dict_and_keep_keys(
+                task, TASK_KWARGS_KEYS_TO_LOG_AS_INPUTS
+            )
             input_dict["task"] = task_dict
-            name = f"Task: {args[0].name}"
-
-        elif name == "completion":
-            metadata["object_type"] = "completion"
-            input_dict = {"messages": kwargs.get("messages")}
-            span_type = "llm"
-            name = "llm call"
+            name = f"Task: {task.name}"
 
         return input_dict, name, span_type
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -151,64 +154,30 @@ 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)
+        if object_type == "crew":
+            output_dict = jsonable_encoder.encode(output)
+            output_dict.pop("token_usage", None)
+        elif object_type == "agent":
+            output_dict = {"output": output}
+        elif object_type == "task":
+            output_dict = _encode_dict_and_keep_keys(
+                output, TASK_KWARGS_KEYS_TO_LOG_AS_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 = output.model_dump()
-            _ = output_dict.pop("token_usage")
-        elif object_type == "agent":
-            output_dict = {"output": output}
-        elif object_type == "task":
-            output_dict = output.model_dump(include=TASK_KWARGS_KEYS_TO_LOG_AS_OUTPUT)
-        elif object_type == "completion":
-            output_dict = output.model_dump()
-            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
-            )
-            output_dict = {}
-
-        return model, provider, output_dict, usage
-
+    @override
     def _streams_handler(
         self,
         output: Any,
@@ -216,3 +185,12 @@ class CrewAITrackDecorator(base_track_decorator.BaseTrackDecorator):
         generations_aggregator: Optional[Callable[[List[Any]], str]],
     ) -> Optional[Union[Generator, AsyncGenerator]]:
         return super()._streams_handler(output, capture_output, generations_aggregator)
+
+
+def _encode_dict_and_keep_keys(
+    dict: Dict[str, Any],
+    keys_to_keep: List[str],
+) -> Dict[str, Any]:
+    encoded_dict = jsonable_encoder.encode(dict)
+    encoded_dict, _ = dict_utils.split_dict_by_keys(encoded_dict, keys_to_keep)
+    return encoded_dict

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)

opik/integrations/crewai/patchers/llm_client.py

@@ -0,0 +1,207 @@
+"""
+Client patchers for CrewAI LLM providers (v1.0.0+).
+
+This module handles patching of LLM clients used by CrewAI agents with Opik tracking.
+Each provider has its own patching function that handles missing dependencies gracefully.
+"""
+
+import logging
+from typing import Any, Optional, TYPE_CHECKING
+
+import crewai
+
+if TYPE_CHECKING:
+    import crewai.llms.providers.openai.completion as openai_completion
+    import crewai.llms.providers.anthropic.completion as anthropic_completion
+    import crewai.llms.providers.gemini.completion as gemini_completion
+    import crewai.llms.providers.bedrock.completion as bedrock_completion
+
+LOGGER = logging.getLogger(__name__)
+
+
+def patch_llm_client(crew: crewai.Crew, project_name: Optional[str]) -> None:
+    """
+    Patches LLM clients used by CrewAI agents with Opik tracking.
+
+    Handles missing provider libraries gracefully by logging warnings instead of failing.
+
+    Args:
+        crew: The Crew instance containing agents to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    for agent in crew.agents:
+        _patch_single_llm_client(agent.llm, project_name)
+
+
+def _patch_single_llm_client(llm: Any, project_name: Optional[str]) -> None:
+    """
+    Patches an LLM client based on its provider type.
+
+    Args:
+        llm: The CrewAI LLM instance to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    if _is_openai_llm(llm):
+        _patch_openai_client(llm, project_name)
+    elif _is_anthropic_llm(llm):
+        _patch_anthropic_client(llm, project_name)
+    elif _is_gemini_llm(llm):
+        _patch_gemini_client(llm, project_name)
+    elif _is_bedrock_llm(llm):
+        _patch_bedrock_client(llm, project_name)
+
+
+def _is_openai_llm(llm: Any) -> bool:
+    """
+    Checks if LLM is an OpenAI provider.
+
+    Args:
+        llm: The CrewAI LLM to check.
+
+    Returns:
+        True if LLM is OpenAI provider, False otherwise.
+    """
+    try:
+        import crewai.llms.providers.openai.completion
+
+        return isinstance(llm, crewai.llms.providers.openai.completion.OpenAICompletion)
+    except ImportError:
+        return False
+
+
+def _is_anthropic_llm(llm: Any) -> bool:
+    """
+    Checks if LLM is an Anthropic provider.
+
+    Args:
+        llm: The CrewAI LLM to check.
+
+    Returns:
+        True if LLM is Anthropic provider, False otherwise.
+    """
+    try:
+        import crewai.llms.providers.anthropic.completion
+
+        return isinstance(
+            llm, crewai.llms.providers.anthropic.completion.AnthropicCompletion
+        )
+    except ImportError:
+        return False
+
+
+def _is_gemini_llm(llm: Any) -> bool:
+    """
+    Checks if LLM is a Gemini provider.
+
+    Args:
+        llm: The CrewAI LLM to check.
+
+    Returns:
+        True if LLM is Gemini provider, False otherwise.
+    """
+    try:
+        import crewai.llms.providers.gemini.completion
+
+        return isinstance(llm, crewai.llms.providers.gemini.completion.GeminiCompletion)
+    except ImportError:
+        return False
+
+
+def _is_bedrock_llm(llm: Any) -> bool:
+    """
+    Checks if LLM is a Bedrock provider.
+
+    Args:
+        llm: The CrewAI LLM to check.
+
+    Returns:
+        True if LLM is Bedrock provider, False otherwise.
+    """
+    try:
+        import crewai.llms.providers.bedrock.completion
+
+        return isinstance(
+            llm, crewai.llms.providers.bedrock.completion.BedrockCompletion
+        )
+    except ImportError:
+        return False
+
+
+def _patch_openai_client(
+    llm: "openai_completion.OpenAICompletion", project_name: Optional[str]
+) -> None:
+    """
+    Patches OpenAI client for the given LLM.
+
+    Args:
+        llm: The CrewAI LLM instance with OpenAI client to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    try:
+        import opik.integrations.openai
+
+        llm.client = opik.integrations.openai.track_openai(
+            llm.client, project_name=project_name
+        )
+    except Exception:
+        LOGGER.warning("Failed to track OpenAI client for LLM", exc_info=True)
+
+
+def _patch_anthropic_client(
+    llm: "anthropic_completion.AnthropicCompletion", project_name: Optional[str]
+) -> None:
+    """
+    Patches Anthropic client for the given LLM.
+
+    Args:
+        llm: The CrewAI LLM instance with Anthropic client to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    try:
+        import opik.integrations.anthropic
+
+        llm.client = opik.integrations.anthropic.track_anthropic(
+            llm.client, project_name=project_name
+        )
+    except Exception:
+        LOGGER.warning("Failed to track Anthropic client for LLM", exc_info=True)
+
+
+def _patch_gemini_client(
+    llm: "gemini_completion.GeminiCompletion", project_name: Optional[str]
+) -> None:
+    """
+    Patches Gemini client for the given LLM.
+
+    Args:
+        llm: The CrewAI LLM instance with Gemini client to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    try:
+        import opik.integrations.genai
+
+        llm.client = opik.integrations.genai.track_genai(
+            llm.client, project_name=project_name
+        )
+    except Exception:
+        LOGGER.warning("Failed to track Gemini client for LLM", exc_info=True)
+
+
+def _patch_bedrock_client(
+    llm: "bedrock_completion.BedrockCompletion", project_name: Optional[str]
+) -> None:
+    """
+    Patches Bedrock client for the given LLM.
+
+    Args:
+        llm: The CrewAI LLM instance with Bedrock client to patch.
+        project_name: The name of the project to associate with tracking.
+    """
+    try:
+        import opik.integrations.bedrock
+
+        llm.client = opik.integrations.bedrock.track_bedrock(
+            llm.client, project_name=project_name
+        )
+    except Exception:
+        LOGGER.warning("Failed to track Bedrock client for LLM", exc_info=True)