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

opik/llm_usage/opik_usage.py

@@ -1,24 +1,27 @@
+import logging
 import pydantic
 from typing import Union, Dict, Any, Optional
 from . import (
-    openai_agent_usage,
-    openai_usage,
+    openai_chat_completions_usage,
     google_usage,
     anthropic_usage,
     unknown_usage,
     bedrock_usage,
+    openai_responses_usage,
 )
-from opik import dict_utils
+import opik.dict_utils as dict_utils
 
 ProviderUsage = Union[
-    openai_usage.OpenAICompletionsUsage,
+    openai_chat_completions_usage.OpenAICompletionsUsage,
     google_usage.GoogleGeminiUsage,
     anthropic_usage.AnthropicUsage,
     bedrock_usage.BedrockUsage,
-    openai_agent_usage.OpenAIAgentUsage,
+    openai_responses_usage.OpenAIResponsesUsage,
     unknown_usage.UnknownUsage,
 ]
 
+LOGGER = logging.getLogger(__name__)
+
 
 class OpikUsage(pydantic.BaseModel):
     """
@@ -68,7 +71,7 @@ class OpikUsage(pydantic.BaseModel):
 
     @classmethod
     def from_openai_completions_dict(cls, usage: Dict[str, Any]) -> "OpikUsage":
-        provider_usage = openai_usage.OpenAICompletionsUsage.from_original_usage_dict(
+        provider_usage = openai_chat_completions_usage.OpenAICompletionsUsage.from_original_usage_dict(
             usage
         )
 
@@ -79,12 +82,35 @@ class OpikUsage(pydantic.BaseModel):
             provider_usage=provider_usage,
         )
 
+    @classmethod
+    def from_groq_completions_dict(cls, usage: Dict[str, Any]) -> "OpikUsage":
+        return cls.from_openai_completions_dict(usage)
+
     @classmethod
     def from_google_dict(cls, usage: Dict[str, Any]) -> "OpikUsage":
         provider_usage = google_usage.GoogleGeminiUsage.from_original_usage_dict(usage)
 
+        # The completions token should include both the candidates token_count and the thought tokens. Usage differs depending on the models and between VertexAI and Google AI
+        # See https://github.com/BerriAI/litellm/pull/10141#discussion_r2052272035
+        # Do something similar as: https://github.com/BerriAI/litellm/blob/4854482af4a2a56060bbfeb4345bce4f1bb7ec41/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py#L980-L995
+        candidates_token_count = provider_usage.candidates_token_count or 0
+
+        total_token_count = provider_usage.prompt_token_count + candidates_token_count
+
+        if provider_usage.total_token_count == total_token_count:
+            completion_tokens = candidates_token_count
+        elif provider_usage.thoughts_token_count is not None:
+            completion_tokens = (
+                candidates_token_count + provider_usage.thoughts_token_count
+            )
+        else:
+            LOGGER.debug(
+                "Something is wrong in Google usage, completion_tokens might be invalid"
+            )
+            completion_tokens = candidates_token_count
+
         return cls(
-            completion_tokens=provider_usage.candidates_token_count,
+            completion_tokens=completion_tokens,
             prompt_tokens=provider_usage.prompt_token_count,
             total_tokens=provider_usage.total_token_count,
             provider_usage=provider_usage,
@@ -130,9 +156,9 @@ class OpikUsage(pydantic.BaseModel):
         )
 
     @classmethod
-    def from_openai_agent_dict(cls, usage: Dict[str, Any]) -> "OpikUsage":
-        provider_usage = openai_agent_usage.OpenAIAgentUsage.from_original_usage_dict(
-            usage
+    def from_openai_responses_dict(cls, usage: Dict[str, Any]) -> "OpikUsage":
+        provider_usage = (
+            openai_responses_usage.OpenAIResponsesUsage.from_original_usage_dict(usage)
         )
 
         return cls(

opik/llm_usage/opik_usage_factory.py

@@ -1,17 +1,23 @@
 import logging
-from . import opik_usage
-from typing import Dict, Any, Callable, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
+
 from opik.types import LLMProvider
+from . import opik_usage
+
 
-_PROVIDER_TO_OPIK_USAGE_BUILDER: Dict[
-    Union[str, LLMProvider], Callable[[Dict[str, Any]], opik_usage.OpikUsage]
+# One provider can have multiple formats of usage dicts, so it can have more than 1 build function
+_PROVIDER_TO_OPIK_USAGE_BUILDERS: Dict[
+    Union[str, LLMProvider],
+    List[Callable[[Dict[str, Any]], opik_usage.OpikUsage]],
 ] = {
-    LLMProvider.OPENAI: opik_usage.OpikUsage.from_openai_completions_dict,
-    LLMProvider.GOOGLE_VERTEXAI: opik_usage.OpikUsage.from_google_dict,
-    LLMProvider.GOOGLE_AI: opik_usage.OpikUsage.from_google_dict,
-    LLMProvider.ANTHROPIC: opik_usage.OpikUsage.from_anthropic_dict,
-    "_bedrock": opik_usage.OpikUsage.from_bedrock_dict,
-    "_openai_agent": opik_usage.OpikUsage.from_openai_agent_dict,
+    LLMProvider.OPENAI: [
+        opik_usage.OpikUsage.from_openai_completions_dict,
+        opik_usage.OpikUsage.from_openai_responses_dict,
+    ],
+    LLMProvider.GOOGLE_VERTEXAI: [opik_usage.OpikUsage.from_google_dict],
+    LLMProvider.GOOGLE_AI: [opik_usage.OpikUsage.from_google_dict],
+    LLMProvider.ANTHROPIC: [opik_usage.OpikUsage.from_anthropic_dict],
+    LLMProvider.BEDROCK: [opik_usage.OpikUsage.from_bedrock_dict],
 }
 
 
@@ -19,22 +25,32 @@ def build_opik_usage(
     provider: Union[str, LLMProvider],
     usage: Dict[str, Any],
 ) -> opik_usage.OpikUsage:
-    build_function = _PROVIDER_TO_OPIK_USAGE_BUILDER[provider]
+    build_functions = _PROVIDER_TO_OPIK_USAGE_BUILDERS[provider]
 
-    result = build_function(usage)
+    exc = None
+    for build_function in build_functions:
+        try:
+            result = build_function(usage)
+            return result
+        except Exception as exc_info:
+            exc = exc_info
+            pass
 
-    return result
+    raise ValueError(
+        f"Failed to build OpikUsage for provider {provider} and usage {usage}, reason: {exc}"
+    )
 
 
 def build_opik_usage_from_unknown_provider(
     usage: Dict[str, Any],
 ) -> opik_usage.OpikUsage:
-    for build_function in _PROVIDER_TO_OPIK_USAGE_BUILDER.values():
-        try:
-            opik_usage_ = build_function(usage)
-            return opik_usage_
-        except Exception:
-            pass
+    for build_functions in _PROVIDER_TO_OPIK_USAGE_BUILDERS.values():
+        for build_function in build_functions:
+            try:
+                opik_usage_ = build_function(usage)
+                return opik_usage_
+            except Exception:
+                pass
 
     return opik_usage.OpikUsage.from_unknown_usage_dict(usage)
 

opik/logging_messages.py

@@ -1,3 +1,6 @@
+# Storing all the messages in this module is considered a deprecated practice.
+# Do it only if your message is used in more than one place.
+
 METADATA_KEY_COLLISION_DURING_DEEPMERGE = (
     "Trace or span metadata value for the sub-key '%s' was overwritten from '%s' to '%s' during the deep merge",
 )
@@ -10,13 +13,11 @@ EXCEPTION_RAISED_FROM_TRACKED_FUNCTION = (
     "Exception raised from tracked function %s.\nInputs: %s\nError message: %s"
 )
 
-FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_OPENAI_LLM_RUN = (
-    "Failed to extract token usage from presumably OpenAI LLM langchain run."
-)
+FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_OPENAI_LLM_RUN = "Failed to extract token usage from presumably OpenAI LLM langchain run. The run dictionary: %s"
 
-FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GOOGLE_LLM_RUN = (
-    "Failed to extract token usage from presumably Google LLM langchain run."
-)
+FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GROQ_LLM_RUN = "Failed to extract token usage from presumably Groq LLM langchain run. The run dictionary: %s"
+
+FAILED_TO_EXTRACT_TOKEN_USAGE_FROM_PRESUMABLY_LANGCHAIN_GOOGLE_LLM_RUN = "Failed to extract token usage from presumably Google LLM langchain run. The run dictionary: %s"
 
 UNEXPECTED_EXCEPTION_ON_SPAN_CREATION_FOR_TRACKED_FUNCTION = "Unexpected exception happened when tried to create a span for function %s.\nInputs: %s\nError message: %s"
 
@@ -28,6 +29,8 @@ FAILED_TO_PARSE_OPENAI_STREAM_CONTENT = "Failed to parse openai Stream content.
 
 FAILED_TO_PROCESS_MESSAGE_IN_BACKGROUND_STREAMER = "Failed to process %s. Error: %s"
 
+MAKE_SURE_OPIK_IS_CONFIGURED_CORRECTLY = "This error may be due to incorrect configuration. For configuration help, see: https://www.comet.com/docs/opik/tracing/sdk_configuration"
+
 HALLUCINATION_DETECTION_FAILED = "Failed hallucination detection"
 
 FACTUALITY_SCORE_CALC_FAILED = "Failed to calculate factuality score"
@@ -36,12 +39,18 @@ ANSWER_RELEVANCE_SCORE_CALC_FAILED = "Failed to calculate answer relevance score
 
 MODERATION_SCORE_CALC_FAILED = "Failed to calculate moderation score"
 
+STRUCTURED_OUTPUT_COMPLIANCE_FAILED = (
+    "Failed to parse structured output compliance metric output."
+)
+
 CONTEXT_RECALL_SCORE_CALC_FAILED = "Failed to calculate context recall score"
 
 GEVAL_SCORE_CALC_FAILED = "Failed to calculate g-eval score"
 
 CONTEXT_PRECISION_SCORE_CALC_FAILED = "Failed to calculate context precision score"
 
+USEFULNESS_SCORE_CALC_FAILED = "Failed to calculate usefulness score"
+
 NESTED_SPAN_PROJECT_NAME_MISMATCH_WARNING_MESSAGE = (
     'You are attempting to log data into a nested span under the project name "{}". '
     'However, the project name "{}" from parent span will be used instead.'
@@ -57,4 +66,7 @@ PARSE_API_KEY_TOO_MANY_PARTS = "Too many parts (%d) found in the Opik API key: %
 
 LLM_PROVIDER_RATE_LIMIT_ERROR_DETECTED_IN_EVALUATE_FUNCTION = "LLM provider rate limit error detected. We recommend reducing the amount of parallel requests by setting `task_threads` evaluation parameter to a smaller number"
 
-WARNING_TOKEN_USAGE_DATA_IS_NOT_AVAILABLE = "You didn't specify argument `stream_usage`=True during LLM initialization. Token usage data is not available for .stream() or .astream() methods."
+WARNING_TOKEN_USAGE_DATA_IS_NOT_AVAILABLE = "You didn't specify argument `stream_usage=True` during LLM initialization. Token usage data is not available for .stream() or .astream() methods."
+
+# Storing all the messages in this module is considered a deprecated practice.
+# Do it only if your message is used in more than one place.

opik/message_processing/arguments_utils.py

@@ -0,0 +1,22 @@
+from typing import Any
+
+
+def recursive_shallow_copy(obj: Any) -> Any:
+    """
+    Creates a recursive shallow copy of a dictionary or list which may include other dictionaries
+    or lists as values.
+
+    Args:
+        obj: The object to copy. Can be a dict, list, or any other type.
+
+    Returns:
+        A recursive shallow copy of the input object where dictionaries and lists are copied,
+        but their non-container elements are referenced.
+    """
+    if isinstance(obj, dict):
+        return {k: recursive_shallow_copy(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [recursive_shallow_copy(item) for item in obj]
+    else:
+        # For non-container types, return the object itself (shallow copy)
+        return obj

opik/message_processing/batching/base_batcher.py

@@ -1,4 +1,3 @@
-import threading
 import time
 import abc
 
@@ -6,45 +5,74 @@ from typing import List, Callable
 from .. import messages
 
 
+BASE_BATCH_MEMORY_LIMIT_MB = 50
+
+
 class BaseBatcher(abc.ABC):
     def __init__(
         self,
         flush_callback: Callable[[messages.BaseMessage], None],
         max_batch_size: int,
         flush_interval_seconds: float,
+        batch_memory_limit_mb: int = BASE_BATCH_MEMORY_LIMIT_MB,
     ):
         self._flush_interval_seconds: float = flush_interval_seconds
         self._flush_callback: Callable[[messages.BaseMessage], None] = flush_callback
         self._accumulated_messages: List[messages.BaseMessage] = []
         self._max_batch_size: int = max_batch_size
+        self._batch_memory_limit_mb: int = batch_memory_limit_mb
 
         self._last_time_flush_callback_called: float = time.time()
-        self._lock = threading.RLock()
 
     def flush(self) -> None:
-        with self._lock:
-            if len(self._accumulated_messages) > 0:
-                batch_message = self._create_batch_from_accumulated_messages()
-                self._accumulated_messages = []
+        if len(self._accumulated_messages) > 0:
+            batch_messages = self._create_batches_from_accumulated_messages()
+            self._accumulated_messages = []
 
+            for batch_message in batch_messages:
                 self._flush_callback(batch_message)
-            self._last_time_flush_callback_called = time.time()
+        self._last_time_flush_callback_called = time.time()
 
     def is_ready_to_flush(self) -> bool:
-        return (
-            time.time() - self._last_time_flush_callback_called
-        ) >= self._flush_interval_seconds
+        elapsed = time.time() - self._last_time_flush_callback_called
+        return elapsed >= self._flush_interval_seconds
 
     def is_empty(self) -> bool:
-        with self._lock:
-            return len(self._accumulated_messages) == 0
+        return len(self._accumulated_messages) == 0
 
     @abc.abstractmethod
-    def _create_batch_from_accumulated_messages(self) -> messages.BaseMessage: ...
+    def _create_batches_from_accumulated_messages(
+        self,
+    ) -> List[messages.BaseMessage]: ...
 
     @abc.abstractmethod
     def add(self, message: messages.BaseMessage) -> None:
-        with self._lock:
-            self._accumulated_messages.append(message)
-            if len(self._accumulated_messages) >= self._max_batch_size:
-                self.flush()
+        self._accumulated_messages.append(message)
+        if len(self._accumulated_messages) >= self._max_batch_size:
+            self.flush()
+
+    def _remove_matching_messages(
+        self, filter_func: Callable[[messages.BaseMessage], bool]
+    ) -> None:
+        """
+        Remove messages from _accumulated_messages that match the provided filter function.
+
+        Args:
+            filter_func: A function that takes a BaseMessage and returns True if the message should be removed
+        """
+        self._accumulated_messages = list(
+            filter(lambda x: not filter_func(x), self._accumulated_messages)
+        )
+
+    def size(self) -> int:
+        """
+        Gets the total number of accumulated messages in the current instance.
+
+        The method calculates and retrieves the count of accumulated messages
+        maintained within an internal structure. Thread safety is guaranteed
+        through the usage of a lock mechanism.
+
+        Returns:
+            int: The total number of accumulated messages.
+        """
+        return len(self._accumulated_messages)

opik/message_processing/batching/batch_manager.py

@@ -1,3 +1,4 @@
+import threading
 from typing import Type, Dict
 from .. import messages
 from . import base_batcher
@@ -15,30 +16,41 @@ class BatchManager:
         self._flushing_thread = flushing_thread.FlushingThread(
             batchers=list(self._message_to_batcher_mapping.values())
         )
+        self._lock = threading.RLock()
 
     def start(self) -> None:
         self._flushing_thread.start()
 
     def stop(self) -> None:
-        self._flushing_thread.close()
+        with self._lock:
+            # stop the flushing thread
+            self._flushing_thread.close()
+            # force flush all pending messages
+            self.flush()
 
     def message_supports_batching(self, message: messages.BaseMessage) -> bool:
+        if message is None:
+            return False
+
         if hasattr(message, "supports_batching"):
             return message.supports_batching
 
         return message.__class__ in self._message_to_batcher_mapping
 
     def process_message(self, message: messages.BaseMessage) -> None:
-        self._message_to_batcher_mapping[type(message)].add(message)
+        with self._lock:
+            self._message_to_batcher_mapping[type(message)].add(message)
 
     def is_empty(self) -> bool:
-        return all(
-            [
-                batcher.is_empty()
-                for batcher in self._message_to_batcher_mapping.values()
-            ]
-        )
+        with self._lock:
+            return all(
+                [
+                    batcher.is_empty()
+                    for batcher in self._message_to_batcher_mapping.values()
+                ]
+            )
 
     def flush(self) -> None:
-        for batcher in self._message_to_batcher_mapping.values():
-            batcher.flush()
+        with self._lock:
+            for batcher in self._message_to_batcher_mapping.values():
+                batcher.flush()

opik/message_processing/batching/batch_manager_constuctors.py

@@ -1,11 +1,7 @@
-import queue
 from typing import Type, Dict
 
-from .. import messages
-
-from . import base_batcher
-from . import batchers
-from . import batch_manager
+from . import base_batcher, batchers, batch_manager
+from .. import messages, message_queue
 
 CREATE_SPANS_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS = 2.0
 CREATE_SPANS_MESSAGE_BATCHER_MAX_BATCH_SIZE = 1000
@@ -16,30 +12,56 @@ CREATE_TRACES_MESSAGE_BATCHER_MAX_BATCH_SIZE = 1000
 FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS = 1.0
 FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE = 1000
 
+GUARDRAIL_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS = 1.0
+GUARDRAIL_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE = 1000
+
+EXPERIMENT_ITEMS_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS = 3.0
+EXPERIMENT_ITEMS_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE = 1000
+
 
-def create_batch_manager(message_queue: queue.Queue) -> batch_manager.BatchManager:
+def create_batch_manager(
+    queue: message_queue.MessageQueue[messages.BaseMessage],
+) -> batch_manager.BatchManager:
     create_span_message_batcher_ = batchers.CreateSpanMessageBatcher(
         flush_interval_seconds=CREATE_SPANS_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
         max_batch_size=CREATE_SPANS_MESSAGE_BATCHER_MAX_BATCH_SIZE,
-        flush_callback=message_queue.put,
+        flush_callback=queue.put,
     )
 
     create_trace_message_batcher_ = batchers.CreateTraceMessageBatcher(
         flush_interval_seconds=CREATE_TRACES_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
         max_batch_size=CREATE_TRACES_MESSAGE_BATCHER_MAX_BATCH_SIZE,
-        flush_callback=message_queue.put,
+        flush_callback=queue.put,
     )
 
     add_span_feedback_scores_batch_message_batcher = batchers.AddSpanFeedbackScoresBatchMessageBatcher(
         flush_interval_seconds=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
         max_batch_size=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE,
-        flush_callback=message_queue.put,
+        flush_callback=queue.put,
     )
 
     add_trace_feedback_scores_batch_message_batcher = batchers.AddTraceFeedbackScoresBatchMessageBatcher(
         flush_interval_seconds=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
         max_batch_size=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE,
-        flush_callback=message_queue.put,
+        flush_callback=queue.put,
+    )
+
+    add_threads_feedback_scores_batch_message_batcher = batchers.AddThreadsFeedbackScoresBatchMessageBatcher(
+        flush_interval_seconds=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
+        max_batch_size=FEEDBACK_SCORES_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE,
+        flush_callback=queue.put,
+    )
+
+    guardrail_batch_message_batcher = batchers.GuardrailBatchMessageBatcher(
+        flush_interval_seconds=GUARDRAIL_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
+        max_batch_size=GUARDRAIL_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE,
+        flush_callback=queue.put,
+    )
+
+    experiment_items_batch_message_batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        flush_interval_seconds=EXPERIMENT_ITEMS_BATCH_MESSAGE_BATCHER_FLUSH_INTERVAL_SECONDS,
+        max_batch_size=EXPERIMENT_ITEMS_BATCH_MESSAGE_BATCHER_MAX_BATCH_SIZE,
+        flush_callback=queue.put,
     )
 
     message_to_batcher_mapping: Dict[
@@ -49,6 +71,9 @@ def create_batch_manager(message_queue: queue.Queue) -> batch_manager.BatchManag
         messages.CreateTraceMessage: create_trace_message_batcher_,
         messages.AddSpanFeedbackScoresBatchMessage: add_span_feedback_scores_batch_message_batcher,
         messages.AddTraceFeedbackScoresBatchMessage: add_trace_feedback_scores_batch_message_batcher,
+        messages.AddThreadsFeedbackScoresBatchMessage: add_threads_feedback_scores_batch_message_batcher,
+        messages.GuardrailBatchMessage: guardrail_batch_message_batcher,
+        messages.CreateExperimentItemsBatchMessage: experiment_items_batch_message_batcher,
     }
 
     batch_manager_ = batch_manager.BatchManager(