enkly-ai/pkg/vllmkit/client_test.go

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")
	}
}