fmt

ddtrace/llmobs/_integrations/base.py

@@ -197,12 +197,19 @@ def trunc(self, text: str) -> str:
             text = text[: self.integration_config.span_char_limit] + "..."
         return text
 
-    def llmobs_set_tags(self, span: Span, args: List[Any], kwargs: Dict[str, Any], **extra_kwargs) -> None:
+    def llmobs_set_tags(
+        self,
+        span: Span,
+        args: List[Any],
+        kwargs: Dict[str, Any],
+        response: Optional[Any] = None,
+        operation: str = "",
+    ) -> None:
         """Extract input/output information from the request and response to be submitted to LLMObs."""
         if not self.llmobs_enabled:
             return
         try:
-            self._llmobs_set_tags(span, args, kwargs, **extra_kwargs)
+            self._llmobs_set_tags(span, args, kwargs, response, operation)
         except Exception:
             log.error("Error extracting LLMObs fields for span %s, likely due to malformed data", span, exc_info=True)
 
@@ -215,4 +222,4 @@ def _llmobs_set_tags(
         response: Optional[Any] = None,
         operation: str = "",
     ) -> None:
-        pass
+        raise NotImplementedError()

ddtrace/llmobs/_integrations/langchain.py

@@ -59,7 +59,7 @@ def _llmobs_set_tags(
         args: List[Any],
         kwargs: Dict[str, Any],
         response: Optional[Any] = None,
-        operation: Optional[str] = None,  # oneof "llm","chat","chain","embedding","retrieval","tool"
+        operation: str = "",  # oneof "llm","chat","chain","embedding","retrieval","tool"
     ) -> None:
         """Sets meta tags and metrics for span events to be sent to LLMObs."""
         if not self.llmobs_enabled:

ddtrace/llmobs/_integrations/openai.py

@@ -144,7 +144,7 @@ def _llmobs_set_tags(
         args: List[Any],
         kwargs: Dict[str, Any],
         response: Optional[Any] = None,
-        operation: str = "completion",  # oneof "completion", "chat", "embedding"
+        operation: str = "",  # oneof "completion", "chat", "embedding"
     ) -> None:
         """Sets meta tags and metrics for span events to be sent to LLMObs."""
         span_kind = "embedding" if operation == "embedding" else "llm"

tests/llmobs/test_llmobs_integrations.py

@@ -198,12 +198,16 @@ def test_llmobs_set_tags(mock_llmobs, mock_log, mock_integration_config):
     span = DummyTracer().trace("Dummy span", service="dummy_service")
     integration = BaseLLMIntegration(mock_integration_config)
     integration._llmobs_set_tags = mock.Mock()
-    integration.llmobs_set_tags(span, kwarg_1="value_1", kwarg_2="value_2")
-    integration._llmobs_set_tags.assert_called_once_with(span, kwarg_1="value_1", kwarg_2="value_2")
+    integration.llmobs_set_tags(span, args=[], kwargs={}, response="response", operation="operation")
+    integration._llmobs_set_tags.assert_called_once_with(span, [], {}, "response", "operation")
 
     integration._llmobs_set_tags = mock.Mock(side_effect=AttributeError("Mocked Exception during _llmobs_set_tags()"))
-    integration.llmobs_set_tags(span, kwarg_1="value_1", kwarg_2="value_2")
-    integration._llmobs_set_tags.assert_called_once_with(span, kwarg_1="value_1", kwarg_2="value_2")
+    integration.llmobs_set_tags(
+        span, args=[1, 2, 3], kwargs={"a": 123}, response=[{"content": "hello"}], operation="operation"
+    )
+    integration._llmobs_set_tags.assert_called_once_with(
+        span, [1, 2, 3], {"a": 123}, [{"content": "hello"}], "operation"
+    )
     mock_log.error.assert_called_once_with(
         "Error extracting LLMObs fields for span %s, likely due to malformed data", span, exc_info=True
     )