(t *testing.T)
| 59 | } |
| 60 | |
| 61 | func TestTelemetryMiddleware_WrapModelCall(t *testing.T) { |
| 62 | tracer := telemetry.NewSimpleTracer() |
| 63 | |
| 64 | m := NewTelemetryMiddleware(&TelemetryMiddlewareConfig{ |
| 65 | Tracer: tracer, |
| 66 | AgentID: "test-agent", |
| 67 | AgentName: "Test Agent", |
| 68 | Provider: genai.ProviderAnthropic, |
| 69 | Model: "claude-3-5-sonnet", |
| 70 | ConversationID: "conv-123", |
| 71 | RecordPrompts: true, |
| 72 | RecordCompletions: true, |
| 73 | }) |
| 74 | |
| 75 | ctx := context.Background() |
| 76 | req := &ModelRequest{ |
| 77 | Messages: []types.Message{ |
| 78 | {Role: types.MessageRoleUser, Content: "Hello"}, |
| 79 | }, |
| 80 | SystemPrompt: "You are a helpful assistant", |
| 81 | Metadata: map[string]any{ |
| 82 | "max_tokens": 1000, |
| 83 | "temperature": 0.7, |
| 84 | }, |
| 85 | } |
| 86 | |
| 87 | // Mock handler |
| 88 | handler := func(ctx context.Context, req *ModelRequest) (*ModelResponse, error) { |
| 89 | return &ModelResponse{ |
| 90 | Message: types.Message{ |
| 91 | Role: types.MessageRoleAssistant, |
| 92 | Content: "Hello! How can I help you?", |
| 93 | }, |
| 94 | Metadata: map[string]any{ |
| 95 | "input_tokens": 10, |
| 96 | "output_tokens": 8, |
| 97 | "response_id": "resp-123", |
| 98 | "stop_reason": "end_turn", |
| 99 | }, |
| 100 | }, nil |
| 101 | } |
| 102 | |
| 103 | resp, err := m.WrapModelCall(ctx, req, handler) |
| 104 | if err != nil { |
| 105 | t.Fatalf("unexpected error: %v", err) |
| 106 | } |
| 107 | if resp == nil { |
| 108 | t.Fatal("expected response") |
| 109 | } |
| 110 | |
| 111 | // Verify spans |
| 112 | spans := tracer.GetSpans() |
| 113 | if len(spans) != 1 { |
| 114 | t.Fatalf("expected 1 span, got %d", len(spans)) |
| 115 | } |
| 116 | |
| 117 | span := spans[0] |
| 118 | if span.Name() != "chat claude-3-5-sonnet" { |
nothing calls this directly
no test coverage detected