opik 1.9.5__py3-none-any.whl → 1.9.39__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,31 +60,13 @@ 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)
-
-        return trace_data
-
     def start_trace(self, trace_id: Optional[str] = None) -> None:
         if (
             self._skip_index_construction_trace
@@ -94,18 +74,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 +106,26 @@ 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
+        root.init_end_time().update(output=last_event_output)
+
+        if isinstance(root, span.SpanData):
+            # We created a wrapper span (external trace existed)
+            self._opik_client.span(**root.as_parameters)
+            self._opik_context_storage.pop_span_data(ensure_id=root.id)
+        elif isinstance(root, trace.TraceData):
+            # We created a trace (no external trace existed)
+            self._opik_client.trace(**root.as_parameters)
+            self._opik_context_storage.pop_trace_data(ensure_id=root.id)
 
-        # 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()
+        # Clean up
+        _llama_root.set(None)
 
     def on_event_start(
         self,
@@ -142,53 +138,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 +202,31 @@ 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__)
+        # 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]
 
-                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)
+            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/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/message_processing/batching/batchers.py

@@ -1,11 +1,11 @@
 from typing import Union, cast, List
 
-from . import base_batcher, sequence_splitter
-from .. import messages
 from opik.rest_api.types import span_write, trace_write
-import opik.jsonable_encoder as jsonable_encoder
 import opik.dict_utils as dict_utils
 
+from . import base_batcher, sequence_splitter
+from .. import messages, encoder_helpers
+
 
 class CreateSpanMessageBatcher(base_batcher.BaseBatcher):
     def _create_batches_from_accumulated_messages(  # type: ignore
@@ -18,8 +18,10 @@ class CreateSpanMessageBatcher(base_batcher.BaseBatcher):
             cleaned_span_write_kwargs = dict_utils.remove_none_from_dict(
                 span_write_kwargs
             )
-            cleaned_span_write_kwargs = jsonable_encoder.encode(
-                cleaned_span_write_kwargs
+            cleaned_span_write_kwargs = encoder_helpers.encode_and_anonymize(
+                cleaned_span_write_kwargs,
+                fields_to_anonymize=messages.CreateSpansBatchMessage.fields_to_anonymize(),
+                object_type="span",
             )
             rest_spans.append(span_write.SpanWrite(**cleaned_span_write_kwargs))
 
@@ -52,8 +54,10 @@ class CreateTraceMessageBatcher(base_batcher.BaseBatcher):
             cleaned_trace_write_kwargs = dict_utils.remove_none_from_dict(
                 trace_write_kwargs
             )
-            cleaned_trace_write_kwargs = jsonable_encoder.encode(
-                cleaned_trace_write_kwargs
+            cleaned_trace_write_kwargs = encoder_helpers.encode_and_anonymize(
+                cleaned_trace_write_kwargs,
+                fields_to_anonymize=messages.CreateTraceBatchMessage.fields_to_anonymize(),
+                object_type="trace",
             )
             rest_traces.append(trace_write.TraceWrite(**cleaned_trace_write_kwargs))
 

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/messages.py

@@ -1,7 +1,7 @@
 import dataclasses
 import datetime
 from dataclasses import field
-from typing import Optional, Any, Dict, List, Union, Literal
+from typing import Optional, Any, Dict, List, Union, Literal, Set
 
 from . import arguments_utils
 from ..rest_api.types import span_write, trace_write
@@ -50,6 +50,10 @@ class CreateTraceMessage(BaseMessage):
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class UpdateTraceMessage(BaseMessage):
@@ -78,6 +82,10 @@ class UpdateTraceMessage(BaseMessage):
         data["id"] = data.pop("trace_id")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class CreateSpanMessage(BaseMessage):
@@ -112,6 +120,10 @@ class CreateSpanMessage(BaseMessage):
         data["total_estimated_cost"] = data.pop("total_cost")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class UpdateSpanMessage(BaseMessage):
@@ -144,6 +156,10 @@ class UpdateSpanMessage(BaseMessage):
         data["total_estimated_cost"] = data.pop("total_cost")
         return data
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class FeedbackScoreMessage(BaseMessage):
@@ -200,11 +216,19 @@ class AddThreadsFeedbackScoresBatchMessage(BaseMessage):
 class CreateSpansBatchMessage(BaseMessage):
     batch: List[span_write.SpanWrite]
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class CreateTraceBatchMessage(BaseMessage):
     batch: List[trace_write.TraceWrite]
 
+    @staticmethod
+    def fields_to_anonymize() -> Set[str]:
+        return {"input", "output", "metadata"}
+
 
 @dataclasses.dataclass
 class GuardrailBatchItemMessage(BaseMessage):

opik/message_processing/online_message_processor.py

@@ -4,12 +4,11 @@ from typing import Callable, Dict, Type, Any
 import pydantic
 import tenacity
 
-from . import messages, message_processors
+from . import encoder_helpers, messages, message_processors
 from .. import (
     exceptions as exceptions,
     logging_messages as logging_messages,
     dict_utils,
-    jsonable_encoder,
 )
 from ..rate_limit import rate_limit
 from ..rest_api import client as rest_api_client, core as rest_api_core
@@ -128,7 +127,12 @@ class OpikMessageProcessor(message_processors.BaseMessageProcessor):
         cleaned_create_span_kwargs = dict_utils.remove_none_from_dict(
             create_span_kwargs
         )
-        cleaned_create_span_kwargs = jsonable_encoder.encode(cleaned_create_span_kwargs)
+        cleaned_create_span_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_create_span_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="span",
+        )
+
         LOGGER.debug("Create span request: %s", cleaned_create_span_kwargs)
         self._rest_client.spans.create_span(**cleaned_create_span_kwargs)
 
@@ -140,9 +144,12 @@ class OpikMessageProcessor(message_processors.BaseMessageProcessor):
         cleaned_create_trace_kwargs = dict_utils.remove_none_from_dict(
             create_trace_kwargs
         )
-        cleaned_create_trace_kwargs = jsonable_encoder.encode(
-            cleaned_create_trace_kwargs
+        cleaned_create_trace_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_create_trace_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="trace",
         )
+
         LOGGER.debug("Create trace request: %s", cleaned_create_trace_kwargs)
         self._rest_client.traces.create_trace(**cleaned_create_trace_kwargs)
 
@@ -155,7 +162,12 @@ class OpikMessageProcessor(message_processors.BaseMessageProcessor):
         cleaned_update_span_kwargs = dict_utils.remove_none_from_dict(
             update_span_kwargs
         )
-        cleaned_update_span_kwargs = jsonable_encoder.encode(cleaned_update_span_kwargs)
+        cleaned_update_span_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_update_span_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="span",
+        )
+
         LOGGER.debug("Update span request: %s", cleaned_update_span_kwargs)
         self._rest_client.spans.update_span(**cleaned_update_span_kwargs)
 
@@ -168,9 +180,12 @@ class OpikMessageProcessor(message_processors.BaseMessageProcessor):
         cleaned_update_trace_kwargs = dict_utils.remove_none_from_dict(
             update_trace_kwargs
         )
-        cleaned_update_trace_kwargs = jsonable_encoder.encode(
-            cleaned_update_trace_kwargs
+        cleaned_update_trace_kwargs = encoder_helpers.encode_and_anonymize(
+            cleaned_update_trace_kwargs,
+            fields_to_anonymize=message.fields_to_anonymize(),
+            object_type="trace",
         )
+
         LOGGER.debug("Update trace request: %s", cleaned_update_trace_kwargs)
         self._rest_client.traces.update_trace(**cleaned_update_trace_kwargs)
         LOGGER.debug("Sent trace %s", message.trace_id)

opik/opik_context.py

@@ -1,7 +1,7 @@
 import contextlib
-from typing import Any, Dict, List, Optional, Union, Iterator
+from typing import Any, Dict, List, Optional, Iterator, Union
 
-import opik.llm_usage as llm_usage
+from opik import llm_usage
 from opik.api_objects import span, trace, opik_client, prompt
 from opik.api_objects.attachment import Attachment
 from opik.types import (
@@ -11,7 +11,7 @@ from opik.types import (
     ErrorInfoDict,
 )
 
-import opik.decorator.tracing_runtime_config as tracing_runtime_config
+from opik import tracing_runtime_config
 
 from . import context_storage, exceptions
 from .decorator import error_info_collector
@@ -68,7 +68,7 @@ def update_current_span(
     total_cost: Optional[float] = None,
     attachments: Optional[List[Attachment]] = None,
     error_info: Optional[ErrorInfoDict] = None,
-    prompts: Optional[List[prompt.Prompt]] = None,
+    prompts: Optional[List[prompt.BasePrompt]] = None,
 ) -> None:
     """
     Update the current span with the provided parameters. This method is usually called within a tracked function.
@@ -97,7 +97,7 @@ def update_current_span(
         return
 
     if prompts is not None:
-        prompts = [prompt.prompt.to_info_dict(p) for p in prompts]
+        prompts = [p.__internal_api__to_info_dict__() for p in prompts]
 
     new_params = {
         "name": name,
@@ -130,7 +130,7 @@ def update_current_trace(
     feedback_scores: Optional[List[FeedbackScoreDict]] = None,
     thread_id: Optional[str] = None,
     attachments: Optional[List[Attachment]] = None,
-    prompts: Optional[List[prompt.Prompt]] = None,
+    prompts: Optional[List[prompt.BasePrompt]] = None,
 ) -> None:
     """
     Update the current trace with the provided parameters. This method is usually called within a tracked function.
@@ -151,7 +151,7 @@ def update_current_trace(
         return
 
     if prompts is not None:
-        prompts = [prompt.prompt.to_info_dict(p) for p in prompts]
+        prompts = [p.__internal_api__to_info_dict__() for p in prompts]
 
     new_params = {
         "name": name,