[OPIK-443] [SDK] Add "model" and "provider" fields to spans (#740)

* OPIK-443 [SDK] Add "model" and "provider" fields to spans * add "model" and "provider" fields to spans in backend emulator * add e2e test * add incorrect response workaround * exclude problematic httpx library version * exclude problematic httpx library version v2 --------- Co-authored-by: Aliaksandr Kuzmik <[email protected]>
comet-ml · Nov 28, 2024 · 2dcab92 · 2dcab92
1 parent a1623c3
commit 2dcab92
Show file tree

Hide file tree

Showing 13 changed files with 75 additions and 2 deletions.
diff --git a/sdks/python/setup.py b/sdks/python/setup.py
@@ -37,7 +37,7 @@
     long_description_content_type="text/markdown",
     install_requires=[
         "click",
-        "httpx<1.0.0",
+        "httpx<0.28.0",
         "levenshtein<1.0.0",
         "litellm",
         "openai<2.0.0",

diff --git a/sdks/python/src/opik/api_objects/opik_client.py b/sdks/python/src/opik/api_objects/opik_client.py
@@ -196,6 +196,8 @@ def span(
         usage: Optional[UsageDict] = None,
         feedback_scores: Optional[List[FeedbackScoreDict]] = None,
         project_name: Optional[str] = None,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
     ) -> span.Span:
         """
         Create and log a new span.
@@ -216,6 +218,8 @@ def span(
             feedback_scores: The list of feedback score dicts associated with the span. Dicts don't require to have an `id` value.
             project_name: The name of the project. If not set, the project name which was configured when Opik instance
                 was created will be used.
+            model: The name of LLM (in this case `type` parameter should be == `llm`)
+            provider: The provider of LLM.
 
         Returns:
             span.Span: The created span object.
@@ -267,6 +271,8 @@ def span(
             metadata=metadata,
             tags=tags,
             usage=parsed_usage.supported_usage,
+            model=model,
+            provider=provider,
         )
         self._streamer.put(create_span_message)
 
@@ -626,7 +632,21 @@ def get_span_content(self, id: str) -> span_public.SpanPublic:
             span_public.SpanPublic: pydantic model object with all the data associated with the span found.
             Raises an error if span was not found.
         """
-        return self._rest_client.spans.get_span_by_id(id)
+        result = self._rest_client.spans.get_span_by_id(id)
+
+        # fixme temporary fix for wrong response payload
+        # because span_public.SpanPublic is frozen we will create a copy and update it
+        new_values: Dict[str, Any] = {}
+
+        if result.model == "":
+            new_values["model"] = None
+        if result.provider == "":
+            new_values["provider"] = None
+
+        if len(new_values) > 0:
+            result = result.model_copy(update=new_values)
+
+        return result
 
     def get_project(self, id: str) -> project_public.ProjectPublic:
         """

diff --git a/sdks/python/src/opik/api_objects/span.py b/sdks/python/src/opik/api_objects/span.py
@@ -78,6 +78,8 @@ def update(
         output: Optional[Dict[str, Any]] = None,
         tags: Optional[List[str]] = None,
         usage: Optional[UsageDict] = None,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
     ) -> None:
         """
         Update the span attributes.
@@ -89,6 +91,8 @@ def update(
             output: The output data for the span.
             tags: A list of tags to be associated with the span.
             usage: Usage information for the span.
+            model: The name of LLM.
+            provider: The provider of LLM.
 
         Returns:
             None
@@ -112,6 +116,8 @@ def update(
             output=output,
             tags=tags,
             usage=parsed_usage.supported_usage,
+            model=model,
+            provider=provider,
         )
         self._streamer.put(end_span_message)
 
@@ -127,6 +133,8 @@ def span(
         output: Optional[Dict[str, Any]] = None,
         tags: Optional[List[str]] = None,
         usage: Optional[UsageDict] = None,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
     ) -> "Span":
         """
         Create a new child span within the current span.
@@ -142,6 +150,8 @@ def span(
             output: The output data for the span.
             tags: A list of tags to be associated with the span.
             usage: Usage information for the span.
+            model: The name of LLM (in this case `type` parameter should be == `llm`)
+            provider: The provider of LLM.
 
         Returns:
             Span: The created child span object.
@@ -172,6 +182,8 @@ def span(
             metadata=metadata,
             tags=tags,
             usage=parsed_usage.supported_usage,
+            model=model,
+            provider=provider,
         )
         self._streamer.put(create_span_message)
 
@@ -255,6 +267,8 @@ class SpanData:
     usage: Optional[UsageDict] = None
     feedback_scores: Optional[List[FeedbackScoreDict]] = None
     project_name: Optional[str] = None
+    model: Optional[str] = None
+    provider: Optional[str] = None
 
     def update(self, **new_data: Any) -> "SpanData":
         for key, value in new_data.items():

diff --git a/sdks/python/src/opik/api_objects/trace.py b/sdks/python/src/opik/api_objects/trace.py
@@ -107,6 +107,8 @@ def span(
         output: Optional[Dict[str, Any]] = None,
         tags: Optional[List[str]] = None,
         usage: Optional[UsageDict] = None,
+        model: Optional[str] = None,
+        provider: Optional[str] = None,
     ) -> span.Span:
         """
         Create a new span within the trace.
@@ -123,6 +125,8 @@ def span(
             output: The output data for the span.
             tags: A list of tags to be associated with the span.
             usage: Usage information for the span.
+            model: The name of LLM (in this case `type` parameter should be == `llm`)
+            provider: The provider of LLM.
 
         Returns:
             span.Span: The created span object.
@@ -153,6 +157,8 @@ def span(
             metadata=metadata,
             tags=tags,
             usage=parsed_usage.supported_usage,
+            model=model,
+            provider=provider,
         )
         self._streamer.put(create_span_message)
 

diff --git a/sdks/python/src/opik/decorator/arguments_helpers.py b/sdks/python/src/opik/decorator/arguments_helpers.py
@@ -28,6 +28,8 @@ class EndSpanParameters(BaseArguments):
     output: Optional[Dict[str, Any]] = None
     tags: Optional[List[str]] = None
     usage: Optional[Dict[str, Any]] = None
+    model: Optional[str] = None
+    provider: Optional[str] = None
 
 
 @dataclasses.dataclass

diff --git a/sdks/python/src/opik/message_processing/message_processors.py b/sdks/python/src/opik/message_processing/message_processors.py
@@ -79,6 +79,8 @@ def _process_create_span_message(self, message: messages.CreateSpanMessage) -> N
             "metadata": message.metadata,
             "tags": message.tags,
             "usage": message.usage,
+            "model": message.model,
+            "provider": message.provider,
         }
 
         cleaned_create_span_kwargs = dict_utils.remove_none_from_dict(
@@ -122,6 +124,8 @@ def _process_update_span_message(self, message: messages.UpdateSpanMessage) -> N
             "metadata": message.metadata,
             "tags": message.tags,
             "usage": message.usage,
+            "model": message.model,
+            "provider": message.provider,
         }
 
         cleaned_update_span_kwargs = dict_utils.remove_none_from_dict(
@@ -202,6 +206,8 @@ def _process_create_span_batch_message(
                 "metadata": item.metadata,
                 "tags": item.tags,
                 "usage": item.usage,
+                "model": item.model,
+                "provider": item.provider,
             }
             cleaned_span_write_kwargs = dict_utils.remove_none_from_dict(
                 span_write_kwargs

diff --git a/sdks/python/src/opik/message_processing/messages.py b/sdks/python/src/opik/message_processing/messages.py
@@ -48,6 +48,8 @@ class CreateSpanMessage(BaseMessage):
     tags: Optional[List[str]]
     type: SpanType
     usage: Optional[UsageDict]
+    model: Optional[str]
+    provider: Optional[str]
 
 
 @dataclasses.dataclass
@@ -62,6 +64,8 @@ class UpdateSpanMessage(BaseMessage):
     metadata: Optional[Dict[str, Any]]
     tags: Optional[List[str]]
     usage: Optional[UsageDict]
+    model: Optional[str]
+    provider: Optional[str]
 
 
 @dataclasses.dataclass

diff --git a/sdks/python/tests/e2e/test_tracing.py b/sdks/python/tests/e2e/test_tracing.py
@@ -172,6 +172,9 @@ def f_inner(y):
 def test_manually_created_trace_and_span__happyflow(
     opik_client: opik.Opik, project_name
 ):
+    model_name = "some-llm"
+    provider_name = "some-llm-provider"
+
     # Call
     trace = opik_client.trace(
         name="trace-name",
@@ -187,6 +190,8 @@ def test_manually_created_trace_and_span__happyflow(
         output={"output": "span-output"},
         tags=["span-tag"],
         metadata={"span-metadata-key": "span-metadata-value"},
+        model=model_name,
+        provider=provider_name,
     )
 
     opik_client.flush()
@@ -215,6 +220,8 @@ def test_manually_created_trace_and_span__happyflow(
         tags=["span-tag"],
         metadata={"span-metadata-key": "span-metadata-value"},
         project_name=project_name or OPIK_E2E_TESTS_PROJECT_NAME,
+        model=model_name,
+        provider=provider_name,
     )
 
 

diff --git a/sdks/python/tests/e2e/verifiers.py b/sdks/python/tests/e2e/verifiers.py
@@ -88,6 +88,8 @@ def verify_span(
     type: str = mock.ANY,  # type: ignore
     feedback_scores: List[FeedbackScoreDict] = mock.ANY,  # type: ignore
     project_name: Optional[str] = mock.ANY,
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
 ):
     if not synchronization.until(
         lambda: (opik_client.get_span_content(id=span_id) is not None),
@@ -115,6 +117,8 @@ def verify_span(
         span.metadata, metadata
     )
     assert span.tags == tags, testlib.prepare_difference_report(span.tags, tags)
+    assert span.model == model
+    assert span.provider == provider
 
     if project_name is not mock.ANY:
         span_project = opik_client.get_project(span.project_id)

diff --git a/sdks/python/tests/testlib/backend_emulator_message_processor.py b/sdks/python/tests/testlib/backend_emulator_message_processor.py
@@ -122,6 +122,8 @@ def process(self, message: messages.BaseMessage) -> None:
                 end_time=message.end_time,
                 usage=message.usage,
                 project_name=message.project_name,
+                model=message.model,
+                provider=message.provider,
             )
 
             self._span_to_parent_span[span.id] = message.parent_span_id

diff --git a/sdks/python/tests/testlib/models.py b/sdks/python/tests/testlib/models.py
@@ -25,6 +25,8 @@ class SpanModel:
     feedback_scores: List["FeedbackScoreModel"] = dataclasses.field(
         default_factory=list
     )
+    model: Optional[str] = None
+    provider: Optional[str] = None
 
 
 @dataclasses.dataclass

diff --git a/sdks/python/tests/unit/message_processing/batching/test_batch_manager.py b/sdks/python/tests/unit/message_processing/batching/test_batch_manager.py
@@ -97,6 +97,8 @@ def test_batch_manager__start_and_stop_were_called__accumulated_data_is_flushed(
         tags=NOT_USED,
         type=NOT_USED,
         usage=NOT_USED,
+        model=NOT_USED,
+        provider=NOT_USED,
     )
 
     example_span_batcher = create_span_message_batcher.CreateSpanMessageBatcher(

diff --git a/sdks/python/tests/unit/message_processing/test_message_streaming.py b/sdks/python/tests/unit/message_processing/test_message_streaming.py
@@ -56,6 +56,8 @@ def test_streamer__batching_disabled__messages_that_support_batching_are_process
             tags=NOT_USED,
             type=NOT_USED,
             usage=NOT_USED,
+            model=NOT_USED,
+            provider=NOT_USED,
         )
 
         tested.put(CREATE_SPAN_MESSAGE)
@@ -97,6 +99,8 @@ def test_streamer__batching_enabled__messages_that_support_batching_are_processe
             tags=NOT_USED,
             type=NOT_USED,
             usage=NOT_USED,
+            model=NOT_USED,
+            provider=NOT_USED,
         )
 
         tested.put(CREATE_SPAN_MESSAGE)