feat(llmobs): support manual instrumentation of prompts (#7257)

* prompt support for hallucinations + basic tests

* update typedocs

* add more tests

* updates from shared tests

* fmt

* fmt

* change template annotation

* update tests

* add tagging for prompt instrumentation method (annotated)

* trigger ci

Author: sabrenner

docs/test.ts

@@ -681,7 +681,12 @@ llmobs.annotate({
     outputTokens: 5,
     totalTokens: 15
   },
-  tags: {}
+  tags: {},
+  prompt: {
+    id: '123',
+    version: '1.0.0',
+    template: 'this is a {message}',
+  }
 })
 llmobs.annotate(span, {
   inputData: 'input',

index.d.ts

@@ -3202,6 +3202,49 @@ declare namespace tracer {
       type?: string
     }
 
+    /**
+     * A Prompt object that represents the prompt template used for an LLM call.
+     * Used to power LLM Observability prompts and hallucination evaluations.
+     */
+    interface Prompt {
+      /**
+       * Version of the prompt
+       */
+      version?: string,
+
+
+      /**
+       * The id of the prompt set by the user. Should be unique per mlApp.
+       */
+      id?: string,
+
+      /**
+       * An object of string key-value pairs that will be used to render the prompt
+       */
+      variables?: Record<string, string>,
+
+      /**
+       * List of tags to add to the prompt run.
+       */
+      tags?: Record<string, string>,
+
+
+      /**
+       * A list of variable key names that contains query information
+       */
+      queryVariables?: string[],
+
+      /**
+       * A list of variable key names that contain ground truth context information.
+       */
+      contextVariables?: string[],
+
+      /**
+       * A template string or chat message template list.
+       */
+      template?: string | Message[]
+    }
+
     /**
      * Annotation options for LLM Observability spans.
      */
@@ -3235,7 +3278,12 @@ declare namespace tracer {
       /**
        * Object of JSON serializable key-value tag pairs to set or update on the LLM Observability span regarding the span's context.
        */
-      tags?: { [key: string]: any }
+      tags?: { [key: string]: any },
+
+      /**
+       * A Prompt object that represents the prompt used for an LLM call. Only used on `llm` spans.
+       */
+      prompt?: Prompt,
     }
 
     interface AnnotationContextOptions {
@@ -3248,6 +3296,11 @@ declare namespace tracer {
        * Set to override the span name for any spans annotated within the returned context.
        */
       name?: string,
+
+      /**
+       * A Prompt object that represents the prompt used for an LLM call. Only used on `llm` spans.
+       */
+      prompt?: Prompt,
     }
 
     interface RoutingContextOptions {

packages/dd-trace/src/llmobs/constants/tags.js

@@ -17,13 +17,17 @@ module.exports = {
   TRACE_ID: '_ml_obs.trace_id',
   PROPAGATED_TRACE_ID_KEY: '_dd.p.llmobs_trace_id',
   ROOT_PARENT_ID: 'undefined',
+  DEFAULT_PROMPT_NAME: 'unnamed-prompt',
+  INTERNAL_CONTEXT_VARIABLE_KEYS: '_dd_context_variable_keys',
+  INTERNAL_QUERY_VARIABLE_KEYS: '_dd_query_variable_keys',
 
   MODEL_NAME: '_ml_obs.meta.model_name',
   MODEL_PROVIDER: '_ml_obs.meta.model_provider',
 
   INPUT_DOCUMENTS: '_ml_obs.meta.input.documents',
   INPUT_MESSAGES: '_ml_obs.meta.input.messages',
   INPUT_VALUE: '_ml_obs.meta.input.value',
+  INPUT_PROMPT: '_ml_obs.meta.input.prompt',
 
   OUTPUT_DOCUMENTS: '_ml_obs.meta.output.documents',
   OUTPUT_MESSAGES: '_ml_obs.meta.output.messages',

packages/dd-trace/src/llmobs/plugins/openai/index.js

@@ -422,16 +422,15 @@ class OpenAiLLMObsPlugin extends LLMObsPlugin {
     // Handle prompt tracking for reusable prompts
     if (inputs.prompt && response?.prompt) {
       const { id, version } = response.prompt // ResponsePrompt
-      // TODO: Add proper tagger API for prompt metadata
       if (id && version) {
         const normalizedVariables = normalizePromptVariables(inputs.prompt.variables)
         const chatTemplate = extractChatTemplateFromInstructions(response.instructions, normalizedVariables)
-        this._tagger._setTag(span, '_ml_obs.meta.input.prompt', {
+        this._tagger.tagPrompt(span, {
           id,
           version,
           variables: normalizedVariables,
-          chat_template: chatTemplate
-        })
+          template: chatTemplate
+        }, true)
         const tags = { [PROMPT_TRACKING_INSTRUMENTATION_METHOD]: INSTRUMENTATION_METHOD_AUTO }
         if (hasMultimodalInputs(inputs.prompt.variables)) {
           tags[PROMPT_MULTIMODAL] = 'true'

packages/dd-trace/src/llmobs/sdk.js

@@ -241,7 +241,7 @@ class LLMObs extends NoopLLMObs {
         throw new Error('LLMObs span must have a span kind specified')
       }
 
-      const { inputData, outputData, metadata, metrics, tags } = options
+      const { inputData, outputData, metadata, metrics, tags, prompt } = options
 
       if (inputData || outputData) {
         if (spanKind === 'llm') {
@@ -264,6 +264,9 @@ class LLMObs extends NoopLLMObs {
       if (tags) {
         this._tagger.tagSpanTags(span, tags)
       }
+      if (prompt) {
+        this._tagger.tagPrompt(span, prompt)
+      }
     } catch (e) {
       if (e.ddErrorTag) {
         err = e.ddErrorTag

packages/dd-trace/src/llmobs/span_processor.js

@@ -27,6 +27,7 @@ const {
   PARENT_ID_KEY,
   SESSION_ID,
   NAME,
+  INPUT_PROMPT,
   ROUTING_API_KEY,
   ROUTING_SITE
 } = require('./constants/tags')
@@ -130,11 +131,6 @@ class LLMObsSpanProcessor {
       inputType = 'value'
     }
 
-    // Handle prompt metadata for reusable prompts
-    if (mlObsTags['_ml_obs.meta.input.prompt']) {
-      input.prompt = mlObsTags['_ml_obs.meta.input.prompt']
-    }
-
     if (spanKind === 'llm' && mlObsTags[OUTPUT_MESSAGES]) {
       llmObsSpan.output = mlObsTags[OUTPUT_MESSAGES]
       outputType = 'messages'
@@ -185,6 +181,12 @@ class LLMObsSpanProcessor {
     if (input) meta.input = input
     if (output) meta.output = output
 
+    const prompt = mlObsTags[INPUT_PROMPT]
+    if (prompt && spanKind === 'llm') {
+      // by this point, we should have logged a warning if the span kind was not llm
+      meta.input.prompt = prompt
+    }
+
     const llmObsSpanEvent = {
       trace_id: span.context().toTraceId(true),
       span_id: span.context().toSpanId(),

packages/dd-trace/src/llmobs/tagger.js

@@ -29,8 +29,14 @@ const {
   INTEGRATION,
   DECORATOR,
   PROPAGATED_ML_APP_KEY,
+  DEFAULT_PROMPT_NAME,
+  INTERNAL_CONTEXT_VARIABLE_KEYS,
+  INTERNAL_QUERY_VARIABLE_KEYS,
+  INPUT_PROMPT,
   ROUTING_API_KEY,
-  ROUTING_SITE
+  ROUTING_SITE,
+  PROMPT_TRACKING_INSTRUMENTATION_METHOD,
+  INSTRUMENTATION_METHOD_ANNOTATED
 } = require('./constants/tags')
 const { storage } = require('./storage')
 
@@ -113,6 +119,10 @@ class LLMObsTagger {
     const annotationContextName = annotationContext?.name
     if (annotationContextName) this._setTag(span, NAME, annotationContextName)
 
+    // apply annotation context prompt
+    const annotationContextPrompt = annotationContext?.prompt
+    if (annotationContextPrompt) this.tagPrompt(span, annotationContextPrompt)
+
     const routing = storage.getStore()?.routingContext
     if (routing) {
       this._setTag(span, ROUTING_API_KEY, routing.apiKey)
@@ -204,6 +214,160 @@ class LLMObsTagger {
     }
   }
 
+  /**
+   * Tags a prompt on an LLMObs span.
+   * @param {import('../opentracing/span')} span
+   * @param {string | Record<string, unknown>} prompt
+   * @param {boolean?} strictValidation
+   *   whether to validate the prompt against the strict schema, used for auto-instrumentation
+   */
+  tagPrompt (span, prompt, strictValidation = false) {
+    const spanKind = registry.get(span)?.[SPAN_KIND]
+    if (spanKind !== 'llm') {
+      log.warn('Dropping prompt on non-LLM span kind, annotating prompts is only supported for LLM span kinds.')
+      return
+    }
+
+    if (!prompt || typeof prompt !== 'object') {
+      this.#handleFailure('Prompt must be an object.', 'invalid_prompt')
+      return
+    }
+
+    const mlApp = registry.get(span)?.[ML_APP] // this should be defined at this point
+    const {
+      id,
+      version,
+      tags,
+      variables,
+      template,
+      contextVariables,
+      queryVariables,
+    } = prompt
+
+    if (strictValidation) {
+      if (id == null) {
+        this.#handleFailure('Prompt ID is required.', 'invalid_prompt')
+        return
+      }
+
+      if (template == null) {
+        this.#handleFailure('Prompt template is required.', 'invalid_prompt')
+        return
+      }
+    }
+
+    const finalPromptId = id ?? `${mlApp}_${DEFAULT_PROMPT_NAME}`
+    const finalCtxVariablesKeys = contextVariables ?? ['context']
+    const finalQueryVariablesKeys = queryVariables ?? ['question']
+
+    // validate prompt id
+    if (typeof finalPromptId !== 'string') {
+      this.#handleFailure('Prompt ID must be a string.', 'invalid_prompt')
+      return
+    }
+
+    // validate prompt context variables keys
+    if (Array.isArray(finalCtxVariablesKeys)) {
+      for (const key of finalCtxVariablesKeys) {
+        if (typeof key !== 'string') {
+          this.#handleFailure('Prompt context variables keys must be an array of strings.', 'invalid_prompt')
+          return
+        }
+      }
+    } else if (finalCtxVariablesKeys) {
+      this.#handleFailure('Prompt context variables keys must be an array.', 'invalid_prompt')
+      return
+    }
+
+    // validate prompt query variables keys
+    if (Array.isArray(finalQueryVariablesKeys)) {
+      for (const key of finalQueryVariablesKeys) {
+        if (typeof key !== 'string') {
+          this.#handleFailure('Prompt query variables keys must be an array of strings.', 'invalid_prompt')
+          return
+        }
+      }
+    } else if (finalQueryVariablesKeys) {
+      this.#handleFailure('Prompt query variables keys must be an array.', 'invalid_prompt')
+      return
+    }
+
+    // validate prompt version
+    if (version && typeof version !== 'string') {
+      this.#handleFailure('Prompt version must be a string.', 'invalid_prompt')
+      return
+    }
+
+    // validate prompt tags
+    if (tags && (typeof tags !== 'object' || tags instanceof Map)) {
+      this.#handleFailure('Prompt tags must be an non-Map object.', 'invalid_prompt')
+      return
+    } else if (tags) {
+      for (const [key, value] of Object.entries(tags)) {
+        if (typeof key !== 'string' || typeof value !== 'string') {
+          this.#handleFailure('Prompt tags must be an object of string key-value pairs.', 'invalid_prompt')
+          return
+        }
+      }
+    }
+
+    // validate prompt template is either string or list of messages
+    if (template && !(typeof template === 'string' || Array.isArray(template))) {
+      this.#handleFailure('Prompt template must be a string or an array of messages.', 'invalid_prompt')
+      return
+    }
+
+    if (Array.isArray(template)) {
+      for (const message of template) {
+        if (typeof message !== 'object' || !message.role || !message.content) {
+          this.#handleFailure(
+            'Prompt chat template must be an array of objects with role and content properties.', 'invalid_prompt'
+          )
+          return
+        }
+      }
+    }
+
+    // validate variables are a string-string mapping
+    if (variables && (typeof variables !== 'object' || variables instanceof Map)) {
+      this.#handleFailure('Prompt variables must be an non-Map object.', 'invalid_prompt')
+      return
+    } else if (variables) {
+      for (const [key, value] of Object.entries(variables)) {
+        if (typeof key !== 'string' || typeof value !== 'string') {
+          this.#handleFailure('Prompt variables must be an object of string key-value pairs.', 'invalid_prompt')
+          return
+        }
+      }
+    }
+
+    let finalTemplate, finalChatTemplate
+    if (typeof template === 'string') {
+      finalTemplate = template
+    } else if (Array.isArray(template)) {
+      finalChatTemplate = template.map(message => ({ role: message.role, content: message.content }))
+    }
+
+    const validatedPrompt = {}
+    if (finalPromptId) validatedPrompt.id = finalPromptId
+    if (version) validatedPrompt.version = version
+    if (variables) validatedPrompt.variables = variables
+    if (finalTemplate) validatedPrompt.template = finalTemplate
+    if (finalChatTemplate?.length) validatedPrompt.chat_template = finalChatTemplate
+    if (tags) validatedPrompt.tags = tags
+    if (finalCtxVariablesKeys) validatedPrompt[INTERNAL_CONTEXT_VARIABLE_KEYS] = finalCtxVariablesKeys
+    if (finalQueryVariablesKeys) validatedPrompt[INTERNAL_QUERY_VARIABLE_KEYS] = finalQueryVariablesKeys
+
+    const currentPrompt = registry.get(span)?.[INPUT_PROMPT]
+    if (currentPrompt) {
+      Object.assign(currentPrompt, validatedPrompt)
+    } else {
+      this._setTag(span, INPUT_PROMPT, validatedPrompt)
+    }
+
+    this.tagSpanTags(span, { [PROMPT_TRACKING_INSTRUMENTATION_METHOD]: INSTRUMENTATION_METHOD_ANNOTATED })
+  }
+
   changeKind (span, newKind) {
     this._setTag(span, SPAN_KIND, newKind)
   }