feat(llmobs): prompt tracking

contrib/mark3labs/mcp-go/intent_capture.go

@@ -111,5 +111,5 @@ func processTelemetry(ctx context.Context, telemetryVal map[string]any) {
 	if !ok {
 		return
 	}
-	toolSpan.Annotate(llmobs.WithIntent(intent))
+	toolSpan.Annotate(llmobs.WithAnnotatedIntent(intent))
 }

contrib/modelcontextprotocol/go-sdk/intent_capture.go

@@ -141,5 +141,5 @@ func annotateSpanWithIntent(ctx context.Context, telemetryVal map[string]any) {
 	if !ok {
 		return
 	}
-	toolSpan.Annotate(llmobs.WithIntent(intent))
+	toolSpan.Annotate(llmobs.WithAnnotatedIntent(intent))
 }

internal/llmobs/llmobs.go

@@ -501,10 +501,6 @@ func (l *LLMObs) llmobsSpanEvent(span *Span) *transport.LLMObsSpanEvent {
 		} else {
 			input["prompt"] = inputPrompt
 		}
-	} else if spanKind == SpanKindLLM {
-		if span.parent != nil && span.parent.llmCtx.prompt != nil {
-			input["prompt"] = span.parent.llmCtx.prompt
-		}
 	}
 
 	if toolDefinitions := span.llmCtx.toolDefinitions; len(toolDefinitions) > 0 {
@@ -720,6 +716,7 @@ func (l *LLMObs) StartSpan(ctx context.Context, kind SpanKind, name string, cfg
 			log.Warn("llmobs: ML App is required for sending LLM Observability data.")
 		}
 	}
+
 	log.Debug("llmobs: starting LLMObs span: %s, span_kind: %s, ml_app: %s", spanName, kind, span.mlApp)
 	return span, contextWithActiveLLMSpan(ctx, span)
 }

internal/llmobs/llmobs_test.go

@@ -652,6 +652,90 @@ func TestSpanAnnotate(t *testing.T) {
 				"intent":    "test intent",
 			},
 		},
+		{
+			name: "llm-span-with-full-prompt",
+			kind: llmobs.SpanKindLLM,
+			annotations: llmobs.SpanAnnotations{
+				Prompt: &llmobs.Prompt{
+					ID:                  "my-prompt",
+					Version:             "1.0.0",
+					Label:               "production",
+					Template:            "Hello {{name}}!",
+					Variables:           map[string]string{"name": "World"},
+					Tags:                map[string]string{"env": "prod"},
+					RAGContextVariables: []string{"context"},
+					RAGQueryVariables:   []string{"query"},
+				},
+			},
+			wantMeta: map[string]any{
+				"span.kind": "llm",
+				"input": map[string]any{
+					"prompt": map[string]any{
+						"id":                        "my-prompt",
+						"version":                   "1.0.0",
+						"label":                     "production",
+						"template":                  "Hello {{name}}!",
+						"variables":                 map[string]any{"name": "World"},
+						"tags":                      map[string]any{"env": "prod"},
+						"_dd_context_variable_keys": []any{"context"},
+						"_dd_query_variable_keys":   []any{"query"},
+						"ml_app":                    mlApp,
+					},
+				},
+			},
+		},
+		{
+			name: "llm-span-with-prompt-defaults-applied",
+			kind: llmobs.SpanKindLLM,
+			annotations: llmobs.SpanAnnotations{
+				Prompt: &llmobs.Prompt{
+					Template: "Answer the question.",
+				},
+			},
+			wantMeta: map[string]any{
+				"span.kind": "llm",
+				"input": map[string]any{
+					"prompt": map[string]any{
+						"id":                        mlApp + "_unnamed-prompt",
+						"template":                  "Answer the question.",
+						"_dd_context_variable_keys": []any{"context"},
+						"_dd_query_variable_keys":   []any{"question"},
+						"ml_app":                    mlApp,
+					},
+				},
+			},
+		},
+		{
+			name: "llm-span-with-prompt-chat-template",
+			kind: llmobs.SpanKindLLM,
+			annotations: llmobs.SpanAnnotations{
+				Prompt: &llmobs.Prompt{
+					ID: "chat-prompt",
+					ChatTemplate: []llmobs.LLMMessage{
+						{Role: "system", Content: "You are a helpful assistant."},
+						{Role: "user", Content: "{{question}}"},
+					},
+					Variables:         map[string]string{"question": "What is Go?"},
+					RAGQueryVariables: []string{"question"},
+				},
+			},
+			wantMeta: map[string]any{
+				"span.kind": "llm",
+				"input": map[string]any{
+					"prompt": map[string]any{
+						"id": "chat-prompt",
+						"chat_template": []any{
+							map[string]any{"role": "system", "content": "You are a helpful assistant."},
+							map[string]any{"role": "user", "content": "{{question}}"},
+						},
+						"variables":                 map[string]any{"question": "What is Go?"},
+						"_dd_context_variable_keys": []any{"context"},
+						"_dd_query_variable_keys":   []any{"question"},
+						"ml_app":                    mlApp,
+					},
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {

internal/llmobs/span.go

@@ -80,18 +80,27 @@ type EvaluationConfig struct {
 
 // Prompt represents a prompt template used with LLM spans.
 type Prompt struct {
-	// Template is the prompt template string.
-	Template string `json:"template,omitempty"`
-	// ID is the unique identifier for the prompt.
+	// ID is the unique identifier for the prompt within the ML app.
 	ID string `json:"id,omitempty"`
 	// Version is the version of the prompt.
 	Version string `json:"version,omitempty"`
+	// Label is the deployment label (e.g., "production", "staging").
+	Label string `json:"label,omitempty"`
+	// Template is the prompt template string.
+	Template string `json:"template,omitempty"`
+	// ChatTemplate is a list of messages forming the prompt (alternative to Template).
+	ChatTemplate []LLMMessage `json:"chat_template,omitempty"`
 	// Variables contains the variables used in the prompt template.
 	Variables map[string]string `json:"variables,omitempty"`
+	// Tags contains custom tags for the prompt.
+	Tags map[string]string `json:"tags,omitempty"`
 	// RAGContextVariables specifies which variables contain RAG context.
-	RAGContextVariables []string `json:"rag_context_variables,omitempty"`
+	RAGContextVariables []string `json:"_dd_context_variable_keys,omitempty"`
 	// RAGQueryVariables specifies which variables contain RAG queries.
-	RAGQueryVariables []string `json:"rag_query_variables,omitempty"`
+	RAGQueryVariables []string `json:"_dd_query_variable_keys,omitempty"`
+
+	// MLApp is the ML application name, set internally from the span's context.
+	MLApp string `json:"ml_app,omitempty"`
 }
 
 // ToolDefinition represents a tool definition for LLM spans.
@@ -352,6 +361,10 @@ func (s *Span) Annotate(a SpanAnnotations) {
 			if a.Prompt.RAGQueryVariables == nil {
 				a.Prompt.RAGQueryVariables = []string{"question"}
 			}
+			if a.Prompt.ID == "" {
+				a.Prompt.ID = s.mlApp + "_unnamed-prompt"
+			}
+			a.Prompt.MLApp = s.mlApp
 			s.llmCtx.prompt = a.Prompt
 		}
 	}

llmobs/llmobs.go

@@ -140,9 +140,6 @@ type (
 	// It is used to annotate output of retrieval spans.
 	RetrievedDocument = illmobs.RetrievedDocument
 
-	// Prompt represents a structured prompt template used with LLMs.
-	Prompt = illmobs.Prompt
-
 	// ToolDefinition represents the definition of a tool/function that an LLM can call.
 	ToolDefinition = illmobs.ToolDefinition
 
@@ -153,6 +150,28 @@ type (
 	ToolResult = illmobs.ToolResult
 )
 
+// Prompt represents a structured prompt template used with LLM spans.
+type Prompt struct {
+	// ID is the unique identifier for the prompt within the ML app.
+	ID string
+	// Version is the version of the prompt.
+	Version string
+	// Label is the deployment label (e.g., "production", "staging").
+	Label string
+	// Template is the prompt template string.
+	Template string
+	// ChatTemplate is a list of messages forming the prompt (alternative to Template).
+	ChatTemplate []LLMMessage
+	// Variables contains the variables used in the prompt template.
+	Variables map[string]string
+	// Tags contains custom tags for the prompt.
+	Tags map[string]string
+	// RAGContextVariables specifies which variables contain RAG context.
+	RAGContextVariables []string
+	// RAGQueryVariables specifies which variables contain RAG queries.
+	RAGQueryVariables []string
+}
+
 // Span represents a generic LLMObs span that can be converted to specific span types.
 type Span interface {
 	sealed() // Prevents external implementations

llmobs/llmobs_test.go

@@ -511,6 +511,53 @@ func TestSpanAnnotations(t *testing.T) {
 		assert.Contains(t, spans[0].Meta, "input")
 		assert.Contains(t, spans[0].Meta, "output")
 	})
+	t.Run("with-annotated-prompt", func(t *testing.T) {
+		tt := testTracer(t)
+		defer tt.Stop()
+
+		span, _ := llmobs.StartLLMSpan(ctx, "test-llm-prompt")
+		span.AnnotateLLMIO(nil, nil,
+			llmobs.WithAnnotatedPrompt(llmobs.Prompt{
+				Template:  "Answer the question: {{question}}",
+				Variables: map[string]string{"question": "What is 2+2?"},
+			}),
+		)
+		span.Finish()
+
+		spans := tt.WaitForLLMObsSpans(t, 1)
+		require.Len(t, spans, 1)
+		inputMeta, ok := spans[0].Meta["input"].(map[string]any)
+		require.True(t, ok)
+		assert.NotNil(t, inputMeta["prompt"])
+	})
+	t.Run("with-annotated-intent", func(t *testing.T) {
+		tt := testTracer(t)
+		defer tt.Stop()
+
+		span, _ := llmobs.StartToolSpan(ctx, "test-tool-intent")
+		span.Annotate(llmobs.WithAnnotatedIntent("fetch weather data for the user's location"))
+		span.Finish()
+
+		spans := tt.WaitForLLMObsSpans(t, 1)
+		require.Len(t, spans, 1)
+		assert.Equal(t, "fetch weather data for the user's location", spans[0].Meta["intent"])
+	})
+	t.Run("with-annotated-tool-definitions", func(t *testing.T) {
+		tt := testTracer(t)
+		defer tt.Stop()
+
+		span, _ := llmobs.StartLLMSpan(ctx, "test-llm-tool-definitions")
+		span.AnnotateLLMIO(nil, nil,
+			llmobs.WithAnnotatedToolDefinitions([]llmobs.ToolDefinition{
+				{Name: "get_weather", Description: "Get the current weather for a location"},
+			}),
+		)
+		span.Finish()
+
+		spans := tt.WaitForLLMObsSpans(t, 1)
+		require.Len(t, spans, 1)
+		assert.NotNil(t, spans[0].Meta["tool_definitions"])
+	})
 }
 
 func TestEvaluationMetrics(t *testing.T) {

llmobs/option.go

@@ -152,8 +152,41 @@ func WithAnnotatedMetrics(metrics map[string]float64) AnnotateOption {
 
 // WithIntent sets the intent for the span.
 // Intent is a description of a reason for calling an MCP tool.
+// Deprecated: Use WithAnnotatedIntent instead.
 func WithIntent(intent string) AnnotateOption {
+	return WithAnnotatedIntent(intent)
+}
+
+// WithAnnotatedIntent sets the intent for the span.
+// Intent is a description of a reason for calling an MCP tool.
+func WithAnnotatedIntent(intent string) AnnotateOption {
 	return func(a *illmobs.SpanAnnotations) {
 		a.Intent = intent
 	}
 }
+
+// WithAnnotatedPrompt sets the prompt for the span annotation.
+// Only applicable to LLM spans.
+func WithAnnotatedPrompt(prompt Prompt) AnnotateOption {
+	return func(a *illmobs.SpanAnnotations) {
+		a.Prompt = &illmobs.Prompt{
+			Template:            prompt.Template,
+			ChatTemplate:        prompt.ChatTemplate,
+			ID:                  prompt.ID,
+			Version:             prompt.Version,
+			Label:               prompt.Label,
+			Variables:           prompt.Variables,
+			Tags:                prompt.Tags,
+			RAGContextVariables: prompt.RAGContextVariables,
+			RAGQueryVariables:   prompt.RAGQueryVariables,
+		}
+	}
+}
+
+// WithAnnotatedToolDefinitions sets the tool definitions for the span annotation.
+// Only applicable to LLM spans.
+func WithAnnotatedToolDefinitions(toolDefinitions []ToolDefinition) AnnotateOption {
+	return func(a *illmobs.SpanAnnotations) {
+		a.ToolDefinitions = toolDefinitions
+	}
+}