diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py index 5f24d932c5..39b05f30c2 100644 --- a/ddtrace/llmobs/_llmobs.py +++ b/ddtrace/llmobs/_llmobs.py @@ -233,16 +233,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: @@ -486,6 +495,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. @@ -494,7 +504,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. @@ -537,15 +549,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) diff --git a/ddtrace/llmobs/_trace_processor.py b/ddtrace/llmobs/_trace_processor.py index 5a654a8fb9..46d8c98695 100644 --- a/ddtrace/llmobs/_trace_processor.py +++ b/ddtrace/llmobs/_trace_processor.py @@ -86,8 +86,14 @@ 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: + prompt_json_str = span._meta.pop(INPUT_PROMPT) + 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(prompt_json_str) if span.error: meta[ERROR_MSG] = span.get_tag(ERROR_MSG) meta[ERROR_STACK] = span.get_tag(ERROR_STACK) diff --git a/releasenotes/notes/annotation-context-modify-name-and-prompt-cc74b3b268983181.yaml b/releasenotes/notes/annotation-context-modify-name-and-prompt-cc74b3b268983181.yaml new file mode 100644 index 0000000000..ef218a6a80 --- /dev/null +++ b/releasenotes/notes/annotation-context-modify-name-and-prompt-cc74b3b268983181.yaml @@ -0,0 +1,5 @@ +--- +features: + - | + LLM Observability: Introduces `prompt` and `name` arguments to ``LLMObs.annotation_context`` to support setting an integration generated span's name and `prompt` field. + For more information on annotation contexts, see https://docs.datadoghq.com/llm_observability/setup/sdk/#annotating-a-span. \ No newline at end of file diff --git a/tests/llmobs/test_llmobs_service.py b/tests/llmobs/test_llmobs_service.py index 4abe99d4f0..807d609c5f 100644 --- a/tests/llmobs/test_llmobs_service.py +++ b/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: @@ -1573,13 +1565,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"}): @@ -1593,13 +1611,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"}): diff --git a/tests/llmobs/test_llmobs_trace_processor.py b/tests/llmobs/test_llmobs_trace_processor.py index c0a199391d..da1544c5e6 100644 --- a/tests/llmobs/test_llmobs_trace_processor.py +++ b/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()