context var keys

DataDog · Oct 30, 2024 · a2f05ca · a2f05ca
1 parent 60cde3d
commit a2f05ca
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git a/ddtrace/llmobs/_utils.py b/ddtrace/llmobs/_utils.py
@@ -55,6 +55,8 @@ def validate_prompt(prompt: dict) -> Dict[str, Union[str, dict, List[str]]]:
         if not all(isinstance(k, str) for k in ctx_variable_keys):
             raise TypeError("Prompt `context_variable_keys` must be a list of strings.")
         validated_prompt[INTERNAL_CONTEXT_VARIABLE_KEYS] = ctx_variable_keys
+    else:
+        validated_prompt[INTERNAL_CONTEXT_VARIABLE_KEYS] = ["context"]
     return validated_prompt
 
 

diff --git a/tests/llmobs/test_llmobs_service.py b/tests/llmobs/test_llmobs_service.py
@@ -777,6 +777,28 @@ def test_annotate_prompt_dict(LLMObs):
             "variables": {"var1": "var1", "var2": "var3"},
             "version": "1.0.0",
             "id": "test_prompt",
+            "_dd_context_variable_keys": ["context"],
+        }
+
+
+def test_annotate_prompt_dict_with_context_var_keys(LLMObs):
+    with LLMObs.llm(model_name="test_model") as span:
+        LLMObs.annotate(
+            span=span,
+            prompt={
+                "template": "{var1} {var3}",
+                "variables": {"var1": "var1", "var2": "var3"},
+                "version": "1.0.0",
+                "id": "test_prompt",
+                "context_variable_keys": ["var1", "var2"],
+            },
+        )
+        assert json.loads(span.get_tag(INPUT_PROMPT)) == {
+            "template": "{var1} {var3}",
+            "variables": {"var1": "var1", "var2": "var3"},
+            "version": "1.0.0",
+            "id": "test_prompt",
+            "_dd_context_variable_keys": ["var1", "var2"],
         }
 
 
@@ -789,13 +811,15 @@ def test_annotate_prompt_typed_dict(LLMObs):
                 variables={"var1": "var1", "var2": "var3"},
                 version="1.0.0",
                 id="test_prompt",
+                context_variable_keys=["var1", "var2"],
             ),
         )
         assert json.loads(span.get_tag(INPUT_PROMPT)) == {
             "template": "{var1} {var3}",
             "variables": {"var1": "var1", "var2": "var3"},
             "version": "1.0.0",
             "id": "test_prompt",
+            "_dd_context_variable_keys": ["var1", "var2"],
         }
 
 
@@ -1769,7 +1793,10 @@ def test_annotation_context_modifies_span_tags(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"}
+            assert json.loads(span.get_tag(INPUT_PROMPT)) == {
+                "template": "test_template",
+                "_dd_context_variable_keys": ["context"],
+            }
 
 
 def test_annotation_context_modifies_name(LLMObs):
@@ -1908,7 +1935,10 @@ async def test_annotation_context_async_modifies_span_tags(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"}
+            assert json.loads(span.get_tag(INPUT_PROMPT)) == {
+                "template": "test_template",
+                "_dd_context_variable_keys": ["context"],
+            }
 
 
 async def test_annotation_context_async_modifies_name(LLMObs):