velox-bot/internal/llm/ollama_test.go

package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestChatRequestAndReply(t *testing.T) {
	var got chatRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/chat" || r.Method != http.MethodPost {
			t.Errorf("unexpected %s %s", r.Method, r.URL.Path)
		}
		_ = json.NewDecoder(r.Body).Decode(&got)
		_, _ = w.Write([]byte(`{"message":{"role":"assistant","content":"<think>hmm</think>\n  the summary \n"}}`))
	}))
	defer srv.Close()

	out, err := New(srv.URL+"/", "some-model", 6).Chat(context.Background(), "sys", "usr")
	if err != nil {
		t.Fatal(err)
	}
	if out != "the summary" {
		t.Errorf("reply = %q (think block and whitespace should be stripped)", out)
	}
	if got.Model != "some-model" || got.Stream || len(got.Messages) != 2 ||
		got.Messages[0].Role != "system" || got.Messages[1].Content != "usr" {
		t.Errorf("request = %+v", got)
	}
	if got.Options["num_ctx"].(float64) != numCtx {
		t.Errorf("num_ctx = %v", got.Options["num_ctx"])
	}
	if got.Options["num_thread"].(float64) != 6 {
		t.Errorf("num_thread = %v, want 6", got.Options["num_thread"])
	}
}

func TestChatOmitsThreadsWhenUnset(t *testing.T) {
	var got chatRequest
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewDecoder(r.Body).Decode(&got)
		_, _ = w.Write([]byte(`{"message":{"content":"ok"}}`))
	}))
	defer srv.Close()

	if _, err := New(srv.URL, "m", 0).Chat(context.Background(), "s", "u"); err != nil {
		t.Fatal(err)
	}
	if _, set := got.Options["num_thread"]; set {
		t.Error("num_thread must be left to Ollama when threads is 0")
	}
}

func TestChatErrors(t *testing.T) {
	cases := map[string]struct {
		status int
		body   string
		want   string
	}{
		"http error":  {http.StatusNotFound, `model not found`, "404"},
		"api error":   {http.StatusOK, `{"error":"out of memory"}`, "out of memory"},
		"empty reply": {http.StatusOK, `{"message":{"content":"<think>x</think>"}}`, "empty"},
		"bad json":    {http.StatusOK, `nope`, "decoding"},
	}
	for name, c := range cases {
		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(c.status)
			_, _ = w.Write([]byte(c.body))
		}))
		_, err := New(srv.URL, "m", 0).Chat(context.Background(), "s", "u")
		srv.Close()
		if err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: err = %v, want containing %q", name, err, c.want)
		}
	}

	if _, err := New("http://127.0.0.1:1", "m", 0).Chat(context.Background(), "s", "u"); err == nil || !strings.Contains(err.Error(), "unreachable") {
		t.Errorf("unreachable server: err = %v", err)
	}
}

func TestChatDisablesThinkingAndFallsBack(t *testing.T) {
	var seen []bool // whether each request carried think:false
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req chatRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		hasThink := req.Think != nil && !*req.Think
		seen = append(seen, hasThink)
		if hasThink {
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":"\"some-model\" does not support thinking"}`))
			return
		}
		_, _ = w.Write([]byte(`{"message":{"content":"fine"}}`))
	}))
	defer srv.Close()

	out, err := New(srv.URL, "some-model", 0).Chat(context.Background(), "s", "u")
	if err != nil || out != "fine" {
		t.Fatalf("out=%q err=%v", out, err)
	}
	if len(seen) != 2 || !seen[0] || seen[1] {
		t.Errorf("want one request with think:false then a retry without it, got %v", seen)
	}
}

func TestChatDoesNotRetryOtherBadRequests(t *testing.T) {
	calls := 0
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		calls++
		w.WriteHeader(http.StatusBadRequest)
		_, _ = w.Write([]byte(`{"error":"model is required"}`))
	}))
	defer srv.Close()

	if _, err := New(srv.URL, "", 0).Chat(context.Background(), "s", "u"); err == nil || calls != 1 {
		t.Errorf("err=%v calls=%d, want an error after exactly one call", err, calls)
	}
}

// streamServer replies to /api/chat with the given NDJSON lines, flushing
// after each so the client sees them as separate reads.
func streamServer(t *testing.T, gotReq *chatRequest, lines ...string) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if gotReq != nil {
			_ = json.NewDecoder(r.Body).Decode(gotReq)
		}
		for _, l := range lines {
			_, _ = w.Write([]byte(l + "\n"))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
		}
	}))
}

func TestChatStreamReportsGrowingTextAndHidesThinking(t *testing.T) {
	var req chatRequest
	srv := streamServer(t, &req,
		`{"message":{"content":"<think>pl"},"done":false}`,
		`{"message":{"content":"an</think>Olá"},"done":false}`,
		`{"message":{"content":" mundo"},"done":false}`,
		`{"message":{"content":""},"done":true}`,
	)
	defer srv.Close()

	var seen []string
	out, err := New(srv.URL, "m", 6).ChatStream(context.Background(), "s", "u", func(s string) { seen = append(seen, s) })
	if err != nil {
		t.Fatal(err)
	}
	if out != "Olá mundo" {
		t.Errorf("final = %q", out)
	}
	// The first chunk is only an unfinished think block, so nothing is shown for it.
	if strings.Join(seen, "|") != "Olá|Olá mundo" {
		t.Errorf("callbacks = %v", seen)
	}
	if !req.Stream {
		t.Error("request must ask for stream:true")
	}
	if req.Think == nil || *req.Think {
		t.Error("request must carry think:false")
	}
}

func TestChatStreamWorksWithoutCallback(t *testing.T) {
	srv := streamServer(t, nil, `{"message":{"content":"ok"},"done":true}`)
	defer srv.Close()
	if out, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err != nil || out != "ok" {
		t.Errorf("out=%q err=%v", out, err)
	}
}

func TestChatStreamErrors(t *testing.T) {
	srv := streamServer(t, nil, `{"message":{"content":"par"},"done":false}`, `{"error":"model crashed"}`)
	if _, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err == nil || !strings.Contains(err.Error(), "model crashed") {
		t.Errorf("mid-stream error: %v", err)
	}
	srv.Close()

	srv = streamServer(t, nil, `not json`)
	if _, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err == nil || !strings.Contains(err.Error(), "decoding") {
		t.Errorf("bad json: %v", err)
	}
	srv.Close()

	srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		_, _ = w.Write([]byte("model not found"))
	}))
	defer srv.Close()
	if _, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err == nil || !strings.Contains(err.Error(), "404") {
		t.Errorf("bad status: %v", err)
	}

	empty := streamServer(t, nil, `{"message":{"content":"<think>only</think>"},"done":true}`)
	defer empty.Close()
	if _, err := New(empty.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err == nil || !strings.Contains(err.Error(), "empty") {
		t.Errorf("empty reply: %v", err)
	}
}

func TestChatStreamKeepsWhatArrivedIfTheStreamEndsEarly(t *testing.T) {
	srv := streamServer(t, nil, `{"message":{"content":"half an ans"},"done":false}`) // no done:true
	defer srv.Close()
	if out, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil); err != nil || out != "half an ans" {
		t.Errorf("out=%q err=%v", out, err)
	}
}

func TestChatStreamFallsBackWhenThinkingUnsupported(t *testing.T) {
	var seen []bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req chatRequest
		_ = json.NewDecoder(r.Body).Decode(&req)
		hasThink := req.Think != nil
		seen = append(seen, hasThink)
		if hasThink {
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":"\"m\" does not support thinking"}`))
			return
		}
		_, _ = w.Write([]byte(`{"message":{"content":"fine"},"done":true}` + "\n"))
	}))
	defer srv.Close()

	out, err := New(srv.URL, "m", 0).ChatStream(context.Background(), "s", "u", nil)
	if err != nil || out != "fine" || len(seen) != 2 || !seen[0] || seen[1] {
		t.Errorf("out=%q err=%v seen=%v", out, err, seen)
	}
}