yahya/experiments-prompts-tracking

ddtrace/llmobs/_llmobs.py

@@ -1,7 +1,7 @@
 import json
 import os
 import time
-from typing import Any
+from typing import Any, Tuple
 from typing import Dict
 from typing import List
 from typing import Optional
@@ -66,10 +66,10 @@
 from ddtrace.llmobs._utils import _get_span_name
 from ddtrace.llmobs._utils import _is_evaluation_span
 from ddtrace.llmobs._utils import safe_json
-from ddtrace.llmobs._utils import validate_prompt
 from ddtrace.llmobs._writer import LLMObsEvalMetricWriter
 from ddtrace.llmobs._writer import LLMObsSpanWriter
 from ddtrace.llmobs.utils import Documents
+from ddtrace.llmobs.utils import Prompt
 from ddtrace.llmobs.utils import ExportedLLMObsSpan
 from ddtrace.llmobs.utils import Messages
 from ddtrace.propagation.http import HTTPPropagator
@@ -463,7 +463,7 @@ def _tag_span_links(self, span, span_links):
 
     @classmethod
     def annotation_context(
-        cls, tags: Optional[Dict[str, Any]] = None, prompt: Optional[dict] = None, name: Optional[str] = None
+        cls, tags: Optional[Dict[str, Any]] = None, prompt: Optional[Prompt] = None, name: Optional[str] = None
     ) -> AnnotationContext:
         """
         Sets specified attributes on all LLMObs spans created while the returned AnnotationContext is active.
@@ -805,11 +805,30 @@ def retrieval(
             log.warning(SPAN_START_WHILE_DISABLED_WARNING)
         return cls._instance._start_span("retrieval", name=name, session_id=session_id, ml_app=ml_app)
 
+    @classmethod
+    def prompt_context(cls,
+        name: str,
+        version: Optional[str]="1.0.0",
+        template: Optional[List[Tuple[str, str]]]=None,
+        variables: Optional[Dict[str, Any]]=None,
+        example_variable_keys: Optional[List[str]]=None,
+        constraint_variable_keys: Optional[List[str]]=None,
+        rag_context_variable_keys: Optional[List[str]]=None,
+        rag_query_variable_keys: Optional[List[str]]=None,
+        ml_app: str="") -> AnnotationContext:
+        """
+        shortcut to create a prompt object and annotate it
+        """
+        # TODO try to check for if the prompt already exists within the span and update it
+        prompt = Prompt(name, version, template, variables, example_variable_keys, constraint_variable_keys,
+                        rag_context_variable_keys, rag_query_variable_keys, ml_app)
+        return cls.annotation_context(prompt=prompt)
+
     @classmethod
     def annotate(
         cls,
         span: Optional[Span] = None,
-        prompt: Optional[dict] = None,
+        prompt: Optional[Prompt] = None,
         input_data: Optional[Any] = None,
         output_data: Optional[Any] = None,
         metadata: Optional[Dict[str, Any]] = None,
@@ -823,15 +842,8 @@ 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": "...", ...}}`.
-                        Can also be set using the `ddtrace.llmobs.utils.Prompt` constructor class.
-                        - This argument is only applicable to LLM spans.
-                        - The dictionary may contain two optional keys relevant to RAG applications:
-                            `rag_context_variables` - a list of variable key names that contain ground
-                                                        truth context information
-                            `rag_query_variables` - a list of variable key names that contains query
-                                                        information for an LLM call
+        :param prompt: An instance of the `ddtrace.llmobs.utils.Prompt` class that represents the prompt used for an LLM call.
+            - 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.
@@ -883,8 +895,8 @@ def annotate(
             span.name = _name
         if prompt is not None:
             try:
-                validated_prompt = validate_prompt(prompt)
-                cls._set_dict_attribute(span, INPUT_PROMPT, validated_prompt)
+                dict_prompt = prompt.prepare_prompt(ml_app=_get_ml_app(span) or "")
+                cls._set_dict_attribute(span, INPUT_PROMPT, dict_prompt)
             except TypeError:
                 log.warning("Failed to validate prompt with error: ", exc_info=True)
         if not span_kind:

ddtrace/llmobs/_utils.py

@@ -20,59 +20,12 @@
 from ddtrace.llmobs._constants import OPENAI_APM_SPAN_NAME
 from ddtrace.llmobs._constants import SESSION_ID
 from ddtrace.llmobs._constants import VERTEXAI_APM_SPAN_NAME
+from ddtrace.llmobs.utils import Prompt
 from ddtrace.trace import Span
 
 
 log = get_logger(__name__)
 
-
-def validate_prompt(prompt: dict) -> Dict[str, Union[str, dict, List[str]]]:
-    validated_prompt = {}  # type: Dict[str, Union[str, dict, List[str]]]
-    if not isinstance(prompt, dict):
-        raise TypeError("Prompt must be a dictionary")
-    variables = prompt.get("variables")
-    template = prompt.get("template")
-    version = prompt.get("version")
-    prompt_id = prompt.get("id")
-    ctx_variable_keys = prompt.get("rag_context_variables")
-    rag_query_variable_keys = prompt.get("rag_query_variables")
-    if variables is not None:
-        if not isinstance(variables, dict):
-            raise TypeError("Prompt variables must be a dictionary.")
-        if not any(isinstance(k, str) or isinstance(v, str) for k, v in variables.items()):
-            raise TypeError("Prompt variable keys and values must be strings.")
-        validated_prompt["variables"] = variables
-    if template is not None:
-        if not isinstance(template, str):
-            raise TypeError("Prompt template must be a string")
-        validated_prompt["template"] = template
-    if version is not None:
-        if not isinstance(version, str):
-            raise TypeError("Prompt version must be a string.")
-        validated_prompt["version"] = version
-    if prompt_id is not None:
-        if not isinstance(prompt_id, str):
-            raise TypeError("Prompt id must be a string.")
-        validated_prompt["id"] = prompt_id
-    if ctx_variable_keys is not None:
-        if not isinstance(ctx_variable_keys, list):
-            raise TypeError("Prompt field `context_variable_keys` must be a list of strings.")
-        if not all(isinstance(k, str) for k in ctx_variable_keys):
-            raise TypeError("Prompt field `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"]
-    if rag_query_variable_keys is not None:
-        if not isinstance(rag_query_variable_keys, list):
-            raise TypeError("Prompt field `rag_query_variables` must be a list of strings.")
-        if not all(isinstance(k, str) for k in rag_query_variable_keys):
-            raise TypeError("Prompt field `rag_query_variables` must be a list of strings.")
-        validated_prompt[INTERNAL_QUERY_VARIABLE_KEYS] = rag_query_variable_keys
-    else:
-        validated_prompt[INTERNAL_QUERY_VARIABLE_KEYS] = ["question"]
-    return validated_prompt
-
-
 class LinkTracker:
     def __init__(self, object_span_links=None):
         self._object_span_links = object_span_links or {}
@@ -185,7 +138,7 @@ def _get_session_id(span: Span) -> Optional[str]:
 def _inject_llmobs_parent_id(span_context):
     """Inject the LLMObs parent ID into the span context for reconnecting distributed LLMObs traces."""
     span = ddtrace.tracer.current_span()
-    
+
     if span is None:
         log.warning("No active span to inject LLMObs parent ID info.")
         return