feat(llmobs): add prompt and name arguments to annotation context

ddtrace/llmobs/_llmobs.py

@@ -230,16 +230,25 @@ def disable(cls) -> None:
         log.debug("%s disabled", cls.__name__)
 
     @classmethod
-    def annotation_context(cls, tags: Optional[Dict[str, Any]] = None) -> AnnotationContext:
+    def annotation_context(
+        cls, tags: Optional[Dict[str, Any]] = None, prompt: Optional[dict] = None, name: Optional[str] = None
+    ) -> AnnotationContext:
         """
         Sets specified attributes on all LLMObs spans created while the returned AnnotationContext is active.
-        Do not use nested annotation contexts to override the same tags since the order in which annotations
+        Do not use nested annotation contexts to override the same attributes since the order in which annotations
         are applied is non-deterministic.
 
         :param tags: Dictionary of JSON serializable key-value tag pairs to set or update on the LLMObs span
                      regarding the span's context.
+        :param prompt: A dictionary that represents the prompt used for an LLM call in the following form:
+                        `{"template": "...", "id": "...", "version": "...", "variables": {"variable_1": "...", ...}}`.
+                        Can also be set using the `ddtrace.llmobs.utils.Prompt` constructor class.
+                        This argument is only applicable to LLM spans.
+        :param name: Set to override the span name for any spans annotated within the returned context.
         """
-        return AnnotationContext(cls._instance.tracer, lambda span: cls.annotate(span, tags=tags))
+        return AnnotationContext(
+            cls._instance.tracer, lambda span: cls.annotate(span, tags=tags, prompt=prompt, _name=name)
+        )
 
     @classmethod
     def flush(cls) -> None:
@@ -483,6 +492,7 @@ def annotate(
         metadata: Optional[Dict[str, Any]] = None,
         metrics: Optional[Dict[str, Any]] = None,
         tags: Optional[Dict[str, Any]] = None,
+        _name: Optional[str] = None,
     ) -> None:
         """
         Sets parameters, inputs, outputs, tags, and metrics as provided for a given LLMObs span.
@@ -491,7 +501,9 @@ def annotate(
         :param Span span: Span to annotate. If no span is provided, the current active span will be used.
                           Must be an LLMObs-type span, i.e. generated by the LLMObs SDK.
         :param prompt: A dictionary that represents the prompt used for an LLM call in the following form:
-                    {"template": "...", "id": "...", "version": "...", "variables": {"variable_1": "value_1", ...}}.
+                        `{"template": "...", "id": "...", "version": "...", "variables": {"variable_1": "...", ...}}`.
+                        Can also be set using the `ddtrace.llmobs.utils.Prompt` constructor class.
+                        This argument is only applicable to LLM spans.
         :param input_data: A single input string, dictionary, or a list of dictionaries based on the span kind:
                            - llm spans: accepts a string, or a dictionary of form {"content": "...", "role": "..."},
                                         or a list of dictionaries with the same signature.
@@ -534,15 +546,13 @@ def annotate(
         if parameters is not None:
             log.warning("Setting parameters is deprecated, please set parameters and other metadata as tags instead.")
             cls._tag_params(span, parameters)
+        if _name is not None:
+            span.name = _name
+        if prompt is not None:
+            cls._tag_prompt(span, prompt)
         if not span_kind:
             log.debug("Span kind not specified, skipping annotation for input/output data")
             return
-        if prompt is not None:
-            if span_kind == "llm":
-                cls._tag_prompt(span, prompt)
-            else:
-                log.warning("Annotating prompts are only supported for LLM span kinds.")
-
         if input_data or output_data:
             if span_kind == "llm":
                 cls._tag_llm_io(span, input_messages=input_data, output_messages=output_data)

ddtrace/llmobs/_trace_processor.py

@@ -86,8 +86,13 @@ def _llmobs_span_event(self, span: Span) -> Dict[str, Any]:
             meta["output"]["value"] = span._meta.pop(OUTPUT_VALUE)
         if span_kind == "retrieval" and span.get_tag(OUTPUT_DOCUMENTS) is not None:
             meta["output"]["documents"] = json.loads(span._meta.pop(OUTPUT_DOCUMENTS))
-        if span_kind == "llm" and span.get_tag(INPUT_PROMPT) is not None:
-            meta["input"]["prompt"] = json.loads(span._meta.pop(INPUT_PROMPT))
+        if span.get_tag(INPUT_PROMPT) is not None:
+            if span_kind != "llm":
+                log.warning(
+                    "Dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds."
+                )
+            else:
+                meta["input"]["prompt"] = json.loads(span._meta.pop(INPUT_PROMPT))
         if span.error:
             meta[ERROR_MSG] = span.get_tag(ERROR_MSG)
             meta[ERROR_STACK] = span.get_tag(ERROR_STACK)

releasenotes/notes/annotation-context-modify-name-and-prompt-cc74b3b268983181.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    LLM Observability: Introduces `prompt` and `name` arguments to ``LLMObs.annotation_context`` to support setting an integration generated span's name and `prompt` field.

tests/llmobs/test_llmobs_service.py

@@ -798,14 +798,6 @@ def test_annotate_prompt_wrong_type(LLMObs, mock_logs):
         mock_logs.reset_mock()
 
 
-def test_annotate_prompt_wrong_kind(LLMObs, mock_logs):
-    with LLMObs.task(name="dummy") as span:
-        LLMObs.annotate(prompt={"variables": {"var1": "var1"}})
-        assert span.get_tag(INPUT_PROMPT) is None
-        mock_logs.warning.assert_called_once_with("Annotating prompts are only supported for LLM span kinds.")
-        mock_logs.reset_mock()
-
-
 def test_span_error_sets_error(LLMObs, mock_llmobs_span_writer):
     with pytest.raises(ValueError):
         with LLMObs.llm(model_name="test_model", model_provider="test_model_provider") as span:
@@ -1531,13 +1523,39 @@ def test_annotation_context_modifies_span_tags(LLMObs):
             assert json.loads(span.get_tag(TAGS)) == {"foo": "bar"}
 
 
-def test_annotation_context_finished_context_does_not_modify_spans(LLMObs):
+def test_annotation_context_modifies_prompt(LLMObs):
+    with LLMObs.annotation_context(prompt={"template": "test_template"}):
+        with LLMObs.llm(name="test_agent", model_name="test") as span:
+            assert json.loads(span.get_tag(INPUT_PROMPT)) == {"template": "test_template"}
+
+
+def test_annotation_context_modifies_name(LLMObs):
+    with LLMObs.annotation_context(name="test_agent_override"):
+        with LLMObs.llm(name="test_agent", model_name="test") as span:
+            assert span.name == "test_agent_override"
+
+
+def test_annotation_context_finished_context_does_not_modify_tags(LLMObs):
     with LLMObs.annotation_context(tags={"foo": "bar"}):
         pass
     with LLMObs.agent(name="test_agent") as span:
         assert span.get_tag(TAGS) is None
 
 
+def test_annotation_context_finished_context_does_not_modify_prompt(LLMObs):
+    with LLMObs.annotation_context(prompt={"template": "test_template"}):
+        pass
+    with LLMObs.llm(name="test_agent", model_name="test") as span:
+        assert span.get_tag(INPUT_PROMPT) is None
+
+
+def test_annotation_context_finished_context_does_not_modify_name(LLMObs):
+    with LLMObs.annotation_context(name="test_agent_override"):
+        pass
+    with LLMObs.agent(name="test_agent") as span:
+        assert span.name == "test_agent"
+
+
 def test_annotation_context_nested(LLMObs):
     with LLMObs.annotation_context(tags={"foo": "bar", "boo": "bar"}):
         with LLMObs.annotation_context(tags={"car": "car"}):
@@ -1551,13 +1569,39 @@ async def test_annotation_context_async_modifies_span_tags(LLMObs):
             assert json.loads(span.get_tag(TAGS)) == {"foo": "bar"}
 
 
-async def test_annotation_context_async_finished_context_does_not_modify_spans(LLMObs):
+async def test_annotation_context_async_modifies_prompt(LLMObs):
+    async with LLMObs.annotation_context(prompt={"template": "test_template"}):
+        with LLMObs.llm(name="test_agent", model_name="test") as span:
+            assert json.loads(span.get_tag(INPUT_PROMPT)) == {"template": "test_template"}
+
+
+async def test_annotation_context_async_modifies_name(LLMObs):
+    async with LLMObs.annotation_context(name="test_agent_override"):
+        with LLMObs.llm(name="test_agent", model_name="test") as span:
+            assert span.name == "test_agent_override"
+
+
+async def test_annotation_context_async_finished_context_does_not_modify_tags(LLMObs):
     async with LLMObs.annotation_context(tags={"foo": "bar"}):
         pass
     with LLMObs.agent(name="test_agent") as span:
         assert span.get_tag(TAGS) is None
 
 
+async def test_annotation_context_async_finished_context_does_not_modify_prompt(LLMObs):
+    async with LLMObs.annotation_context(prompt={"template": "test_template"}):
+        pass
+    with LLMObs.llm(name="test_agent", model_name="test") as span:
+        assert span.get_tag(INPUT_PROMPT) is None
+
+
+async def test_annotation_context_finished_context_async_does_not_modify_name(LLMObs):
+    async with LLMObs.annotation_context(name="test_agent_override"):
+        pass
+    with LLMObs.agent(name="test_agent") as span:
+        assert span.name == "test_agent"
+
+
 async def test_annotation_context_async_nested(LLMObs):
     async with LLMObs.annotation_context(tags={"foo": "bar", "boo": "bar"}):
         async with LLMObs.annotation_context(tags={"car": "car"}):

tests/llmobs/test_llmobs_trace_processor.py

@@ -341,6 +341,19 @@ def test_prompt_is_set():
         assert tp._llmobs_span_event(llm_span)["meta"]["input"]["prompt"] == {"variables": {"var1": "var2"}}
 
 
+def test_prompt_is_not_set_for_non_llm_spans():
+    """Test that prompt is NOT set on the span event if the span is not an LLM span."""
+    dummy_tracer = DummyTracer()
+    mock_llmobs_span_writer = mock.MagicMock()
+    with override_global_config(dict(_llmobs_ml_app="unnamed-ml-app")):
+        with dummy_tracer.trace("task_span", span_type=SpanTypes.LLM) as task_span:
+            task_span.set_tag(SPAN_KIND, "task")
+            task_span.set_tag(INPUT_VALUE, "ival")
+            task_span.set_tag(INPUT_PROMPT, json.dumps({"variables": {"var1": "var2"}}))
+        tp = LLMObsTraceProcessor(llmobs_span_writer=mock_llmobs_span_writer)
+        assert tp._llmobs_span_event(task_span)["meta"]["input"].get("prompt") is None
+
+
 def test_metadata_is_set():
     """Test that metadata is set on the span event if it is present on the span."""
     dummy_tracer = DummyTracer()