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

opik/integrations/llama_index/callback.py

@@ -1,21 +1,26 @@
+import contextvars
 import logging
-from typing import Optional, Dict, List, Any
+from typing import Optional, Dict, Any, List, Union
 import uuid
 
 from llama_index.core.callbacks import schema as llama_index_schema
 from llama_index.core.callbacks import base_handler
 
-import opik.opik_context as opik_context
-import opik.decorator.tracing_runtime_config as tracing_runtime_config
+from opik import context_storage, tracing_runtime_config
+from opik.decorator import arguments_helpers, span_creation_handler
 from opik.api_objects import opik_client, span, trace
 
 from . import event_parsing_utils
-from ...api_objects import helpers
 
 LOGGER = logging.getLogger(__name__)
 
-
 INDEX_CONSTRUCTION_TRACE_NAME = "index_construction"
+LLAMA_INDEX_METADATA = {"created_from": "llama_index"}
+
+# Context variable for root trace/span created by LlamaIndex
+_llama_root: contextvars.ContextVar[Optional[Union[span.SpanData, trace.TraceData]]] = (
+    contextvars.ContextVar("_llama_root", default=None)
+)
 
 
 def _get_last_event(trace_map: Dict[str, List[str]]) -> str:
@@ -36,20 +41,13 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         project_name: Optional[str] = None,
         skip_index_construction_trace: bool = False,
     ):
-        """
-        Initialize the instance with optional customization to define event filters and project-
-        specific data handling. The constructor sets up the necessary client and data mappings
-        for operational processing.
-
-        Parameters:
-            event_starts_to_ignore: Optional list of event start types to be ignored during
-                processing.
-            event_ends_to_ignore: Optional list of event end types to be ignored during
-                processing.
-            project_name: Optional string representing the project name to establish context in
-                client operations.
-            skip_index_construction_trace: A boolean value determining whether to skip creation of trace/spans of index
-                construction.
+        """Initialize LlamaIndex callback handler for Opik tracing.
+
+        Args:
+            event_starts_to_ignore: Event start types to be ignored during processing.
+            event_ends_to_ignore: Event end types to be ignored during processing.
+            project_name: Project name for trace/span context.
+            skip_index_construction_trace: Whether to skip index construction traces.
         """
         event_starts_to_ignore = (
             event_starts_to_ignore if event_starts_to_ignore else []
@@ -62,30 +60,27 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
 
         self._skip_index_construction_trace = skip_index_construction_trace
         self._project_name = project_name
-        self._opik_client = opik_client.Opik(
-            _use_batching=True,
-            project_name=project_name,
-        )
-
-        self._opik_trace_data: Optional[trace.TraceData] = None
+        self._opik_client = opik_client.get_client_cached()
+        self._opik_context_storage = context_storage.get_current_context_instance()
 
+        # Event tracking - shared across contexts, but events have unique IDs
         self._map_event_id_to_span_data: Dict[str, span.SpanData] = {}
         self._map_event_id_to_output: Dict[str, Any] = {}
 
-    def _create_trace_data(self, trace_name: Optional[str]) -> trace.TraceData:
-        trace_data = trace.TraceData(
-            name=trace_name,
-            metadata={"created_from": "llama_index"},
-            project_name=self._project_name,
-        )
-
-        if (
-            self._opik_client.config.log_start_trace_span
-            and tracing_runtime_config.is_tracing_active()
-        ):
-            self._opik_client.trace(**trace_data.as_start_parameters)
+        # For streaming: end_trace may be called before event_end, so we need to
+        # defer the trace output update until the event output is available
+        self._pending_root_output_updates: Dict[
+            str, Union[span.SpanData, trace.TraceData]
+        ] = {}
 
-        return trace_data
+    def _send_root_to_backend(
+        self, root: Union[span.SpanData, trace.TraceData]
+    ) -> None:
+        """Send root trace or span data to the backend."""
+        if isinstance(root, span.SpanData):
+            self._opik_client.span(**root.as_parameters)
+        elif isinstance(root, trace.TraceData):
+            self._opik_client.trace(**root.as_parameters)
 
     def start_trace(self, trace_id: Optional[str] = None) -> None:
         if (
@@ -94,18 +89,29 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         ):
             return
 
-        # When a new LLama Index trace is started, create a new trace in Opik
-        existing_trace_data = opik_context.get_current_trace_data()
-        if existing_trace_data is not None:
-            self._opik_trace_data = existing_trace_data
-        else:
-            self._opik_trace_data = self._create_trace_data(trace_name=trace_id)
+        trace_name = trace_id if trace_id else "llama_index_operation"
 
-        if (
-            self._opik_client.config.log_start_trace_span
-            and tracing_runtime_config.is_tracing_active()
-        ):
-            self._opik_client.trace(**self._opik_trace_data.as_start_parameters)
+        span_creation_result = span_creation_handler.create_span_respecting_context(
+            start_span_arguments=arguments_helpers.StartSpanParameters(
+                name=trace_name,
+                type="general",
+                project_name=self._project_name,
+                metadata=LLAMA_INDEX_METADATA,
+            ),
+            distributed_trace_headers=None,
+            opik_context_storage=self._opik_context_storage,
+        )
+
+        if span_creation_result.trace_data is not None:
+            self._opik_context_storage.set_trace_data(span_creation_result.trace_data)
+            self._opik_client.trace(
+                **span_creation_result.trace_data.as_start_parameters
+            )
+            _llama_root.set(span_creation_result.trace_data)
+        else:
+            self._opik_context_storage.add_span_data(span_creation_result.span_data)
+            self._opik_client.span(**span_creation_result.span_data.as_start_parameters)
+            _llama_root.set(span_creation_result.span_data)
 
     def end_trace(
         self,
@@ -115,21 +121,38 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         if not trace_map:
             return
 
-        # When a trace finishes, we first get the last event output
+        root = _llama_root.get()
+        if root is None:
+            return
+
         last_event = _get_last_event(trace_map)
-        last_event_output = self._map_event_id_to_output.get(last_event, None)
 
-        # And then end the trace with the optional output
-        if self._opik_trace_data is not None:
-            self._opik_trace_data.init_end_time().update(output=last_event_output)
-            if tracing_runtime_config.is_tracing_active():
-                self._opik_client.trace(**self._opik_trace_data.as_parameters)
-            self._opik_trace_data = None
+        # Check if the output for the last event is already available.
+        # For streaming calls, LlamaIndex calls end_trace() BEFORE event_end(),
+        # so the output won't be stored yet.
+        if last_event in self._map_event_id_to_output:
+            last_event_output = self._map_event_id_to_output.get(last_event)
+            root.init_end_time().update(output=last_event_output)
 
-        # Do not clean _map_event_id_to_span_data as streaming LLM events can
-        # end after this method is called. _map_event_id_to_span_data is
-        # individually cleaned after each event is ended
-        self._map_event_id_to_output.clear()
+            # Send the trace/span with output
+            self._send_root_to_backend(root)
+        else:
+            # Output not available yet (streaming scenario).
+            # Store the root so we can update it when event_end is called.
+            # Don't send the trace/span yet - it will be sent in on_event_end
+            # with the output and correct end_time to avoid race conditions.
+            # Note: We don't set end_time here because the actual end is when
+            # the last event ends, not when LlamaIndex calls end_trace().
+            self._pending_root_output_updates[last_event] = root
+
+        # Clean up context storage
+        if isinstance(root, span.SpanData):
+            self._opik_context_storage.pop_span_data(ensure_id=root.id)
+        elif isinstance(root, trace.TraceData):
+            self._opik_context_storage.pop_trace_data(ensure_id=root.id)
+
+        # Clean up
+        _llama_root.set(None)
 
     def on_event_start(
         self,
@@ -142,53 +165,60 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         if not event_id:
             event_id = str(uuid.uuid4())
 
-        # the event is not part of a trace probably because we are skipping the index construction trace
-        if self._opik_trace_data is None:
+        root_span_or_trace = _llama_root.get()
+
+        if root_span_or_trace is None:
             if not self._skip_index_construction_trace:
                 LOGGER.warning(
-                    "No trace data found in context for event start. "
-                    "This is likely due to the fact that the trace is not started properly. "
-                    f"The parent_id: {parent_id}, event_type: {event_type}, event_id: {event_id}."
+                    "No active LlamaIndex trace/span found in context. "
+                    "parent_id=%s, event_type=%s, event_id=%s",
+                    parent_id,
+                    event_type,
+                    event_id,
                 )
-
             return event_id
 
-        # Get parent span Id if it exists
-        if parent_id and parent_id in self._map_event_id_to_span_data:
-            opik_parent_id = self._map_event_id_to_span_data[parent_id].id
-        else:
-            opik_parent_id = None
-
-        # Compute the span input based on the event payload
         span_input = event_parsing_utils.get_span_input_from_events(event_type, payload)
 
-        project_name = helpers.resolve_child_span_project_name(
-            parent_project_name=self._opik_trace_data.project_name,
-            child_project_name=self._project_name,
-            show_warning=self._opik_trace_data.created_by != "evaluation",
+        # Skip creating span if event duplicates root operation name
+        root_name = root_span_or_trace.name if root_span_or_trace else None
+        event_duplicates_root = (
+            parent_id == llama_index_schema.BASE_TRACE_EVENT
+            and event_type.value == root_name
         )
+        if event_duplicates_root:
+            if span_input:
+                root_span_or_trace.update(input=span_input)
+            return event_id
 
-        # Create a new span for this event
-        span_data = span.SpanData(
-            trace_id=self._opik_trace_data.id,
-            name=event_type.value,
-            parent_span_id=opik_parent_id,
-            type=(
-                "llm" if event_type == llama_index_schema.CBEventType.LLM else "general"
+        span_creation_result = span_creation_handler.create_span_respecting_context(
+            start_span_arguments=arguments_helpers.StartSpanParameters(
+                name=event_type.value,
+                input=span_input,
+                type=(
+                    "llm"
+                    if event_type == llama_index_schema.CBEventType.LLM
+                    else "general"
+                ),
+                project_name=self._project_name,
+                metadata=LLAMA_INDEX_METADATA,
             ),
-            input=span_input,
-            project_name=project_name,
+            distributed_trace_headers=None,
+            opik_context_storage=self._opik_context_storage,
         )
+        span_data = span_creation_result.span_data
         self._map_event_id_to_span_data[event_id] = span_data
+        self._opik_context_storage.add_span_data(span_data)
+
         if (
             self._opik_client.config.log_start_trace_span
             and tracing_runtime_config.is_tracing_active()
         ):
             self._opik_client.span(**span_data.as_start_parameters)
 
-        # If the parent_id is a BASE_TRACE_EVENT, update the trace with the span input
-        if parent_id == llama_index_schema.BASE_TRACE_EVENT and span_input:
-            self._opik_trace_data.update(input=span_input)
+        # Update root input from first child event
+        if parent_id == llama_index_schema.BASE_TRACE_EVENT and span_input is not None:
+            root_span_or_trace.update(input=span_input)
 
         return event_id
 
@@ -199,32 +229,41 @@ class LlamaIndexCallbackHandler(base_handler.BaseCallbackHandler):
         event_id: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
-        # Get the span output from the event and store it so we can use it if needed
-        # when finishing the trace
         span_output = event_parsing_utils.get_span_output_from_event(
             event_type, payload
         )
-
         error_info = event_parsing_utils.get_span_error_info(payload)
 
-        if event_id:
-            self._map_event_id_to_output[event_id] = span_output
+        if not event_id:
+            return
 
-            # Log the output to the span with the matching id
-            if event_id in self._map_event_id_to_span_data:
-                span_data = self._map_event_id_to_span_data[event_id]
+        # Store output for end_trace
+        self._map_event_id_to_output[event_id] = span_output
 
-                llm_usage_info = event_parsing_utils.get_usage_data(payload)
-                span_data.update(**llm_usage_info.__dict__)
+        # Check if there's a pending root trace/span output update for this event.
+        # This happens when end_trace() was called before event_end() (streaming scenario).
+        if event_id in self._pending_root_output_updates:
+            root = self._pending_root_output_updates.pop(event_id)
+            # Set end_time now (the actual end) and update with output
+            root.init_end_time().update(output=span_output)
 
-                span_data.update(
-                    output=span_output, error_info=error_info
-                ).init_end_time()
-                if tracing_runtime_config.is_tracing_active():
-                    self._opik_client.span(**span_data.as_parameters)
+            # Send the trace/span to the backend with correct end_time and output
+            self._send_root_to_backend(root)
+
+        # Finalize span if it exists
+        if event_id in self._map_event_id_to_span_data:
+            span_data = self._map_event_id_to_span_data[event_id]
+
+            llm_usage_info = event_parsing_utils.get_usage_data(payload)
+            span_data.update(**llm_usage_info.__dict__)
+            span_data.update(output=span_output, error_info=error_info).init_end_time()
+
+            if tracing_runtime_config.is_tracing_active():
+                self._opik_client.span(**span_data.as_parameters)
 
-                del self._map_event_id_to_span_data[event_id]
+            self._opik_context_storage.pop_span_data(ensure_id=span_data.id)
+            del self._map_event_id_to_span_data[event_id]
 
     def flush(self) -> None:
-        """Sends pending Opik data to the backend"""
+        """Flush pending Opik data to backend."""
         self._opik_client.flush()

opik/integrations/openai/agents/opik_tracing_processor.py

@@ -3,12 +3,11 @@ from agents import tracing
 
 import logging
 
+from opik import context_storage, tracing_runtime_config
 from opik.api_objects.span import span_data
 from opik.api_objects.trace import trace_data
 from opik.api_objects import opik_client
 from opik.decorator import span_creation_handler, arguments_helpers
-import opik.decorator.tracing_runtime_config as tracing_runtime_config
-import opik.context_storage as context_storage
 
 from . import span_data_parsers
 

opik/integrations/openai/openai_chat_completions_decorator.py

@@ -56,9 +56,7 @@ class OpenaiChatCompletionsTrackDecorator(base_track_decorator.BaseTrackDecorato
         ), "Expected kwargs to be not None in chat.completion.create(**kwargs), chat.completion.parse(**kwargs) or chat.completion.stream(**kwargs)"
 
         name = track_options.name if track_options.name is not None else func.__name__
-        if _is_completions_stream_call(
-            name_passed_to_track_decorator=name, kwargs=kwargs
-        ):
+        if kwargs.get("stream") is True:
             kwargs = _remove_not_given_sentinel_values(kwargs)
             name = "chat_completion_stream"
 
@@ -196,17 +194,5 @@ def _remove_not_given_sentinel_values(dict_: Dict[str, Any]) -> Dict[str, Any]:
         key: value
         for key, value in dict_.items()
         if value is not _openai_types.NOT_GIVEN
+        and not isinstance(value, _openai_types.Omit)
     }
-
-
-def _is_completions_stream_call(
-    name_passed_to_track_decorator: str, kwargs: Dict[str, Any]
-) -> bool:
-    if not name_passed_to_track_decorator == "chat_completion_create":
-        return False
-
-    for _, value in kwargs.items():
-        if value is _openai_types.NOT_GIVEN:
-            return True
-
-    return False

opik/integrations/openai/opik_tracker.py

@@ -18,7 +18,7 @@ def track_openai(
     """Adds Opik tracking wrappers to an OpenAI client.
 
     The client is always patched; however every wrapped call checks
-    `opik.decorator.tracing_runtime_config.is_tracing_active()` before emitting
+    `opik.is_tracing_active()` before emitting
     any telemetry. If tracing is disabled at call time, the wrapped function
     executes normally but no span/trace is sent.
 

opik/integrations/sagemaker/auth.py

@@ -43,4 +43,8 @@ def setup_aws_sagemaker_session_hook() -> None:
 
         client.auth = sagemaker_auth
 
-    opik.hooks.register_httpx_client_hook(sagemaker_auth_client_hook)
+    opik.hooks.add_httpx_client_hook(
+        opik.hooks.HttpxClientHook(
+            client_modifier=sagemaker_auth_client_hook, client_init_arguments=None
+        )
+    )

opik/llm_usage/google_usage.py

@@ -1,11 +1,13 @@
 from typing import Optional, Dict, Any
+
+
 from . import base_original_provider_usage
 
 
 class GoogleGeminiUsage(base_original_provider_usage.BaseOriginalProviderUsage):
     """Google AI / VertexAI calls token usage data. Updated 11.03.2025"""
 
-    candidates_token_count: int
+    candidates_token_count: Optional[int]
     """Number of tokens in the response(s)."""
 
     prompt_token_count: int

opik/llm_usage/opik_usage.py

@@ -93,22 +93,21 @@ class OpikUsage(pydantic.BaseModel):
         # 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
 
-        if (
-            provider_usage.total_token_count
-            == provider_usage.prompt_token_count + provider_usage.candidates_token_count
-        ):
-            completion_tokens = provider_usage.candidates_token_count
+        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 = (
-                provider_usage.candidates_token_count
-                + provider_usage.thoughts_token_count
+                candidates_token_count + provider_usage.thoughts_token_count
             )
         else:
             LOGGER.debug(
                 "Something is wrong in Google usage, completion_tokens might be invalid"
             )
-            completion_tokens = provider_usage.candidates_token_count
+            completion_tokens = candidates_token_count
 
         return cls(
             completion_tokens=completion_tokens,

opik/llm_usage/opik_usage_factory.py

@@ -27,15 +27,17 @@ def build_opik_usage(
 ) -> opik_usage.OpikUsage:
     build_functions = _PROVIDER_TO_OPIK_USAGE_BUILDERS[provider]
 
+    exc = None
     for build_function in build_functions:
         try:
             result = build_function(usage)
             return result
-        except Exception:
+        except Exception as exc_info:
+            exc = exc_info
             pass
 
     raise ValueError(
-        f"Failed to build OpikUsage for provider {provider} and usage {usage}"
+        f"Failed to build OpikUsage for provider {provider} and usage {usage}, reason: {exc}"
     )
 
 

opik/logging_messages.py

@@ -29,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"
@@ -37,6 +39,10 @@ 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"

opik/message_processing/batching/base_batcher.py

@@ -1,4 +1,3 @@
-import threading
 import time
 import abc
 
@@ -24,25 +23,22 @@ class BaseBatcher(abc.ABC):
         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_messages = self._create_batches_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()
+            for batch_message in batch_messages:
+                self._flush_callback(batch_message)
+        self._last_time_flush_callback_called = time.time()
 
     def is_ready_to_flush(self) -> bool:
         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_batches_from_accumulated_messages(
@@ -51,10 +47,9 @@ class BaseBatcher(abc.ABC):
 
     @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]
@@ -65,10 +60,9 @@ class BaseBatcher(abc.ABC):
         Args:
             filter_func: A function that takes a BaseMessage and returns True if the message should be removed
         """
-        with self._lock:
-            self._accumulated_messages = list(
-                filter(lambda x: not filter_func(x), self._accumulated_messages)
-            )
+        self._accumulated_messages = list(
+            filter(lambda x: not filter_func(x), self._accumulated_messages)
+        )
 
     def size(self) -> int:
         """
@@ -81,5 +75,4 @@ class BaseBatcher(abc.ABC):
         Returns:
             int: The total number of accumulated messages.
         """
-        with self._lock:
-            return len(self._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

@@ -15,6 +15,9 @@ 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(
     queue: message_queue.MessageQueue[messages.BaseMessage],
@@ -55,6 +58,12 @@ def create_batch_manager(
         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[
         Type[messages.BaseMessage], base_batcher.BaseBatcher
     ] = {
@@ -64,6 +73,7 @@ def create_batch_manager(
         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(