package vllmkit import ( "bufio" "context" "encoding/json" "fmt" "net/http" "net/http/httptest" "testing" ) func TestThinkWire(t *testing.T) { cases := []struct { name string setting, override string wantEnable *bool // nil means "chat_template_kwargs omitted entirely" }{ {"both empty", "", "", nil}, {"setting only, bool true", "true", "", boolPtr(true)}, {"setting only, bool false", "false", "", boolPtr(false)}, {"setting only, level", "high", "", boolPtr(true)}, {"override wins over setting", "low", "high", boolPtr(true)}, {"override alone", "", "medium", boolPtr(true)}, {"override explicit false wins", "true", "false", boolPtr(false)}, {"case/space insensitive", " TRUE ", "", boolPtr(true)}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { got := thinkWire(tc.setting, tc.override) if tc.wantEnable == nil { if got != nil { t.Errorf("thinkWire(%q, %q) = %#v, want nil", tc.setting, tc.override, got) } return } gotEnable, ok := got["enable_thinking"].(bool) if !ok || gotEnable != *tc.wantEnable { t.Errorf("thinkWire(%q, %q) = %#v, want enable_thinking=%v", tc.setting, tc.override, got, *tc.wantEnable) } }) } } func boolPtr(b bool) *bool { return &b } // TestChatStream_ThinkOverride confirms the per-call think parameter // reaches the wire request's chat_template_kwargs.enable_thinking, // taking precedence over the Client's configured default exactly as // thinkWire specifies — this is the knob agentrun.Agent.ThinkPolicy // relies on to vary reasoning on or off by round. func TestChatStream_ThinkOverride(t *testing.T) { var gotKwargs map[string]any sawKwargs := false srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var req chatRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { t.Fatalf("decode request: %v", err) } gotKwargs = req.ChatTemplateKwargs sawKwargs = req.ChatTemplateKwargs != nil w.Header().Set("Content-Type", "text/event-stream") fmt.Fprintf(w, "data: %s\n\n", mustJSON(finishedChunk())) fmt.Fprintf(w, "data: [DONE]\n\n") })) defer srv.Close() run := func(clientThink, callThink string) (map[string]any, bool) { c := New(srv.URL, "", "test-model", "") c.Think = clientThink events, err := c.ChatStream(context.Background(), []Message{{Role: RoleUser, Content: "hi"}}, nil, callThink) if err != nil { t.Fatalf("ChatStream: %v", err) } for ev := range events { if ev.Err != nil { t.Fatalf("stream error: %v", ev.Err) } } return gotKwargs, sawKwargs } if got, saw := run("", ""); saw || got != nil { t.Errorf("client default empty, no override: got kwargs=%#v saw=%v, want omitted", got, saw) } if got, _ := run("low", ""); got["enable_thinking"] != true { t.Errorf("client default only: got kwargs=%#v, want enable_thinking=true", got) } if got, _ := run("false", "high"); got["enable_thinking"] != true { t.Errorf("override should win: got kwargs=%#v, want enable_thinking=true", got) } if got, _ := run("true", "false"); got["enable_thinking"] != false { t.Errorf("override alone (explicit false): got kwargs=%#v, want enable_thinking=false", got) } } func finishedChunk() chatStreamChunk { var c chatStreamChunk c.Choices = make([]struct { Delta struct { Content string `json:"content"` ReasoningContent string `json:"reasoning_content"` ToolCalls []wireToolCall `json:"tool_calls"` } `json:"delta"` FinishReason *string `json:"finish_reason"` }, 1) reason := "stop" c.Choices[0].FinishReason = &reason return c } func mustJSON(v any) string { b, err := json.Marshal(v) if err != nil { panic(err) } return string(b) } // TestChatStream_StreamsTokens is a sanity check that the SSE-based // stream reader in ChatStream still behaves with an // bufio.Scanner-friendly response (guards against accidentally // breaking the "data: " framing while touching this file). func TestChatStream_StreamsTokens(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") bw := bufio.NewWriter(w) write := func(content string, done bool) { var c chatStreamChunk c.Choices = make([]struct { Delta struct { Content string `json:"content"` ReasoningContent string `json:"reasoning_content"` ToolCalls []wireToolCall `json:"tool_calls"` } `json:"delta"` FinishReason *string `json:"finish_reason"` }, 1) c.Choices[0].Delta.Content = content if done { reason := "stop" c.Choices[0].FinishReason = &reason } fmt.Fprintf(bw, "data: %s\n\n", mustJSON(c)) } write("hel", false) write("lo", true) fmt.Fprint(bw, "data: [DONE]\n\n") bw.Flush() })) defer srv.Close() c := New(srv.URL, "", "test-model", "") events, err := c.ChatStream(context.Background(), []Message{{Role: RoleUser, Content: "hi"}}, nil, "") if err != nil { t.Fatalf("ChatStream: %v", err) } var got string for ev := range events { if ev.Err != nil { t.Fatalf("stream error: %v", ev.Err) } got += ev.Token } if got != "hello" { t.Errorf("got tokens %q, want %q", got, "hello") } } // TestToolCallAccumulator_Streamed confirms incremental, index-keyed // tool_calls deltas (arguments split across several chunks, as a real // vLLM server streams them) reassemble into one complete call. func TestToolCallAccumulator_Streamed(t *testing.T) { idx := 0 acc := newToolCallAccumulator() acc.add([]wireToolCall{{Index: &idx, ID: "call_1", Function: wireToolCallFunction{Name: "grep"}}}) acc.add([]wireToolCall{{Index: &idx, Function: wireToolCallFunction{Arguments: `{"pat`}}}) acc.add([]wireToolCall{{Index: &idx, Function: wireToolCallFunction{Arguments: `tern":"foo"}`}}}) calls, err := acc.finish() if err != nil { t.Fatalf("finish: %v", err) } if len(calls) != 1 { t.Fatalf("got %d calls, want 1", len(calls)) } if calls[0].ID != "call_1" || calls[0].Name != "grep" { t.Errorf("got %+v, want id=call_1 name=grep", calls[0]) } if calls[0].Arguments["pattern"] != "foo" { t.Errorf("got arguments %#v, want pattern=foo", calls[0].Arguments) } } func TestIsToolCallParseError(t *testing.T) { if !IsToolCallParseError(fmt.Errorf("wrap: %w", errMalformedToolArgs)) { t.Error("wrapped errMalformedToolArgs should be detected") } if !IsToolCallParseError(fmt.Errorf("llm: vllm error: BadRequestError: Failed to parse tool calls from model output")) { t.Error("server-side tool-call parse failure should be detected") } if IsToolCallParseError(fmt.Errorf("llm: vllm returned 500 Internal Server Error: out of memory")) { t.Error("unrelated server error should not be classified as a tool-call parse error") } if IsToolCallParseError(nil) { t.Error("nil error should never be a tool-call parse error") } }