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

opik/integrations/dspy/parsers.py

@@ -0,0 +1,168 @@
+"""
+Parsers and data structures for extracting information from DSPy LM responses.
+
+This module contains utilities for parsing DSPy LM history entries and
+extracting relevant information like usage, provider, and cost data.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Tuple
+import logging
+
+import dspy
+
+from opik import llm_usage, types
+
+LOGGER = logging.getLogger(__name__)
+
+
+@dataclass
+class LMHistoryInfo:
+    """
+    Information extracted from a DSPy LM history entry.
+
+    This dataclass holds the parsed information from an LM call's history,
+    including usage statistics, cache status, provider information, and cost.
+
+    Attributes:
+        usage: Token usage information (prompt, completion, total tokens)
+        cache_hit: Whether the response was served from cache.
+            True if cached, False if not, None if unknown.
+        actual_provider: The actual provider that served the request.
+            This is useful for LLM routers like OpenRouter that may route
+            to different underlying providers (e.g., "Novita", "Together").
+        total_cost: The total cost of the request from the provider.
+            This includes accurate pricing for all token types.
+    """
+
+    usage: Optional[llm_usage.OpikUsage]
+    cache_hit: Optional[bool]
+    actual_provider: Optional[str]
+    total_cost: Optional[float]
+
+    def as_tuple(
+        self,
+    ) -> Tuple[
+        Optional[llm_usage.OpikUsage],
+        Optional[bool],
+        Optional[str],
+        Optional[float],
+    ]:
+        """Return the info as a tuple for backwards compatibility."""
+        return (self.usage, self.cache_hit, self.actual_provider, self.total_cost)
+
+
+def get_span_type(instance: Any) -> types.SpanType:
+    """
+    Determine the span type based on the DSPy instance type.
+
+    Args:
+        instance: A DSPy module, LM, or tool instance.
+
+    Returns:
+        The appropriate span type: "llm" for Predict/LM, "tool" for Tool,
+        or "general" for other types.
+    """
+    if isinstance(instance, dspy.Predict):
+        return "llm"
+    elif isinstance(instance, dspy.LM):
+        return "llm"
+    elif isinstance(instance, dspy.Tool):
+        return "tool"
+    return "general"
+
+
+def extract_lm_info_from_history(
+    lm_instance: Any,
+    expected_messages: Optional[Any],
+) -> LMHistoryInfo:
+    """
+    Extract token usage, cache status, actual provider, and cost from the LM's history.
+
+    DSPy stores usage information in the LM's history after each call.
+    We verify the history entry matches our expected messages to handle
+    potential race conditions with concurrent LM calls.
+
+    For routers like OpenRouter, the response contains the actual provider
+    that served the request (e.g., "Novita", "Together"), which differs from
+    the router name used in the model string (e.g., "openrouter").
+
+    The cost field is provided by providers like OpenRouter and includes
+    accurate pricing for all token types (reasoning, cache, multimodal).
+
+    Args:
+        lm_instance: The DSPy LM instance that has the history.
+        expected_messages: The expected messages to match in the history entry.
+
+    Returns:
+        LMHistoryInfo containing usage, cache_hit, actual_provider, and total_cost.
+    """
+    empty_result = LMHistoryInfo(
+        usage=None,
+        cache_hit=None,
+        actual_provider=None,
+        total_cost=None,
+    )
+
+    if not hasattr(lm_instance, "history") or not lm_instance.history:
+        return empty_result
+
+    try:
+        last_entry = lm_instance.history[-1]
+
+        # Verify we have the correct history entry by checking messages match
+        if last_entry.get("messages") != expected_messages:
+            LOGGER.debug(
+                "History entry messages don't match expected messages, "
+                "skipping usage extraction (possibly due to concurrent LM calls)"
+            )
+            return empty_result
+
+        response = last_entry.get("response")
+        usage_dict = last_entry.get("usage")
+
+        # Extract actual provider from response (useful for routers like OpenRouter)
+        # The response is a LiteLLM ModelResponse object with a 'provider' attribute
+        # when using routers like OpenRouter
+        actual_provider: Optional[str] = None
+        if response is not None and hasattr(response, "provider"):
+            actual_provider = response.provider
+
+        # Extract cost from history entry or usage dict
+        # OpenRouter and other providers return accurate cost including all token types
+        total_cost: Optional[float] = None
+        if last_entry.get("cost") is not None:
+            total_cost = last_entry.get("cost")
+        elif usage_dict and usage_dict.get("cost") is not None:
+            total_cost = usage_dict.get("cost")
+
+        # Get explicit cache_hit if set, otherwise infer from usage (empty = cached)
+        if response is None:
+            cache_hit = not usage_dict
+        elif hasattr(response, "cache_hit") and response.cache_hit is not None:
+            cache_hit = response.cache_hit
+        else:
+            # Fallback: infer from usage (empty = cached)
+            cache_hit = not usage_dict
+
+        if usage_dict:
+            usage = llm_usage.build_opik_usage_from_unknown_provider(usage_dict)
+            return LMHistoryInfo(
+                usage=usage,
+                cache_hit=cache_hit,
+                actual_provider=actual_provider,
+                total_cost=total_cost,
+            )
+        else:
+            return LMHistoryInfo(
+                usage=None,
+                cache_hit=cache_hit,
+                actual_provider=actual_provider,
+                total_cost=total_cost,
+            )
+    except Exception:
+        LOGGER.debug(
+            "Failed to extract info from DSPy LM history",
+            exc_info=True,
+        )
+        return empty_result

opik/integrations/genai/encoder_extension.py

@@ -2,16 +2,12 @@ import base64
 from typing import Any, Dict, Union
 
 from google.genai import types as genai_types
-from opik import jsonable_encoder
+import opik.jsonable_encoder as jsonable_encoder
 
 
 def register() -> None:
     def encoder_extension(obj: genai_types.Blob) -> Union[str, Dict[str, Any]]:
-        if (
-            obj.mime_type is not None
-            and obj.data is not None
-            and obj.mime_type.startswith("image")
-        ):
+        if obj.mime_type is not None and obj.data is not None:
             return {
                 "data": base64.b64encode(obj.data).decode("utf-8"),
                 "mime_type": obj.mime_type,

opik/integrations/genai/generate_content_decorator.py

@@ -10,10 +10,12 @@ from typing import (
     Tuple,
     Union,
 )
+from typing_extensions import override
 
 from google.genai import types as genai_types
 
-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
@@ -41,18 +43,22 @@ class GenerateContentTrackDecorator(base_track_decorator.BaseTrackDecorator):
         super().__init__()
         self.provider = 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
         ), "Expected kwargs to be not None in client.models.generate_content(**kwargs), client.aio.models.generate_content(**kwargs)"
 
+        model = kwargs.get("model")
+
         name = track_options.name if track_options.name is not None else func.__name__
+        name = f"{name}: {model}"  # Add model to the name for better viewing UX
 
         metadata = track_options.metadata if track_options.metadata is not None else {}
 
@@ -71,12 +77,13 @@ class GenerateContentTrackDecorator(base_track_decorator.BaseTrackDecorator):
             tags=tags,
             metadata=metadata,
             project_name=track_options.project_name,
-            model=kwargs.get("model", None),
+            model=model,
             provider=self.provider,
         )
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -93,7 +100,11 @@ class GenerateContentTrackDecorator(base_track_decorator.BaseTrackDecorator):
             result_dict, RESPONSE_KEYS_TO_LOG_AS_OUTPUT
         )
 
-        model = result_dict["model_version"]
+        if result_dict.get("model_version") is not None:
+            # Gemini **may** add "models/" prefix to some model versions
+            model = result_dict["model_version"].split("/")[-1]
+        else:
+            model = None
 
         usage = llm_usage.try_build_opik_usage_or_log_error(
             provider=LLMProvider(self.provider),
@@ -101,9 +112,9 @@ class GenerateContentTrackDecorator(base_track_decorator.BaseTrackDecorator):
             logger=LOGGER,
             error_message="Failed to log token usage from genai generate_response call",
         )
-
+        span_name_without_model = current_span_data.name.split(":")[0]  # type: ignore
         result = arguments_helpers.EndSpanParameters(
-            name=f"{current_span_data.name}: {model}",  # Add model to the name for better viewing UX
+            name=f"{span_name_without_model}: {model}",
             output=output,
             usage=usage,
             metadata=metadata,
@@ -113,16 +124,12 @@ class GenerateContentTrackDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(
         self,
         output: Any,
         capture_output: bool,
-        generations_aggregator: Optional[
-            Callable[
-                [List[genai_types.GenerateContentResponse]],
-                genai_types.GenerateContentResponse,
-            ]
-        ],
+        generations_aggregator: Optional[Callable[[List[Any]], Any]],
     ) -> Union[
         None,
         Iterator[genai_types.GenerateContentResponse],

opik/integrations/guardrails/guardrails_decorator.py

@@ -10,6 +10,7 @@ from typing import (
     Tuple,
     Union,
 )
+from typing_extensions import override
 
 from guardrails import validators
 
@@ -23,6 +24,7 @@ RESPONSE_KEYS_TO_LOG_AS_OUTPUT = ["output"]
 
 
 class GuardrailsValidatorValidateDecorator(base_track_decorator.BaseTrackDecorator):
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
@@ -59,6 +61,7 @@ class GuardrailsValidatorValidateDecorator(base_track_decorator.BaseTrackDecorat
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -79,6 +82,7 @@ class GuardrailsValidatorValidateDecorator(base_track_decorator.BaseTrackDecorat
 
         return result
 
+    @override
     def _streams_handler(
         self,
         output: Any,

opik/integrations/harbor/__init__.py

@@ -0,0 +1,17 @@
+"""
+Opik integration for Harbor benchmark evaluation framework.
+
+Example:
+    >>> from opik.integrations.harbor import track_harbor
+    >>> job = Job(config)
+    >>> tracked_job = track_harbor(job)
+    >>> result = await tracked_job.run()
+
+Or enable tracking globally (for CLI usage):
+    >>> from opik.integrations.harbor import track_harbor
+    >>> track_harbor()
+"""
+
+from .opik_tracker import track_harbor, reset_harbor_tracking
+
+__all__ = ["track_harbor", "reset_harbor_tracking"]

opik/integrations/harbor/experiment_service.py

@@ -0,0 +1,269 @@
+"""
+Experiment service for Harbor integration.
+
+This module manages the connection between Harbor benchmark runs and Opik experiments,
+enabling evaluation tracking and result visualization.
+
+Harbor Terminology Mapping to Opik:
+-----------------------------------
+- **Harbor Job**: A benchmark run that evaluates one or more agents on a dataset.
+  Maps to an Opik Experiment.
+
+- **Harbor Trial**: A single agent run on a single task within a job.
+  Each trial produces one Opik Trace (capturing the agent's execution).
+
+- **Harbor Source**: The benchmark dataset being used (e.g., "terminal-bench", "swe-bench").
+  Maps to an Opik Dataset. Each source gets its own dataset.
+
+- **Harbor Task**: A specific problem/challenge within a dataset (e.g., "fix-git" task).
+  Maps to an Opik Dataset Item.
+
+Flow Overview:
+--------------
+1. When a Harbor job starts, this service is initialized with an experiment name.
+2. For each unique source (benchmark dataset), we create/get an Opik Dataset and Experiment.
+3. For each trial (agent run on a task), we:
+   a. Create a dataset item for the task (or reuse existing one if task was run before)
+   b. Link the trial's trace to the experiment via ExperimentItemReferences
+4. This allows viewing all trial results in Opik's experiment comparison UI.
+"""
+
+import logging
+from datetime import datetime
+from typing import Any, Dict, Optional, Set, TYPE_CHECKING
+
+from opik.api_objects import opik_client
+from opik.api_objects.experiment import experiment_item
+
+if TYPE_CHECKING:
+    from opik.api_objects.experiment.experiment import Experiment
+    from opik.api_objects.dataset.dataset import Dataset
+
+LOGGER = logging.getLogger(__name__)
+
+# Global singleton service instance (one per Harbor job)
+_SERVICE: Optional["HarborExperimentService"] = None
+
+
+class HarborExperimentService:
+    """
+    Manages Opik datasets and experiments for Harbor benchmark jobs.
+
+    This service handles the mapping between Harbor's evaluation structure and Opik's
+    experiment tracking:
+
+    - Each Harbor source (benchmark dataset) → One Opik Dataset + One Opik Experiment
+    - Each Harbor task → One Opik Dataset Item
+    - Each Harbor trial → One Opik Trace, linked to the experiment
+
+    The service uses lazy initialization - datasets and experiments are created
+    on-demand when the first trial for a source is linked.
+
+    Attributes:
+        _experiment_name: Name for experiments created by this service.
+        _experiment_config: Config dict stored on experiments (agent/model info).
+        _client: Cached Opik client instance.
+        _datasets: Map of source name → Opik Dataset.
+        _experiments: Map of source name → Opik Experiment.
+        _linked_trials: Set of trial names already linked (prevents duplicates).
+    """
+
+    def __init__(
+        self,
+        experiment_name: str,
+        experiment_config: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        Initialize the experiment service.
+
+        Args:
+            experiment_name: Name for experiments. Typically includes job_id
+                for uniqueness (e.g., "harbor-job-abc123").
+            experiment_config: Optional config dict to store on experiments.
+                Typically contains agent/model info (e.g., {"agent_name": "terminus",
+                "model_name": "gpt-4o"}).
+        """
+        self._experiment_name = experiment_name
+        self._experiment_config = experiment_config or {}
+        self._experiment_config["created_from"] = "harbor"
+        self._client = opik_client.get_client_cached()
+
+        # Lazy-initialized per source (benchmark dataset)
+        self._datasets: Dict[str, "Dataset"] = {}
+        self._experiments: Dict[str, "Experiment"] = {}
+
+        # Track which trials have been linked to avoid duplicates
+        self._linked_trials: Set[str] = set()
+
+    def _ensure_dataset_and_experiment(self, source: str) -> None:
+        """
+        Ensure an Opik Dataset and Experiment exist for the given source.
+
+        Creates them lazily on first access. Each Harbor source (benchmark dataset)
+        gets its own Opik Dataset and Experiment pair.
+
+        Args:
+            source: The Harbor source/benchmark name (e.g., "terminal-bench").
+        """
+        if source in self._experiments:
+            return
+
+        try:
+            # Create or get the dataset for this benchmark source
+            dataset = self._client.get_or_create_dataset(
+                name=source,
+                description=f"Harbor benchmark dataset: {source}",
+            )
+            self._datasets[source] = dataset
+            LOGGER.info("Using dataset '%s' for Harbor source", source)
+
+            # Create a new experiment for this job run
+            experiment = self._client.create_experiment(
+                dataset_name=source,
+                name=self._experiment_name,
+                experiment_config=self._experiment_config,
+            )
+            self._experiments[source] = experiment
+            LOGGER.info(
+                "Created experiment '%s' for dataset '%s'",
+                self._experiment_name,
+                source,
+            )
+        except Exception as e:
+            LOGGER.warning(
+                "Failed to create dataset/experiment for source '%s': %s",
+                source,
+                e,
+            )
+
+    def link_trial_to_experiment(
+        self,
+        trial_name: str,
+        trace_id: str,
+        source: Optional[str] = None,
+        task_name: Optional[str] = None,
+    ) -> None:
+        """
+        Link a Harbor trial's trace to the Opik experiment.
+
+        This creates the connection between a trial's execution trace and the
+        experiment, enabling the trial to appear in Opik's experiment comparison UI.
+
+        The flow:
+        1. Ensure dataset and experiment exist for the source
+        2. Create or find the dataset item for this task
+        3. Link the trace to the experiment via the dataset item
+
+        Args:
+            trial_name: Unique identifier for the trial (e.g., "task__abc123").
+                Used to prevent duplicate linking.
+            trace_id: The Opik trace ID for this trial's execution.
+            source: The Harbor source/benchmark name. Defaults to "harbor-default".
+            task_name: The task name within the benchmark (e.g., "fix-git").
+                Used to create/find the dataset item.
+        """
+        source = source or "harbor-default"
+
+        # Prevent duplicate linking of the same trial
+        if trial_name in self._linked_trials:
+            return
+
+        # Ensure we have a dataset and experiment for this source
+        self._ensure_dataset_and_experiment(source)
+
+        experiment = self._experiments.get(source)
+        dataset = self._datasets.get(source)
+
+        if experiment is None or dataset is None:
+            LOGGER.warning(
+                "Failed to create experiment/dataset for source '%s', "
+                "trial '%s' will not be linked",
+                source,
+                trial_name,
+            )
+            return
+
+        try:
+            # Insert the task as a dataset item (idempotent - duplicates are handled)
+            dataset.insert([{"task_name": task_name}])
+
+            # Find the dataset item ID for this task.
+            # We search because the same task may have been inserted in a previous run,
+            # and we want to reuse the existing item ID for proper experiment linking.
+            items = dataset.get_items()
+            dataset_item_id = None
+            for item in items:
+                if item.get("task_name") == task_name:
+                    dataset_item_id = item.get("id")
+                    break
+
+            if dataset_item_id is None:
+                LOGGER.warning("Could not find dataset item for task '%s'", task_name)
+                return
+
+            # Link the trace to the experiment via the dataset item
+            experiment.insert(
+                [
+                    experiment_item.ExperimentItemReferences(
+                        dataset_item_id=dataset_item_id,
+                        trace_id=trace_id,
+                    )
+                ]
+            )
+
+            self._linked_trials.add(trial_name)
+            LOGGER.debug(
+                "Linked trial '%s' (trace %s) to experiment '%s'",
+                trial_name,
+                trace_id,
+                self._experiment_name,
+            )
+        except Exception as e:
+            LOGGER.warning("Failed to link trial '%s' to experiment: %s", trial_name, e)
+
+
+def setup_lazy(
+    experiment_name: Optional[str] = None,
+    experiment_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    """
+    Setup the experiment service lazily.
+
+    Called when the first Harbor trial runs. Creates the global service instance
+    that will be used for all subsequent trial linking. Datasets and experiments
+    are created on-demand when trials are linked.
+
+    Args:
+        experiment_name: Name for the experiment. If None, auto-generates
+            a timestamped name like "harbor-20241209-143000".
+        experiment_config: Optional config dict to store on experiments.
+            Typically contains agent/model info (e.g., {"agent_name": "terminus",
+            "model_name": "gpt-4o"}).
+    """
+    global _SERVICE
+
+    if _SERVICE is not None:
+        LOGGER.debug("Experiment service already setup, skipping")
+        return
+
+    if experiment_name is None:
+        timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+        experiment_name = f"harbor-{timestamp}"
+
+    _SERVICE = HarborExperimentService(
+        experiment_name=experiment_name,
+        experiment_config=experiment_config,
+    )
+
+    LOGGER.info("Experiment service setup for '%s'", experiment_name)
+
+
+def get_service() -> Optional[HarborExperimentService]:
+    """Get the current experiment service instance, or None if not initialized."""
+    return _SERVICE
+
+
+def reset() -> None:
+    """Reset the experiment service. Used for testing."""
+    global _SERVICE
+    _SERVICE = None