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

opik/api_objects/search_helpers.py

@@ -0,0 +1,89 @@
+from typing import Optional, List, Callable, Any
+
+from opik import synchronization
+from opik.api_objects import rest_stream_parser
+from opik.api_objects.helpers import OptionalFilterParsedItemList
+from opik.rest_api import client as rest_api_client
+from opik.rest_api.types import span_public, trace_public
+
+
+def search_spans_with_filters(
+    rest_client: rest_api_client.OpikApi,
+    trace_id: Optional[str],
+    project_name: str,
+    filters: Optional[OptionalFilterParsedItemList],
+    max_results: int,
+    truncate: bool,
+) -> List[span_public.SpanPublic]:
+    spans = rest_stream_parser.read_and_parse_full_stream(
+        read_source=lambda current_batch_size,
+        last_retrieved_id: rest_client.spans.search_spans(
+            trace_id=trace_id,
+            project_name=project_name,
+            filters=filters,
+            limit=current_batch_size,
+            truncate=truncate,
+            last_retrieved_id=last_retrieved_id,
+        ),
+        max_results=max_results,
+        parsed_item_class=span_public.SpanPublic,
+    )
+
+    return spans
+
+
+def search_traces_with_filters(
+    rest_client: rest_api_client.OpikApi,
+    project_name: Optional[str],
+    filters: Optional[OptionalFilterParsedItemList],
+    max_results: int,
+    truncate: bool,
+) -> List[trace_public.TracePublic]:
+    traces = rest_stream_parser.read_and_parse_full_stream(
+        read_source=lambda current_batch_size,
+        last_retrieved_id: rest_client.traces.search_traces(
+            project_name=project_name,
+            filters=filters,
+            limit=current_batch_size,
+            truncate=truncate,
+            last_retrieved_id=last_retrieved_id,
+        ),
+        max_results=max_results,
+        parsed_item_class=trace_public.TracePublic,
+    )
+    return traces
+
+
+def search_and_wait_for_done(
+    search_functor: Callable[[], List[Any]],
+    wait_for_at_least: int,
+    wait_for_timeout: int,
+    sleep_time: float,
+) -> List[Any]:
+    """
+    The expected behavior is to keep making repeated calls until either the specified number of
+    results is found or the timeout is reached. The function will then return the best possible
+    attempt results to meet these conditions.
+    Args:
+        search_functor: The function to call to retrieve the results.
+        wait_for_at_least: The minimum number of results to return.
+        wait_for_timeout: The timeout for waiting for results.
+        sleep_time: The time to sleep between calls to search_functor.
+
+    Returns:
+        The function returns the results of the best possible attempt to meet both waiting conditions.
+    """
+    result: List[Any] = []
+
+    def search() -> List[Any]:
+        nonlocal result
+        result = search_functor()
+        return result
+
+    synchronization.wait_for_done(
+        check_function=lambda: len(search()) >= wait_for_at_least,
+        timeout=wait_for_timeout,
+        sleep_time=sleep_time,
+    )
+
+    return result

opik/api_objects/span/span_client.py

@@ -2,9 +2,14 @@ import datetime
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-from opik import datetime_helpers, id_helpers, llm_usage
+import opik.datetime_helpers as datetime_helpers
+import opik.id_helpers as id_helpers
+import opik.llm_usage as llm_usage
+import opik.api_objects.attachment as attachment
 from opik.message_processing import messages, streamer
 
+from ..attachment import converters as attachment_converters
+
 from opik.types import (
     DistributedTraceHeadersDict,
     ErrorInfoDict,
@@ -23,6 +28,7 @@ class Span:
         trace_id: str,
         project_name: str,
         message_streamer: streamer.Streamer,
+        url_override: str,
         parent_span_id: Optional[str] = None,
     ):
         """
@@ -33,6 +39,7 @@ class Span:
         self.parent_span_id = parent_span_id
         self._streamer = message_streamer
         self._project_name = project_name
+        self._url_override = url_override
 
     def end(
         self,
@@ -126,32 +133,24 @@ class Span:
         Returns:
             None
         """
-        backend_compatible_usage = validation_helpers.validate_and_parse_usage(
-            usage=usage,
-            logger=LOGGER,
-            provider=provider,
-        )
-
-        if backend_compatible_usage is not None:
-            metadata = helpers.add_usage_to_metadata(usage=usage, metadata=metadata)
-
-        end_span_message = messages.UpdateSpanMessage(
-            span_id=self.id,
+        update_span(
+            id=self.id,
             trace_id=self.trace_id,
             parent_span_id=self.parent_span_id,
+            url_override=self._url_override,
+            message_streamer=self._streamer,
             project_name=self._project_name,
             end_time=end_time,
             metadata=metadata,
             input=input,
             output=output,
             tags=tags,
-            usage=backend_compatible_usage,
+            usage=usage,
             model=model,
             provider=provider,
             error_info=error_info,
             total_cost=total_cost,
         )
-        self._streamer.put(end_span_message)
 
     def span(
         self,
@@ -169,74 +168,57 @@ class Span:
         provider: LLMProvider = LLMProvider.OPENAI,
         error_info: Optional[ErrorInfoDict] = None,
         total_cost: Optional[float] = None,
+        attachments: Optional[List[attachment.Attachment]] = None,
     ) -> "Span":
         """
         Create a new child span within the current span.
 
         Args:
-            id: The ID of the span, should be in UUIDv7 format. If not provided, a new ID will be generated.
+            id: The ID of the span should be in UUIDv7 format. If not provided, a new ID will be generated.
             name: The name of the span.
             type: The type of the span. Defaults to "general".
-            start_time: The start time of the span. If not provided, current time will be used.
+            start_time: The start time of the span. If not provided, the current time will be used.
             end_time: The end time of the span.
             metadata: Additional metadata to be associated with the span.
             input: The input data for the span.
             output: The output data for the span.
             tags: A list of tags to be associated with the span.
             usage: Usage data for the span. In order for input, output and total tokens to be visible in the UI,
-                the usage must contain OpenAI-formatted keys (they can be passed additionaly to original usage on the top level of the dict):  prompt_tokens, completion_tokens and total_tokens.
+                the usage must contain OpenAI-formatted keys (they can be passed additionally to the original usage on the top level of the dict): prompt_tokens, completion_tokens and total_tokens.
                 If OpenAI-formatted keys were not found, Opik will try to calculate them automatically if the usage
                 format is recognized (you can see which provider's formats are recognized in opik.LLMProvider enum), but it is not guaranteed.
             model: The name of LLM (in this case `type` parameter should be == `llm`)
             provider: The provider of LLM. You can find providers officially supported by Opik for cost tracking
-                in `opik.LLMProvider` enum. If your provider is not here, please open an issue in our github - https://github.com/comet-ml/opik.
-                If your provider not in the list, you can still specify it but the cost tracking will not be available
+                in `opik.LLMProvider` enum. If your provider is not here, please open an issue in our GitHub - https://github.com/comet-ml/opik.
+                If your provider is not in the list, you can still specify it, but the cost tracking will not be available
             error_info: The dictionary with error information (typically used when the span function has failed).
             total_cost: The cost of the span in USD. This value takes priority over the cost calculated by Opik from the usage.
+            attachments: The list of attachments to be uploaded to the span.
 
         Returns:
             Span: The created child span object.
         """
-        span_id = id if id is not None else id_helpers.generate_id()
-        start_time = (
-            start_time if start_time is not None else datetime_helpers.local_timestamp()
-        )
-        backend_compatible_usage = validation_helpers.validate_and_parse_usage(
-            usage=usage,
-            logger=LOGGER,
-            provider=provider,
-        )
-
-        if backend_compatible_usage is not None:
-            metadata = helpers.add_usage_to_metadata(usage=usage, metadata=metadata)
-
-        create_span_message = messages.CreateSpanMessage(
-            span_id=span_id,
+        return create_span(
             trace_id=self.trace_id,
             project_name=self._project_name,
+            url_override=self._url_override,
+            message_streamer=self._streamer,
+            span_id=id,
             parent_span_id=self.id,
             name=name,
             type=type,
             start_time=start_time,
             end_time=end_time,
+            metadata=metadata,
             input=input,
             output=output,
-            metadata=metadata,
             tags=tags,
-            usage=backend_compatible_usage,
+            usage=usage,
             model=model,
             provider=provider,
             error_info=error_info,
             total_cost=total_cost,
-        )
-        self._streamer.put(create_span_message)
-
-        return Span(
-            id=span_id,
-            parent_span_id=self.id,
-            trace_id=self.trace_id,
-            message_streamer=self._streamer,
-            project_name=self._project_name,
+            attachments=attachments,
         )
 
     def log_feedback_score(
@@ -280,3 +262,141 @@ class Span:
         function on remote node.
         """
         return {"opik_parent_span_id": self.id, "opik_trace_id": self.trace_id}
+
+
+def create_span(
+    trace_id: str,
+    project_name: str,
+    url_override: str,
+    message_streamer: streamer.Streamer,
+    span_id: Optional[str] = None,
+    parent_span_id: Optional[str] = None,
+    name: Optional[str] = None,
+    type: SpanType = "general",
+    start_time: Optional[datetime.datetime] = None,
+    end_time: Optional[datetime.datetime] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    input: Optional[Dict[str, Any]] = None,
+    output: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    usage: Optional[Union[Dict[str, Any], llm_usage.OpikUsage]] = None,
+    model: Optional[str] = None,
+    provider: Optional[Union[LLMProvider, str]] = None,
+    error_info: Optional[ErrorInfoDict] = None,
+    total_cost: Optional[float] = None,
+    attachments: Optional[List[attachment.Attachment]] = None,
+) -> Span:
+    span_id = span_id if span_id is not None else id_helpers.generate_id()
+    start_time = (
+        start_time if start_time is not None else datetime_helpers.local_timestamp()
+    )
+
+    backend_compatible_usage = validation_helpers.validate_and_parse_usage(
+        usage=usage,
+        logger=LOGGER,
+        provider=provider,
+    )
+
+    if backend_compatible_usage is not None:
+        metadata = helpers.add_usage_to_metadata(usage=usage, metadata=metadata)
+
+    create_span_message = messages.CreateSpanMessage(
+        span_id=span_id,
+        trace_id=trace_id,
+        project_name=project_name,
+        parent_span_id=parent_span_id,
+        name=name,
+        type=type,
+        start_time=start_time,
+        end_time=end_time,
+        input=input,
+        output=output,
+        metadata=metadata,
+        tags=tags,
+        usage=backend_compatible_usage,
+        model=model,
+        provider=provider,
+        error_info=error_info,
+        total_cost=total_cost,
+        last_updated_at=datetime_helpers.local_timestamp(),
+    )
+    message_streamer.put(create_span_message)
+
+    if attachments is not None:
+        for attachment_data in attachments:
+            create_attachment_message = attachment_converters.attachment_to_message(
+                attachment_data=attachment_data,
+                entity_type="span",
+                entity_id=span_id,
+                project_name=project_name,
+                url_override=url_override,
+            )
+            message_streamer.put(create_attachment_message)
+
+    return Span(
+        id=span_id,
+        parent_span_id=parent_span_id,
+        trace_id=trace_id,
+        message_streamer=message_streamer,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+
+def update_span(
+    id: str,
+    trace_id: str,
+    parent_span_id: Optional[str],
+    project_name: str,
+    url_override: str,
+    message_streamer: streamer.Streamer,
+    end_time: Optional[datetime.datetime] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    input: Optional[Dict[str, Any]] = None,
+    output: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    usage: Optional[Union[Dict[str, Any], llm_usage.OpikUsage]] = None,
+    model: Optional[str] = None,
+    provider: Optional[Union[LLMProvider, str]] = None,
+    error_info: Optional[ErrorInfoDict] = None,
+    total_cost: Optional[float] = None,
+    attachments: Optional[List[attachment.Attachment]] = None,
+) -> None:
+    backend_compatible_usage = validation_helpers.validate_and_parse_usage(
+        usage=usage,
+        logger=LOGGER,
+        provider=provider,
+    )
+
+    if backend_compatible_usage is not None:
+        metadata = helpers.add_usage_to_metadata(usage=usage, metadata=metadata)
+
+    update_span_message = messages.UpdateSpanMessage(
+        span_id=id,
+        trace_id=trace_id,
+        parent_span_id=parent_span_id,
+        project_name=project_name,
+        end_time=end_time,
+        metadata=metadata,
+        input=input,
+        output=output,
+        tags=tags,
+        usage=backend_compatible_usage,
+        model=model,
+        provider=provider,
+        error_info=error_info,
+        total_cost=total_cost,
+    )
+
+    if attachments is not None:
+        for attachment_data in attachments:
+            create_attachment_message = attachment_converters.attachment_to_message(
+                attachment_data=attachment_data,
+                entity_type="span",
+                entity_id=id,
+                project_name=project_name,
+                url_override=url_override,
+            )
+            message_streamer.put(create_attachment_message)
+
+    message_streamer.put(update_span_message)

opik/api_objects/span/span_data.py

@@ -3,15 +3,17 @@ import datetime
 import logging
 from typing import Any, Dict, List, Optional, Union
 
-from opik import datetime_helpers, dict_utils, llm_usage
+import opik.api_objects.attachment as attachment
+import opik.datetime_helpers as datetime_helpers
+import opik.llm_usage as llm_usage
 from opik.types import (
+    DistributedTraceHeadersDict,
     ErrorInfoDict,
     FeedbackScoreDict,
     LLMProvider,
     SpanType,
 )
-
-from .. import helpers
+from .. import helpers, data_helpers
 
 LOGGER = logging.getLogger(__name__)
 
@@ -49,13 +51,69 @@ class SpanData:
     provider: Optional[Union[str, LLMProvider]] = None
     error_info: Optional[ErrorInfoDict] = None
     total_cost: Optional[float] = None
+    attachments: Optional[List[attachment.Attachment]] = None
+
+    def create_child_span_data(
+        self,
+        name: Optional[str] = None,
+        type: SpanType = "general",
+        start_time: Optional[datetime.datetime] = None,
+        end_time: Optional[datetime.datetime] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        input: Optional[Dict[str, Any]] = None,
+        output: Optional[Dict[str, Any]] = None,
+        tags: Optional[List[str]] = None,
+        usage: Optional[Union[Dict[str, Any], llm_usage.OpikUsage]] = None,
+        feedback_scores: Optional[List[FeedbackScoreDict]] = None,
+        model: Optional[str] = None,
+        provider: Optional[Union[str, LLMProvider]] = None,
+        error_info: Optional[ErrorInfoDict] = None,
+        total_cost: Optional[float] = None,
+        attachments: Optional[List[attachment.Attachment]] = None,
+    ) -> "SpanData":
+        start_time = (
+            start_time if start_time is not None else datetime_helpers.local_timestamp()
+        )
+        return SpanData(
+            trace_id=self.trace_id,
+            parent_span_id=self.id,
+            project_name=self.project_name,
+            name=name,
+            type=type,
+            start_time=start_time,
+            end_time=end_time,
+            metadata=metadata,
+            input=input,
+            output=output,
+            tags=tags,
+            usage=usage,
+            feedback_scores=feedback_scores,
+            model=model,
+            provider=provider,
+            error_info=error_info,
+            total_cost=total_cost,
+            attachments=attachments,
+        )
 
     def update(self, **new_data: Any) -> "SpanData":
+        """
+        Updates the attributes of the object with the provided key-value pairs. This method checks if
+        an attribute exists before updating it and merges the data appropriately for specific
+        keywords like metadata, output, input, attachments, and tags. If a key doesn't correspond
+        to an attribute of the object or the provided value is None, the update is skipped.
+
+        Args:
+            **new_data: Key-value pairs of attributes to update. Keys should match existing
+                attributes on the object, and values that are None will not update.
+
+        Returns:
+            SpanData: The updated object instance.
+        """
         for key, value in new_data.items():
             if value is None:
                 continue
 
-            if key not in self.__dict__:
+            if key not in self.__dict__ and key != "prompts":
                 LOGGER.debug(
                     "An attempt to update span with parameter name it doesn't have: %s",
                     key,
@@ -63,38 +121,91 @@ class SpanData:
                 continue
 
             if key == "metadata":
-                self._update_metadata(value)
+                self.metadata = data_helpers.merge_metadata(
+                    self.metadata, new_metadata=value
+                )
                 continue
             elif key == "output":
-                self._update_output(value)
+                self.output = data_helpers.merge_outputs(self.output, new_outputs=value)
                 continue
             elif key == "input":
-                self._update_input(value)
+                self.input = data_helpers.merge_inputs(self.input, new_inputs=value)
+                continue
+            elif key == "attachments":
+                self._update_attachments(value)
+                continue
+            elif key == "tags":
+                self.tags = data_helpers.merge_tags(self.tags, new_tags=value)
+                continue
+            elif key == "prompts":
+                self.metadata = data_helpers.merge_metadata(
+                    self.metadata, new_metadata=new_data.get("metadata"), prompts=value
+                )
                 continue
 
             self.__dict__[key] = value
 
         return self
 
-    def _update_metadata(self, new_metadata: Dict[str, Any]) -> None:
-        if self.metadata is None:
-            self.metadata = new_metadata
-        else:
-            self.metadata = dict_utils.deepmerge(self.metadata, new_metadata)
-
-    def _update_output(self, new_output: Dict[str, Any]) -> None:
-        if self.output is None:
-            self.output = new_output
-        else:
-            self.output = dict_utils.deepmerge(self.output, new_output)
-
-    def _update_input(self, new_input: Dict[str, Any]) -> None:
-        if self.input is None:
-            self.input = new_input
-        else:
-            self.input = dict_utils.deepmerge(self.input, new_input)
-
     def init_end_time(self) -> "SpanData":
         self.end_time = datetime_helpers.local_timestamp()
 
         return self
+
+    def _update_attachments(self, attachments: List[attachment.Attachment]) -> None:
+        if self.attachments is None:
+            self.attachments = attachments
+        else:
+            self.attachments.extend(attachments)
+
+    @property
+    def as_start_parameters(self) -> Dict[str, Any]:
+        """Returns parameters of this span to be sent to the server when starting a new span."""
+        start_parameters: Dict[str, Any] = {
+            "id": self.id,
+            "start_time": self.start_time,
+            "project_name": self.project_name,
+            "trace_id": self.trace_id,
+        }
+
+        if self.name is not None:
+            start_parameters["name"] = self.name
+        if self.input is not None:
+            start_parameters["input"] = self.input
+        if self.metadata is not None:
+            start_parameters["metadata"] = self.metadata
+        if self.tags is not None:
+            start_parameters["tags"] = self.tags
+
+        return start_parameters
+
+    @property
+    def as_parameters(self) -> Dict[str, Any]:
+        """Returns all parameters of this span to be sent to the server."""
+        return {
+            "trace_id": self.trace_id,
+            "id": self.id,
+            "parent_span_id": self.parent_span_id,
+            "name": self.name,
+            "type": self.type,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "metadata": self.metadata,
+            "input": self.input,
+            "output": self.output,
+            "tags": self.tags,
+            "usage": self.usage,
+            "feedback_scores": self.feedback_scores,
+            "project_name": self.project_name,
+            "model": self.model,
+            "provider": self.provider,
+            "error_info": self.error_info,
+            "total_cost": self.total_cost,
+            "attachments": self.attachments,
+        }
+
+    def get_distributed_trace_headers(self) -> DistributedTraceHeadersDict:
+        return DistributedTraceHeadersDict(
+            opik_trace_id=self.trace_id,
+            opik_parent_span_id=self.id,
+        )