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

opik/api_objects/experiment/rest_operations.py

@@ -1,13 +1,14 @@
+import json
 from typing import List, Optional
 
-import opik.exceptions as exceptions
+from . import experiment_item
 from .. import rest_stream_parser
-from opik.rest_api import OpikApi
-from opik.rest_api.types import experiment_public
+from ... import exceptions, rest_api
+from ...rest_api.types import experiment_public
 
 
 def get_experiment_data_by_name(
-    rest_client: OpikApi, name: str
+    rest_client: rest_api.OpikApi, name: str
 ) -> experiment_public.ExperimentPublic:
     # TODO: this method is deprecated and should be removed after
     #  deprecated Opik.get_experiment_by_name() will be removed.
@@ -22,7 +23,7 @@ def get_experiment_data_by_name(
 
 
 def get_experiments_data_by_name(
-    rest_client: OpikApi,
+    rest_client: rest_api.OpikApi,
     name: str,
     max_results: Optional[int] = None,
 ) -> List[experiment_public.ExperimentPublic]:
@@ -43,3 +44,62 @@ def get_experiments_data_by_name(
         )
 
     return experiments
+
+
+def find_experiment_items_for_dataset(
+    rest_client: rest_api.OpikApi,
+    dataset_id: str,
+    experiment_ids: List[str],
+    max_results: int,
+    truncate: bool,
+    filter_expression: Optional[str] = None,
+) -> List[experiment_item.ExperimentItemContent]:
+    PAGE_SIZE = 100
+
+    collected_items: List[experiment_item.ExperimentItemContent] = []
+    experiment_ids_json = json.dumps(experiment_ids)
+
+    page_number = 1
+    while len(collected_items) < max_results:
+        page_dataset_items = (
+            rest_client.datasets.find_dataset_items_with_experiment_items(
+                id=dataset_id,
+                page=page_number,
+                size=PAGE_SIZE,
+                experiment_ids=experiment_ids_json,
+                truncate=truncate,
+                filters=filter_expression,
+            )
+        )
+
+        if not page_dataset_items.content:
+            break
+
+        # Build a flat list of experiment-item compares for this page in a readable way
+        experiment_items = []
+        for dataset_item in page_dataset_items.content:
+            # Guard if dataset_item.experiment_items might be None
+            if dataset_item.experiment_items is not None:
+                for experiment_item_compare in dataset_item.experiment_items:
+                    # Convert to domain objects
+                    dataset_item_data = dataset_item.data
+                    if dataset_item_data is not None:
+                        dataset_item_data.update({"id": dataset_item.id})
+                    experiment_item_content = experiment_item.ExperimentItemContent.from_rest_experiment_item_compare(
+                        value=experiment_item_compare,
+                        dataset_item_data=dataset_item_data,
+                    )
+
+                    experiment_items.append(experiment_item_content)
+
+        if not experiment_items:
+            page_number += 1
+            continue
+
+        remaining = max_results - len(collected_items)
+
+        collected_items.extend(experiment_items[:remaining])
+
+        page_number += 1
+
+    return collected_items

opik/api_objects/helpers.py

@@ -13,11 +13,10 @@ from ..rest_api.types import (
     trace_filter_public,
     trace_thread_filter,
 )
-from ..types import FeedbackScoreDict
+from ..types import BatchFeedbackScoreDict
 
 LOGGER = logging.getLogger(__name__)
 
-
 FilterParsedItemT = TypeVar(
     "FilterParsedItemT",
     bound=Union[
@@ -140,7 +139,7 @@ def parse_search_expressions(
 
 
 def parse_feedback_score_messages(
-    scores: List[FeedbackScoreDict],
+    scores: List[BatchFeedbackScoreDict],
     project_name: str,
     parsed_item_class: Type[ScoreMessageT],
     logger: logging.Logger,
@@ -156,9 +155,13 @@ def parse_feedback_score_messages(
 
     score_messages = [
         parsed_item_class(
+            id=score_dict["id"],
+            name=score_dict["name"],
+            value=score_dict["value"],
             source=constants.FEEDBACK_SCORE_SOURCE_SDK,
-            project_name=project_name,
-            **score_dict,
+            project_name=score_dict.get("project_name") or project_name,
+            reason=score_dict.get("reason"),
+            category_name=score_dict.get("category_name"),
         )
         for score_dict in valid_scores
     ]

opik/api_objects/local_recording.py

@@ -0,0 +1,81 @@
+import contextlib
+from typing import Iterator, List
+from typing import Optional
+
+from . import opik_client
+from ..message_processing.emulation import local_emulator_message_processor, models
+from ..message_processing.processors import message_processors_chain
+
+
+class _LocalRecordingHandle:
+    """Lightweight handle that exposes locally recorded data.
+
+    Accessors flush the client to ensure all in-flight messages are processed
+    before reading the in-memory models.
+    """
+
+    def __init__(
+        self,
+        client: opik_client.Opik,
+        local_processor: local_emulator_message_processor.LocalEmulatorMessageProcessor,
+    ) -> None:
+        self._client = client
+        self._local = local_processor
+
+    @property
+    def span_trees(self) -> List[models.SpanModel]:
+        self._client.flush()
+        return self._local.span_trees
+
+    @property
+    def trace_trees(self) -> List[models.TraceModel]:
+        self._client.flush()
+        return self._local.trace_trees
+
+
+@contextlib.contextmanager
+def record_traces_locally(
+    client: Optional[opik_client.Opik] = None,
+) -> Iterator[_LocalRecordingHandle]:
+    """Enable local recording of traces/spans within the context.
+
+    Args:
+        client: Optional Opik client to use for recording. If not provided, the default session client will be used.
+    Usage:
+        with opik.record_traces_locally() as storage:
+            # your code that creates traces/spans
+            spans = storage.span_trees
+
+    Yields a handle with `span_trees` and `trace_trees` properties that flush
+    the client before reading, ensuring all events are captured.
+    """
+    if client is None:
+        client = opik_client.get_client_cached()
+
+    # Disallow nested/local concurrent recordings in the same process
+    existing_local = message_processors_chain.get_local_emulator_message_processor(
+        chain=client.__internal_api__message_processor__
+    )
+    if existing_local is not None and existing_local.is_active():
+        raise RuntimeError(
+            "record_traces_locally() is already active in the current context; nested usage is not allowed."
+        )
+
+    message_processors_chain.toggle_local_emulator_message_processor(
+        active=True, chain=client.__internal_api__message_processor__, reset=True
+    )
+    local = message_processors_chain.get_local_emulator_message_processor(
+        chain=client.__internal_api__message_processor__
+    )
+    if local is None:
+        # Should not happen given the default chain, but guard just in case
+        raise RuntimeError("Local emulator message processor not available")
+
+    handle = _LocalRecordingHandle(client=client, local_processor=local)
+    try:
+        yield handle
+    finally:
+        client.flush()
+        message_processors_chain.toggle_local_emulator_message_processor(
+            active=False, chain=client.__internal_api__message_processor__, reset=True
+        )