package compat import ( "context" "encoding/json" "io" "log/slog" "net/http" "net/http/httptest" "sync" "testing" ) func TestExtractMetadataFromStreamingResponse(t *testing.T) { t.Parallel() server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() var req chatCompletionsRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } if req.Stream == nil || !*req.Stream { t.Fatalf("stream flag = %v, want true", req.Stream) } w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"{\\\"people\\\":[],\"}}]}\n\n") _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"\\\"action_items\\\":[],\\\"dates_mentioned\\\":[],\"}}]}\n\n") _, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"\\\"topics\\\":[\\\"android\\\"],\\\"type\\\":\\\"idea\\\",\\\"source\\\":\\\"stream\\\"}\"}}]}\n\n") _, _ = io.WriteString(w, "data: [DONE]\n\n") })) defer server.Close() client := New(Config{ Name: "litellm", BaseURL: server.URL, APIKey: "test-key", MetadataModel: "qwen3.5:latest", Temperature: 0.1, HTTPClient: server.Client(), Log: slog.New(slog.NewTextHandler(io.Discard, nil)), EmbeddingModel: "unused", }) metadata, err := client.ExtractMetadata(context.Background(), "Project idea: Build an Android companion app.") if err != nil { t.Fatalf("ExtractMetadata() error = %v", err) } if metadata.Type != "idea" { t.Fatalf("metadata type = %q, want idea", metadata.Type) } if metadata.Source != "stream" { t.Fatalf("metadata source = %q, want stream", metadata.Source) } if len(metadata.Topics) != 1 || metadata.Topics[0] != "android" { t.Fatalf("metadata topics = %#v, want [android]", metadata.Topics) } } func TestExtractMetadataRetriesWithoutJSONMode(t *testing.T) { t.Parallel() var mu sync.Mutex jsonModeCalls := 0 plainCalls := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() var req chatCompletionsRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } if req.ResponseFormat != nil && req.ResponseFormat.Type == "json_object" { mu.Lock() jsonModeCalls++ mu.Unlock() _, _ = io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":""}}]}`) return } mu.Lock() plainCalls++ mu.Unlock() _, _ = io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"{\"people\":[],\"action_items\":[],\"dates_mentioned\":[],\"topics\":[\"android\"],\"type\":\"idea\",\"source\":\"test\"}"}}]}`) })) defer server.Close() client := New(Config{ Name: "litellm", BaseURL: server.URL, APIKey: "test-key", MetadataModel: "qwen3.5:latest", Temperature: 0.1, HTTPClient: server.Client(), Log: slog.New(slog.NewTextHandler(io.Discard, nil)), EmbeddingModel: "unused", }) metadata, err := client.ExtractMetadata(context.Background(), "Project idea: Build an Android companion app.") if err != nil { t.Fatalf("ExtractMetadata() error = %v", err) } if metadata.Type != "idea" { t.Fatalf("metadata type = %q, want idea", metadata.Type) } if metadata.Source != "test" { t.Fatalf("metadata source = %q, want test", metadata.Source) } mu.Lock() defer mu.Unlock() if jsonModeCalls != maxMetadataAttempts { t.Fatalf("json mode calls = %d, want %d", jsonModeCalls, maxMetadataAttempts) } if plainCalls != 1 { t.Fatalf("plain calls = %d, want 1", plainCalls) } } func TestExtractMetadataBypassesInvalidFallbackModelAfterFirstFailure(t *testing.T) { t.Parallel() var mu sync.Mutex primaryCalls := 0 invalidFallbackCalls := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { defer r.Body.Close() var req chatCompletionsRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } switch req.Model { case "empty-primary": _, _ = io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":""}}]}`) case "qwen3.5:latest": mu.Lock() primaryCalls++ mu.Unlock() _, _ = io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"{\"people\":[],\"action_items\":[],\"dates_mentioned\":[],\"topics\":[\"metadata\"],\"type\":\"observation\",\"source\":\"primary\"}"}}]}`) case "qwen3": mu.Lock() invalidFallbackCalls++ mu.Unlock() w.WriteHeader(http.StatusBadRequest) _, _ = io.WriteString(w, "{\"error\":{\"message\":\"{'error': '/chat/completions: Invalid model name passed in model=qwen3. Call `/v1/models` to view available models for your key.'}\"}}") default: t.Fatalf("unexpected model %q", req.Model) } })) defer server.Close() client := New(Config{ Name: "litellm", BaseURL: server.URL, APIKey: "test-key", MetadataModel: "empty-primary", FallbackMetadataModels: []string{"qwen3", "qwen3.5:latest"}, Temperature: 0.1, HTTPClient: server.Client(), Log: slog.New(slog.NewTextHandler(io.Discard, nil)), EmbeddingModel: "unused", }) for i := 0; i < 2; i++ { metadata, err := client.ExtractMetadata(context.Background(), "A short note about metadata.") if err != nil { t.Fatalf("ExtractMetadata() error = %v", err) } if metadata.Source != "primary" { t.Fatalf("metadata source = %q, want primary", metadata.Source) } } mu.Lock() defer mu.Unlock() if invalidFallbackCalls != 1 { t.Fatalf("invalid fallback calls = %d, want 1", invalidFallbackCalls) } if primaryCalls != 2 { t.Fatalf("valid fallback calls = %d, want 2", primaryCalls) } }