first commit

core/providers/vllm/chat_test.go

@@ -0,0 +1,110 @@
+package vllm
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	schemas "github.com/maximhq/bifrost/core/schemas"
+)
+
+// TestChatCompletion_ExtraParamsForwardedAutomatically verifies that provider-specific
+// extra params (e.g. chat_template_kwargs) are forwarded to vLLM without requiring
+// the caller to set BifrostContextKeyPassthroughExtraParams on the context.
+func TestChatCompletion_ExtraParamsForwardedAutomatically(t *testing.T) {
+	t.Parallel()
+
+	var capturedBody map[string]interface{}
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			http.Error(w, "read error", http.StatusInternalServerError)
+			return
+		}
+		if err := json.Unmarshal(body, &capturedBody); err != nil {
+			http.Error(w, "json error", http.StatusInternalServerError)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		fmt.Fprint(w, `{
+			"id": "chatcmpl-test",
+			"object": "chat.completion",
+			"created": 1234567890,
+			"model": "gemma",
+			"choices": [{
+				"index": 0,
+				"message": {"role": "assistant", "content": "Hello!"},
+				"finish_reason": "stop"
+			}],
+			"usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}
+		}`)
+	}))
+	defer server.Close()
+
+	provider := newTestVLLMProvider()
+	key := schemas.Key{
+		ID:    "test-key",
+		Value: schemas.EnvVar{Val: "test-api-key"},
+		VLLMKeyConfig: &schemas.VLLMKeyConfig{
+			URL: schemas.EnvVar{Val: server.URL},
+		},
+	}
+
+	// Intentionally do NOT set BifrostContextKeyPassthroughExtraParams — the provider
+	// should set it automatically.
+	ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+
+	hello := "Hello"
+	req := &schemas.BifrostChatRequest{
+		Provider: schemas.VLLM,
+		Model:    "gemma",
+		Input: []schemas.ChatMessage{
+			{
+				Role:    schemas.ChatMessageRoleUser,
+				Content: &schemas.ChatMessageContent{ContentStr: &hello},
+			},
+		},
+		Params: &schemas.ChatParameters{
+			ExtraParams: map[string]interface{}{
+				"chat_template_kwargs": map[string]interface{}{
+					"enable_thinking": true,
+				},
+			},
+		},
+	}
+
+	_, bifrostErr := provider.ChatCompletion(ctx, key, req)
+	if bifrostErr != nil {
+		t.Fatalf("ChatCompletion returned error: %v", bifrostErr.Error.Message)
+	}
+
+	if capturedBody == nil {
+		t.Fatal("mock server did not receive a request body")
+	}
+
+	rawKwargs, ok := capturedBody["chat_template_kwargs"]
+	if !ok {
+		t.Fatalf("chat_template_kwargs missing from outgoing request body; got keys: %v", keys(capturedBody))
+	}
+
+	kwargsMap, ok := rawKwargs.(map[string]interface{})
+	if !ok {
+		t.Fatalf("expected chat_template_kwargs to be an object, got %T", rawKwargs)
+	}
+	if kwargsMap["enable_thinking"] != true {
+		t.Fatalf("expected enable_thinking=true, got %v", kwargsMap["enable_thinking"])
+	}
+}
+
+func keys(m map[string]interface{}) []string {
+	out := make([]string, 0, len(m))
+	for k := range m {
+		out = append(out, k)
+	}
+	return out
+}