{agent, telemetry}: add output for invoke_agent span (#419)

Author: liuzengh

agent/llmagent/extras_test.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/trace/noop"
 
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/event"
@@ -133,7 +134,7 @@ func TestLLMAgent_AfterCb(t *testing.T) {
 	inv := &agent.Invocation{InvocationID: "id", AgentName: "agent"}
 
 	llm := &LLMAgent{agentCallbacks: cb}
-	wrapped := llm.wrapEventChannel(context.Background(), inv, orig)
+	wrapped := llm.wrapEventChannel(context.Background(), inv, orig, noop.Span{})
 
 	var objs []string
 	for e := range wrapped {
@@ -157,7 +158,7 @@ func TestLLMAgent_AfterCbNoResp(t *testing.T) {
 	inv := &agent.Invocation{InvocationID: "id2", AgentName: "agent2"}
 
 	llm := &LLMAgent{}
-	wrapped := llm.wrapEventChannel(context.Background(), inv, orig)
+	wrapped := llm.wrapEventChannel(context.Background(), inv, orig, noop.Span{})
 
 	// Expect exactly one event propagated from original channel and no extras.
 	count := 0

agent/llmagent/llm_agent.go

@@ -16,6 +16,8 @@ import (
 	"reflect"
 	"sync"
 
+	sdktrace "go.opentelemetry.io/otel/trace"
+
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/agent/llmagent/internal/jsonschema"
 	"trpc.group/trpc-go/trpc-agent-go/codeexecutor"
@@ -562,16 +564,22 @@ func registerTools(options *Options) []tool.Tool {
 
 // Run implements the agent.Agent interface.
 // It executes the LLM agent flow and returns a channel of events.
-func (a *LLMAgent) Run(ctx context.Context, invocation *agent.Invocation) (<-chan *event.Event, error) {
-	ctx, span := trace.Tracer.Start(ctx, fmt.Sprintf("%s [%s]", itelemetry.SpanNamePrefixAgentRun, a.name))
-	defer span.End()
+func (a *LLMAgent) Run(ctx context.Context, invocation *agent.Invocation) (e <-chan *event.Event, err error) {
+	ctx, span := trace.Tracer.Start(ctx, fmt.Sprintf("invoke_agent {%s}", a.name))
+	itelemetry.TraceBeforeInvokeAgent(span, invocation)
+	defer func() {
+		if err != nil {
+			span.End()
+		}
+	}()
 
 	// Setup invocation
 	a.setupInvocation(invocation)
 
 	// Run before agent callbacks if they exist.
 	if a.agentCallbacks != nil {
-		customResponse, err := a.agentCallbacks.RunBeforeAgent(ctx, invocation)
+		var customResponse *model.Response
+		customResponse, err = a.agentCallbacks.RunBeforeAgent(ctx, invocation)
 		if err != nil {
 			return nil, fmt.Errorf("before agent callback failed: %w", err)
 		}
@@ -591,13 +599,7 @@ func (a *LLMAgent) Run(ctx context.Context, invocation *agent.Invocation) (<-cha
 	if err != nil {
 		return nil, err
 	}
-
-	// If we have after agent callbacks, we need to wrap the event channel.
-	if a.agentCallbacks != nil {
-		return a.wrapEventChannel(ctx, invocation, flowEventChan), nil
-	}
-
-	return flowEventChan, nil
+	return a.wrapEventChannel(ctx, invocation, flowEventChan, span), nil
 }
 
 // setupInvocation sets up the invocation
@@ -621,14 +623,25 @@ func (a *LLMAgent) wrapEventChannel(
 	ctx context.Context,
 	invocation *agent.Invocation,
 	originalChan <-chan *event.Event,
+	span sdktrace.Span,
 ) <-chan *event.Event {
 	wrappedChan := make(chan *event.Event, 256) // Use default buffer size
 
 	go func() {
-		defer close(wrappedChan)
+		var fullResponse *model.Response
+		defer func() {
+			if fullResponse != nil {
+				itelemetry.TraceAfterInvokeAgent(span, fullResponse)
+			}
+			span.End()
+			close(wrappedChan)
+		}()
 
 		// Forward all events from the original channel
 		for evt := range originalChan {
+			if evt != nil && evt.Response != nil && !evt.Response.IsPartial {
+				fullResponse = evt.Response
+			}
 			if err := event.EmitEvent(ctx, wrappedChan, evt); err != nil {
 				return
 			}
@@ -647,6 +660,7 @@ func (a *LLMAgent) wrapEventChannel(
 					err.Error(),
 				)
 			} else if customResponse != nil {
+				fullResponse = customResponse
 				// Create an event from the custom response.
 				evt = event.NewResponseEvent(invocation.InvocationID, invocation.AgentName, customResponse)
 			}

internal/telemetry/telemetry.go

@@ -35,12 +35,10 @@ const (
 
 	SpanNameCallLLM           = "call_llm"
 	SpanNamePrefixExecuteTool = "execute_tool"
-	SpanNamePrefixAgentRun    = "agent_run"
-	SpanNameInvocation        = "invocation"
 
 	OperationExecuteTool = "execute_tool"
 	OperationCallLLM     = "call_llm"
-	OperationRunRunner   = "run_runner" // attribute of SpanNameInvocation
+	OperationInvokeAgent = "invoke_agent"
 )
 
 const (
@@ -75,11 +73,19 @@ var (
 	KeyToolID       = "trpc.go.agent.tool_id"
 
 	// GenAI operation attributes
-	KeyGenAIOperationName = "gen_ai.operation.name"
-	KeyGenAISystem        = "gen_ai.system"
-	KeyGenAIToolName      = "gen_ai.tool.name"
-	KeyGenAIToolDesc      = "gen_ai.tool.description"
-	KeyGenAIRequestModel  = "gen_ai.request.model"
+	KeyGenAIOperationName     = "gen_ai.operation.name"
+	KeyGenAISystem            = "gen_ai.system"
+	KeyGenAIToolName          = "gen_ai.tool.name"
+	KeyGenAIToolDesc          = "gen_ai.tool.description"
+	KeyGenAIRequestModel      = "gen_ai.request.model"
+	KeyGenAIInputMessages     = "gen_ai.input.messages"
+	KeyGenAIOutputMessages    = "gen_ai.output.messages"
+	KeyGenAIAgentName         = "gen_ai.agent.name"
+	KeyGenAIConversationID    = "gen_ai.conversation.id"
+	KeyGenAIResponseModel     = "gen_ai.response.model"
+	KeyGenAIUsageOutputTokens = "gen_ai.usage.output_tokens"
+	KeyGenAIResponseID        = "gen_ai.response.id"
+	KeyGenAIUsageInputTokens  = "gen_ai.usage.input_tokens"
 
 	// System value
 	SystemTRPCGoAgent = "trpc.go.agent"
@@ -142,24 +148,45 @@ func TraceMergedToolCalls(span trace.Span, rspEvent *event.Event) {
 	)
 }
 
-// TraceRunner traces the invocation of a runner.
-func TraceRunner(span trace.Span, appName string, invoke *agent.Invocation, message model.Message) {
-	if bts, err := json.Marshal(&model.Request{Messages: []model.Message{message}}); err == nil {
+// TraceBeforeInvokeAgent traces the before invocation of an agent.
+func TraceBeforeInvokeAgent(span trace.Span, invoke *agent.Invocation) {
+	if bts, err := json.Marshal(&model.Request{Messages: []model.Message{invoke.Message}}); err == nil {
 		span.SetAttributes(
-			attribute.String(KeyRunnerInput, string(bts)),
+			attribute.String(KeyGenAIInputMessages, string(bts)),
 		)
 	} else {
-		span.SetAttributes(attribute.String(KeyRunnerInput, "<not json serializable>"))
+		span.SetAttributes(attribute.String(KeyGenAIInputMessages, "<not json serializable>"))
 	}
-
 	span.SetAttributes(
 		attribute.String(KeyGenAISystem, SystemTRPCGoAgent),
-		attribute.String(KeyGenAIOperationName, OperationRunRunner),
-		attribute.String(KeyRunnerName, fmt.Sprintf("[trpc-go-agent]: %s/%s", appName, invoke.AgentName)),
+		attribute.String(KeyGenAIOperationName, OperationInvokeAgent),
+		attribute.String(KeyGenAIAgentName, invoke.AgentName),
 		attribute.String(KeyInvocationID, invoke.InvocationID),
-		attribute.String(KeyRunnerSessionID, invoke.Session.ID),
-		attribute.String(KeyRunnerUserID, invoke.Session.UserID),
 	)
+	if invoke.Session != nil {
+		span.SetAttributes(
+			attribute.String(KeyRunnerUserID, invoke.Session.UserID),
+			attribute.String(KeyGenAIConversationID, invoke.Session.ID),
+		)
+	}
+}
+
+// TraceAfterInvokeAgent traces the after invocation of an agent.
+func TraceAfterInvokeAgent(span trace.Span, rsp *model.Response) {
+	if len(rsp.Choices) > 0 {
+		if bts, err := json.Marshal(rsp.Choices[0].Message); err == nil {
+			span.SetAttributes(
+				attribute.String(KeyGenAIOutputMessages, string(bts)),
+			)
+		}
+	}
+	span.SetAttributes(attribute.String(KeyGenAIResponseModel, rsp.Model))
+	if rsp.Usage != nil {
+		span.SetAttributes(attribute.Int(KeyGenAIUsageInputTokens, rsp.Usage.PromptTokens))
+		span.SetAttributes(attribute.Int(KeyGenAIUsageOutputTokens, rsp.Usage.CompletionTokens))
+	}
+	span.SetAttributes(attribute.String(KeyGenAIResponseID, rsp.ID))
+
 }
 
 // TraceCallLLM traces the invocation of an LLM call.

knowledge/vectorstore/inmemory/inmemory_test.go

@@ -201,7 +201,7 @@ func TestVectorStore_DeleteByFilter(t *testing.T) {
 	// Test delete all
 	require.NoError(t, store.Add(ctx, doc1, embedding))
 	require.NoError(t, store.Add(ctx, doc2, embedding))
-	
+
 	err = store.DeleteByFilter(ctx, vectorstore.WithDeleteAll(true))
 	require.NoError(t, err)
 
@@ -313,9 +313,9 @@ func TestVectorStore_GetMetadata(t *testing.T) {
 	metadata, err = store.GetMetadata(ctx, vectorstore.WithGetMetadataLimit(1), vectorstore.WithGetMetadataOffset(1))
 	require.NoError(t, err)
 	require.Equal(t, 1, len(metadata))
-	
+
 	// Get metadata with combined filters
-	metadata, err = store.GetMetadata(ctx, 
+	metadata, err = store.GetMetadata(ctx,
 		vectorstore.WithGetMetadataFilter(map[string]any{"type": "test"}),
 		vectorstore.WithGetMetadataLimit(10),
 		vectorstore.WithGetMetadataOffset(0))
@@ -344,7 +344,7 @@ func TestVectorStore_SearchModes(t *testing.T) {
 		},
 	}
 	doc2 := &document.Document{
-		ID:      "doc2", 
+		ID:      "doc2",
 		Content: "bonjour monde",
 		Metadata: map[string]any{
 			"lang": "fr",
@@ -358,10 +358,10 @@ func TestVectorStore_SearchModes(t *testing.T) {
 
 	// Test vector search mode
 	result, err := store.Search(ctx, &vectorstore.SearchQuery{
-		Vector:      []float64{0.9, 0.1, 0.2},
-		Limit:       5,
-		SearchMode:  vectorstore.SearchModeVector,
-		MinScore:    0.5,
+		Vector:     []float64{0.9, 0.1, 0.2},
+		Limit:      5,
+		SearchMode: vectorstore.SearchModeVector,
+		MinScore:   0.5,
 	})
 	require.NoError(t, err)
 	require.Len(t, result.Results, 1)
@@ -383,10 +383,10 @@ func TestVectorStore_SearchModes(t *testing.T) {
 
 	// Test hybrid search mode (falls back to vector search)
 	result, err = store.Search(ctx, &vectorstore.SearchQuery{
-		Vector:      []float64{0.1, 0.9, 0.2},
-		SearchMode:  vectorstore.SearchModeHybrid,
-		Limit:       5,
-		MinScore:    0.5,
+		Vector:     []float64{0.1, 0.9, 0.2},
+		SearchMode: vectorstore.SearchModeHybrid,
+		Limit:      5,
+		MinScore:   0.5,
 	})
 	require.NoError(t, err)
 	require.Len(t, result.Results, 1)
@@ -395,8 +395,8 @@ func TestVectorStore_SearchModes(t *testing.T) {
 
 	// Test keyword search mode (falls back to filter search)
 	result, err = store.Search(ctx, &vectorstore.SearchQuery{
-		Query:       "bonjour",
-		SearchMode:  vectorstore.SearchModeKeyword,
+		Query:      "bonjour",
+		SearchMode: vectorstore.SearchModeKeyword,
 		Filter: &vectorstore.SearchFilter{
 			Metadata: map[string]any{"lang": "fr"},
 		},

runner/runner.go

@@ -20,13 +20,11 @@ import (
 	"trpc.group/trpc-go/trpc-agent-go/agent"
 	"trpc.group/trpc-go/trpc-agent-go/artifact"
 	"trpc.group/trpc-go/trpc-agent-go/event"
-	itelemetry "trpc.group/trpc-go/trpc-agent-go/internal/telemetry"
 	"trpc.group/trpc-go/trpc-agent-go/log"
 	"trpc.group/trpc-go/trpc-agent-go/memory"
 	"trpc.group/trpc-go/trpc-agent-go/model"
 	"trpc.group/trpc-go/trpc-agent-go/session"
 	"trpc.group/trpc-go/trpc-agent-go/session/inmemory"
-	"trpc.group/trpc-go/trpc-agent-go/telemetry/trace"
 )
 
 // Author types for events.
@@ -190,8 +188,6 @@ func (r *runner) Run(
 	// transfer_to_agent that rely on agent.InvocationFromContext(ctx).
 	ctx = agent.NewInvocationContext(ctx, invocation)
 
-	ctx, span := trace.Tracer.Start(ctx, itelemetry.SpanNameInvocation)
-	defer span.End()
 	// Run the agent and get the event channel.
 	agentEventCh, err := r.agent.Run(ctx, invocation)
 	if err != nil {
@@ -255,8 +251,6 @@ func (r *runner) Run(
 		agent.EmitEvent(ctx, invocation, processedEventCh, runnerCompletionEvent)
 	}()
 
-	itelemetry.TraceRunner(span, r.appName, invocation, message)
-
 	return processedEventCh, nil
 }
 

telemetry/langfuse/exporter.go

@@ -108,8 +108,7 @@ func transformSpan(span *tracepb.Span) {
 		transformCallLLM(span)
 	case itelemetry.OperationExecuteTool:
 		transformExecuteTool(span)
-	case itelemetry.OperationRunRunner:
-		transformRunRunner(span)
+	default:
 	}
 }
 
@@ -247,63 +246,6 @@ func transformExecuteTool(span *tracepb.Span) {
 	span.Attributes = newAttributes
 }
 
-// transformRunRunner transforms runner spans for Langfuse
-func transformRunRunner(span *tracepb.Span) {
-	var newAttributes []*commonpb.KeyValue
-
-	newAttributes = append(newAttributes, &commonpb.KeyValue{
-		Key: observationType,
-		Value: &commonpb.AnyValue{
-			Value: &commonpb.AnyValue_StringValue{StringValue: "agent"},
-		},
-	})
-	// Process existing attributes
-	for _, attr := range span.Attributes {
-		switch attr.Key {
-		case itelemetry.KeyRunnerInput:
-			if attr.Value != nil {
-				newAttributes = append(newAttributes, &commonpb.KeyValue{
-					Key: observationInput,
-					Value: &commonpb.AnyValue{
-						Value: &commonpb.AnyValue_StringValue{StringValue: attr.Value.GetStringValue()},
-					},
-				})
-			} else {
-				newAttributes = append(newAttributes, &commonpb.KeyValue{
-					Key: observationInput,
-					Value: &commonpb.AnyValue{
-						Value: &commonpb.AnyValue_StringValue{StringValue: "N/A"},
-					},
-				})
-			}
-			// Skip this attribute (delete it)
-		case itelemetry.KeyRunnerOutput:
-			if attr.Value != nil {
-				newAttributes = append(newAttributes, &commonpb.KeyValue{
-					Key: observationOutput,
-					Value: &commonpb.AnyValue{
-						Value: &commonpb.AnyValue_StringValue{StringValue: attr.Value.GetStringValue()},
-					},
-				})
-			} else {
-				newAttributes = append(newAttributes, &commonpb.KeyValue{
-					Key: observationOutput,
-					Value: &commonpb.AnyValue{
-						Value: &commonpb.AnyValue_StringValue{StringValue: "N/A"},
-					},
-				})
-			}
-			// Skip this attribute (delete it)
-		default:
-			// Keep other attributes
-			newAttributes = append(newAttributes, attr)
-		}
-	}
-
-	// Replace span attributes
-	span.Attributes = newAttributes
-}
-
 func (e *exporter) Shutdown(ctx context.Context) error {
 	e.mu.RLock()
 	started := e.started