From 7b310512ffcce7889054873efe6938454b68d5cf Mon Sep 17 00:00:00 2001 From: Rodrigo Arguello Date: Thu, 26 Mar 2026 12:32:12 +0100 Subject: [PATCH 1/3] feat(llmobs): prompt tracking --- contrib/mark3labs/mcp-go/intent_capture.go | 2 +- .../go-sdk/intent_capture.go | 2 +- internal/llmobs/llmobs.go | 5 +- internal/llmobs/span.go | 19 ++++++-- llmobs/llmobs.go | 25 ++++++++-- llmobs/llmobs_test.go | 47 +++++++++++++++++++ llmobs/option.go | 33 +++++++++++++ 7 files changed, 121 insertions(+), 12 deletions(-) diff --git a/contrib/mark3labs/mcp-go/intent_capture.go b/contrib/mark3labs/mcp-go/intent_capture.go index 35d02352f15..f4837e48e44 100644 --- a/contrib/mark3labs/mcp-go/intent_capture.go +++ b/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)) } diff --git a/contrib/modelcontextprotocol/go-sdk/intent_capture.go b/contrib/modelcontextprotocol/go-sdk/intent_capture.go index 5780be24f4d..2251c43898f 100644 --- a/contrib/modelcontextprotocol/go-sdk/intent_capture.go +++ b/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)) } diff --git a/internal/llmobs/llmobs.go b/internal/llmobs/llmobs.go index 3b382ff75e0..46acd28ea72 100644 --- a/internal/llmobs/llmobs.go +++ b/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) } diff --git a/internal/llmobs/span.go b/internal/llmobs/span.go index 54503fa576a..c65fa2dbb2d 100644 --- a/internal/llmobs/span.go +++ b/internal/llmobs/span.go @@ -82,16 +82,25 @@ type EvaluationConfig struct { type Prompt struct { // Template is the prompt template string. Template string `json:"template,omitempty"` - // ID is the unique identifier for the prompt. + // ChatTemplate is a list of messages forming the prompt (alternative to Template). + ChatTemplate []LLMMessage `json:"chat_template,omitempty"` + // 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"` // 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 } } diff --git a/llmobs/llmobs.go b/llmobs/llmobs.go index 1422f7bf372..7e1e78bff5a 100644 --- a/llmobs/llmobs.go +++ b/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 { + // Template is the prompt template string. + Template string + // ChatTemplate is a list of messages forming the prompt (alternative to Template). + ChatTemplate []LLMMessage + // 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 + // 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 diff --git a/llmobs/llmobs_test.go b/llmobs/llmobs_test.go index ab3c5ba4ad2..00f8b42647d 100644 --- a/llmobs/llmobs_test.go +++ b/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) { diff --git a/llmobs/option.go b/llmobs/option.go index c373a167069..56205e16621 100644 --- a/llmobs/option.go +++ b/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 + } +} From 7a2fe8ea6e6ff121a9773d4438f85790ef62596d Mon Sep 17 00:00:00 2001 From: Rodrigo Arguello Date: Thu, 26 Mar 2026 16:29:45 +0100 Subject: [PATCH 2/3] add tests --- internal/llmobs/llmobs_test.go | 84 ++++++++++++++++++++++++++++++++++ internal/llmobs/span.go | 8 ++-- llmobs/llmobs.go | 8 ++-- 3 files changed, 92 insertions(+), 8 deletions(-) diff --git a/internal/llmobs/llmobs_test.go b/internal/llmobs/llmobs_test.go index dac9993f6ba..efb0a8b6a3d 100644 --- a/internal/llmobs/llmobs_test.go +++ b/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 { diff --git a/internal/llmobs/span.go b/internal/llmobs/span.go index c65fa2dbb2d..3a9b7a8b79e 100644 --- a/internal/llmobs/span.go +++ b/internal/llmobs/span.go @@ -80,16 +80,16 @@ 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"` - // ChatTemplate is a list of messages forming the prompt (alternative to Template). - ChatTemplate []LLMMessage `json:"chat_template,omitempty"` // 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. diff --git a/llmobs/llmobs.go b/llmobs/llmobs.go index 7e1e78bff5a..0635d966a9f 100644 --- a/llmobs/llmobs.go +++ b/llmobs/llmobs.go @@ -152,16 +152,16 @@ type ( // Prompt represents a structured prompt template used with LLM spans. type Prompt struct { - // Template is the prompt template string. - Template string - // ChatTemplate is a list of messages forming the prompt (alternative to Template). - ChatTemplate []LLMMessage // 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. From 0a66bd37b0143e99764c199afc03954298b25877 Mon Sep 17 00:00:00 2001 From: Rodrigo Arguello Date: Thu, 26 Mar 2026 17:29:26 +0100 Subject: [PATCH 3/3] fix fmt --- internal/llmobs/llmobs_test.go | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/internal/llmobs/llmobs_test.go b/internal/llmobs/llmobs_test.go index efb0a8b6a3d..acfd095fbac 100644 --- a/internal/llmobs/llmobs_test.go +++ b/internal/llmobs/llmobs_test.go @@ -671,15 +671,15 @@ func TestSpanAnnotate(t *testing.T) { "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"}, + "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, + "ml_app": mlApp, }, }, }, @@ -696,11 +696,11 @@ func TestSpanAnnotate(t *testing.T) { "span.kind": "llm", "input": map[string]any{ "prompt": map[string]any{ - "id": mlApp + "_unnamed-prompt", - "template": "Answer the question.", + "id": mlApp + "_unnamed-prompt", + "template": "Answer the question.", "_dd_context_variable_keys": []any{"context"}, "_dd_query_variable_keys": []any{"question"}, - "ml_app": mlApp, + "ml_app": mlApp, }, }, }, @@ -715,8 +715,8 @@ func TestSpanAnnotate(t *testing.T) { {Role: "system", Content: "You are a helpful assistant."}, {Role: "user", Content: "{{question}}"}, }, - Variables: map[string]string{"question": "What is Go?"}, - RAGQueryVariables: []string{"question"}, + Variables: map[string]string{"question": "What is Go?"}, + RAGQueryVariables: []string{"question"}, }, }, wantMeta: map[string]any{ @@ -728,10 +728,10 @@ func TestSpanAnnotate(t *testing.T) { 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?"}, + "variables": map[string]any{"question": "What is Go?"}, "_dd_context_variable_keys": []any{"context"}, "_dd_query_variable_keys": []any{"question"}, - "ml_app": mlApp, + "ml_app": mlApp, }, }, },