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

opik/message_processing/emulation/local_emulator_message_processor.py

@@ -0,0 +1,140 @@
+import datetime
+import logging
+from typing import List, Dict, Optional, Any, Union
+
+from opik.types import ErrorInfoDict, SpanType
+from . import models, emulator_message_processor
+from .. import messages
+from ...rest_api.types import span_write, trace_write
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class LocalEmulatorMessageProcessor(
+    emulator_message_processor.EmulatorMessageProcessor
+):
+    """This class serves as a replacement for the real backend and collects all logged messages
+    locally in memory to be used for evaluation.
+    """
+
+    def __init__(self, active: bool, merge_duplicates: bool = True) -> None:
+        super().__init__(active=active, merge_duplicates=merge_duplicates)
+
+    def process(
+        self,
+        message: Union[
+            messages.BaseMessage, span_write.SpanWrite, trace_write.TraceWrite
+        ],
+    ) -> None:
+        if not self.is_active():
+            return
+
+        if hasattr(message, "delivery_attempts") and message.delivery_attempts > 1:
+            # skip retries
+            LOGGER.debug("Skipping retry of the message: %s", message)
+            return
+
+        super().process(message)
+
+    def create_trace_model(
+        self,
+        trace_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        project_name: str,
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        error_info: Optional[ErrorInfoDict],
+        thread_id: Optional[str],
+        last_updated_at: Optional[datetime.datetime] = None,
+    ) -> models.TraceModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.TraceModel(
+            id=trace_id,
+            start_time=start_time,
+            name=name,
+            project_name=project_name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            end_time=end_time,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            error_info=error_info,
+            thread_id=thread_id,
+            last_updated_at=last_updated_at,
+        )
+
+    def create_span_model(
+        self,
+        span_id: str,
+        start_time: datetime.datetime,
+        name: Optional[str],
+        input: Any,
+        output: Any,
+        tags: Optional[List[str]],
+        metadata: Optional[Dict[str, Any]],
+        type: SpanType,
+        usage: Optional[Dict[str, Any]],
+        end_time: Optional[datetime.datetime],
+        project_name: str,
+        spans: Optional[List[models.SpanModel]],
+        feedback_scores: Optional[List[models.FeedbackScoreModel]],
+        model: Optional[str],
+        provider: Optional[str],
+        error_info: Optional[ErrorInfoDict],
+        total_cost: Optional[float],
+        last_updated_at: Optional[datetime.datetime],
+    ) -> models.SpanModel:
+        if spans is None:
+            spans = []
+        if feedback_scores is None:
+            feedback_scores = []
+
+        return models.SpanModel(
+            id=span_id,
+            start_time=start_time,
+            name=name,
+            input=input,
+            output=output,
+            tags=tags,
+            metadata=metadata,
+            type=type,
+            usage=usage,
+            end_time=end_time,
+            project_name=project_name,
+            spans=spans,
+            feedback_scores=feedback_scores,
+            model=model,
+            provider=provider,
+            error_info=error_info,
+            total_cost=total_cost,
+            last_updated_at=last_updated_at,
+        )
+
+    def create_feedback_score_model(
+        self,
+        score_id: str,
+        name: str,
+        value: float,
+        category_name: Optional[str],
+        reason: Optional[str],
+    ) -> models.FeedbackScoreModel:
+        return models.FeedbackScoreModel(
+            id=score_id,
+            name=name,
+            value=value,
+            category_name=category_name,
+            reason=reason,
+        )

opik/message_processing/emulation/models.py

@@ -0,0 +1,162 @@
+from typing import List, Any, Optional, Dict
+
+from opik.config import OPIK_PROJECT_DEFAULT_NAME
+from opik.types import ErrorInfoDict
+
+import dataclasses
+import datetime
+
+
+@dataclasses.dataclass
+class FeedbackScoreModel:
+    """
+    Represents a model for a feedback score used to evaluate specific spans or traces.
+
+    This class stores and manages feedback scores linked to defined criteria, including
+    identifiers, names, values, categories, and explanations for each score.
+
+    Attributes:
+        id: Unique identifier for the feedback score.
+        name: Name associated with the feedback score.
+        value: The numerical value of the feedback score.
+        category_name: Category to which the feedback score belongs, if any.
+        reason: Reason or explanation for the feedback score, if available.
+    """
+
+    id: str
+    name: str
+    value: float
+    category_name: Optional[str] = None
+    reason: Optional[str] = None
+
+
+@dataclasses.dataclass
+class SpanModel:
+    """
+    Represents a span model used to describe specific points in a process, their metadata,
+    and associated data.
+
+    This class is used to store and manipulate structured data for events or
+    spans, including metadata, time markers, associated input/output, tags,
+    and additional properties. It serves as a representative structure for recording
+    and organizing event-specific information, often used in applications like
+    logging, distributed tracing, or data processing pipelines.
+
+    Attributes:
+        id: Unique identifier for the span.
+        start_time: Start time of the span.
+        name: Name of the span, if provided.
+        input: Input data associated with the span, if any.
+        output: Output data associated with the span, if any.
+        tags: List of tags linked to the span.
+        metadata: Additional metadata for the span.
+        type: Type of the span, defaulting to "general".
+        usage: Usage-related information for the span.
+        end_time: End time of the span, if available.
+        project_name: Name of the project the span is associated with,
+            defaulting to a predefined project name.
+        spans: List of nested spans related to this span.
+        feedback_scores: List of feedback scores associated
+            with the span.
+        model: Model identification used, if applicable.
+        provider: Provider of the span or associated services, if any.
+        error_info: Error information or diagnostics
+            for the span, if applicable.
+        total_cost: Total cost incurred associated with this span,
+            if relevant.
+        last_updated_at: Timestamp of when the span was
+            last updated, if available.
+    """
+
+    id: str
+    start_time: datetime.datetime
+    name: Optional[str] = None
+    input: Optional[Dict[str, Any]] = None
+    output: Optional[Dict[str, Any]] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    type: str = "general"
+    usage: Optional[Dict[str, Any]] = None
+    end_time: Optional[datetime.datetime] = None
+    project_name: str = OPIK_PROJECT_DEFAULT_NAME
+    spans: List["SpanModel"] = dataclasses.field(default_factory=list)
+    feedback_scores: List[FeedbackScoreModel] = dataclasses.field(default_factory=list)
+    model: Optional[str] = None
+    provider: Optional[str] = None
+    error_info: Optional[ErrorInfoDict] = None
+    total_cost: Optional[float] = None
+    last_updated_at: Optional[datetime.datetime] = None
+
+
+@dataclasses.dataclass
+class ExperimentItemModel:
+    """
+    Represents an experiment item model that links a trace to a dataset item in an experiment.
+
+    Attributes:
+        id: Unique identifier for the experiment item.
+        experiment_id: The ID of the experiment this item belongs to.
+        trace_id: The ID of the trace associated with this experiment item.
+        dataset_item_id: The ID of the dataset item associated with this experiment item.
+    """
+
+    id: str
+    experiment_id: str
+    trace_id: str
+    dataset_item_id: str
+
+
+@dataclasses.dataclass
+class TraceModel:
+    """
+    Represents a trace model that encapsulates data about a trace, its related metadata,
+    and associated spans. It is used for tracking and analyzing data during execution
+    or processing tasks.
+
+    This class provides a structure to represent trace information, including the start
+    and end times, associated project details, input/output data, feedback scores, error
+    information, and thread association. It is designed to handle optional fields for
+    flexible use across various scenarios.
+
+    Attributes:
+        id: Unique identifier for the trace.
+        start_time: Timestamp representing the start of the trace.
+        name: Optional name for the trace, which can provide a descriptive
+            label.
+        project_name: Name of the project associated with the trace.
+        input: Optional dictionary containing the input data
+            associated with the trace.
+        output: Optional dictionary containing the output data
+            generated by the trace.
+        tags: Optional list of tags associated with the trace for
+            classification or filtering purposes.
+        metadata: Optional metadata providing additional
+            information about the trace.
+        end_time: Timestamp representing the end of the
+            trace.
+        spans: List of spans associated with the trace, representing
+            individual processing parts or segments within the trace.
+        feedback_scores: List of feedback scores associated
+            with the trace.
+        error_info: Optional dictionary containing information
+            about errors encountered during the trace.
+        thread_id: Optional identifier of the thread associated with the
+            trace.
+        last_updated_at: Timestamp for when the trace was
+            last updated.
+    """
+
+    id: str
+    start_time: datetime.datetime
+    name: Optional[str]
+    project_name: str
+    input: Optional[Dict[str, Any]] = None
+    output: Optional[Dict[str, Any]] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    end_time: Optional[datetime.datetime] = None
+    spans: List[SpanModel] = dataclasses.field(default_factory=list)
+    feedback_scores: List[FeedbackScoreModel] = dataclasses.field(default_factory=list)
+    error_info: Optional[ErrorInfoDict] = None
+    thread_id: Optional[str] = None
+    last_updated_at: Optional[datetime.datetime] = None

opik/message_processing/encoder_helpers.py

@@ -0,0 +1,79 @@
+from typing import Dict, Any, Set, List, Literal
+
+import opik.hooks
+
+from .. import jsonable_encoder
+from ..anonymizer import anonymizer
+
+
+def encode_and_anonymize(
+    kwargs_dict: Dict[str, Any],
+    fields_to_anonymize: Set[str],
+    object_type: Literal["span", "trace"],
+) -> Dict[str, Any]:
+    """
+    Encodes and anonymizes the data in the given dictionary based on the specified
+    fields using registered anonymizers. If no anonymizers are registered, the
+    function simply encodes the dictionary without anonymization.
+
+    Args:
+        kwargs_dict: The dictionary containing the data to encode
+            and anonymize.
+        fields_to_anonymize: The set of fields within the dictionary to
+            anonymize.
+        object_type: A string indicating the type of object ('span' or 'trace')
+            that was used to create the kwargs_dict. This is passed to anonymizers
+            to provide context about the source object.
+
+    Returns:
+        A dictionary that has been encoded and, if applicable, anonymized.
+    """
+    # check if any anonymizer was registered
+    encoded_obj = jsonable_encoder.encode(kwargs_dict)
+    if not opik.hooks.has_anonymizers():
+        return encoded_obj
+
+    anonymizers = opik.hooks.get_anonymizers()
+    return anonymize_encoded_obj(
+        obj=encoded_obj,
+        fields_to_anonymize=fields_to_anonymize,
+        anonymizers=anonymizers,
+        object_type=object_type,
+    )
+
+
+def anonymize_encoded_obj(
+    obj: Dict[str, Any],
+    fields_to_anonymize: Set[str],
+    anonymizers: List[anonymizer.Anonymizer],
+    object_type: Literal["span", "trace"],
+) -> Dict[str, Any]:
+    """
+    Anonymizes specified fields in an encoded dictionary using the provided anonymizers.
+    This function iterates over the given set of field names and applies each anonymizer
+    to the corresponding field in the dictionary, if present. The anonymizers are expected
+    to implement an `anonymize` method that takes the field value, field name, and object type
+    as arguments. Only fields present in the dictionary and listed in `fields_to_anonymize`
+    are anonymized.
+
+    Args:
+        obj: The encoded dictionary whose fields are to be anonymized.
+        fields_to_anonymize: A set of field names within the dictionary to anonymize.
+        anonymizers: A list of anonymizer instances to apply to each field.
+        object_type: A string indicating the type of object ('span' or 'trace'),
+            providing context for anonymization.
+
+    Returns:
+        The dictionary with specified fields anonymized using the provided anonymizers.
+    """
+    if isinstance(obj, dict):
+        for field_name in fields_to_anonymize:
+            if field_name in obj:
+                for anonymizer_instance in anonymizers:
+                    obj[field_name] = anonymizer_instance.anonymize(
+                        obj[field_name],
+                        field_name=field_name,
+                        object_type=object_type,
+                    )
+
+    return obj

opik/message_processing/message_queue.py

@@ -0,0 +1,79 @@
+import collections
+from queue import Empty
+import threading
+from typing import TypeVar, Optional, Generic
+
+T = TypeVar("T")
+
+
+class MessageQueue(Generic[T]):
+    """A thread-safe message queue implementation using a double-ended queue (deque).
+
+    This queue allows putting messages at either end and retrieving them in a FIFO manner.
+    It supports optional maximum length limiting and blocking retrieval with timeout.
+
+    Once a bounded length queue is full, when new items are added, a corresponding number
+    of items are discarded from the opposite end.
+
+    Args:
+        max_length (Optional[int]): Maximum number of messages the queue can hold.
+            If None, the queue size is unlimited.
+
+    The queue implements the following operations:
+    - put(): Add a message to the front of the queue
+    - put_back(): Add a message to the back of the queue
+    - get(): Remove and return message from the back of the queue
+    - empty(): Check if the queue is empty
+    """
+
+    def __init__(self, max_length: Optional[int] = None):
+        self._deque: collections.deque[T] = collections.deque(maxlen=max_length)
+        self._not_empty = threading.Condition()
+        self.max_size = max_length
+
+    def put(self, message: T) -> None:
+        with self._not_empty:
+            self._deque.appendleft(message)
+            self._not_empty.notify()
+
+    def put_back(self, message: T) -> None:
+        with self._not_empty:
+            self._deque.append(message)
+            self._not_empty.notify()
+
+    def get(self, timeout: float) -> T:
+        with self._not_empty:
+            if timeout is None or timeout < 0:
+                raise ValueError("'timeout' must be a non-negative number")
+
+            self._not_empty.wait_for(lambda: len(self._deque) > 0, timeout=timeout)
+
+            if len(self._deque) == 0:
+                raise Empty
+
+            return self._deque.pop()
+
+    def empty(self) -> bool:
+        return len(self._deque) == 0
+
+    def accept_put_without_discarding(self) -> bool:
+        if self.max_size is None:
+            return True
+        with self._not_empty:
+            return len(self._deque) < self.max_size
+
+    def size(self) -> int:
+        return len(self._deque)
+
+    def __len__(self) -> int:
+        return len(self._deque)
+
+
+def calculate_max_queue_size(
+    maximal_queue_size: int,
+    batch_factor: int,
+) -> int:
+    if batch_factor > 0:
+        return int(maximal_queue_size / batch_factor)
+    else:
+        return maximal_queue_size