llmgw/internal/router/chat_models_test.go

package router

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// chatRouter mirrors what the control plane publishes: an endpoint advertises
// every name a client may address it by, which is the model's own name and all
// of its aliases.
func chatRouter(t *testing.T, engineURL string) (*Router, store.User) {
	t.Helper()

	cfg := baseConfig(config.Endpoint{
		ID: "a", URL: engineURL,
		Models: []string{"qwen2.5-coder-7b-instruct", "coder-7b", "coder"},
	})
	cfg.Models = []catalog.Model{{
		Name: "qwen2.5-coder-7b-instruct", Aliases: []string{"coder-7b", "coder"},
	}}
	cfg.Auth.Enabled = true
	cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	rt := newTestRouter(t, cfg)

	org, err := rt.store.CreateOrg(t.Context(), "acme")
	if err != nil {
		t.Fatal(err)
	}
	team, err := rt.store.CreateTeam(t.Context(), org, "platform")
	if err != nil {
		t.Fatal(err)
	}
	project, err := rt.store.CreateProject(t.Context(), team, "gateway")
	if err != nil {
		t.Fatal(err)
	}
	return rt, store.User{ID: "usr-1", ProjectID: project}
}

func TestTheChatPickerOffersAModelOnceRatherThanOncePerAlias(t *testing.T) {
	// An endpoint advertises the model's name and every alias, because that is
	// what the router matches a request against. Listing those strings as they
	// came gave the picker "coder", "coder-7b" and the full name as three
	// separate choices for one model.
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, user := chatRouter(t, e.srv.URL)
	waitForEndpoint(t, rt)

	got := rt.modelsFor(httptest.NewRequest("GET", "/chat", nil), user)
	if len(got) != 1 {
		names := make([]string, len(got))
		for i, m := range got {
			names[i] = m.Name
		}
		t.Fatalf("the picker offers %d entries for one model: %v", len(got), names)
	}
	// The catalog name, not whichever alias happened to be listed first: it is
	// the one that appears in the model list, the usage table and the logs.
	if got[0].Name != "qwen2.5-coder-7b-instruct" {
		t.Errorf("offered as %q, expected the catalog name", got[0].Name)
	}
}

func waitForEndpoint(t *testing.T, rt *Router) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if len(rt.reg.Snapshots("")) > 0 {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal("no endpoint became available")
}

func TestTheChatPageDoesNotAskForMoreThanTheModelAllows(t *testing.T) {
	// The page imposes its own ceiling so a playground turn cannot occupy a
	// replica for ten minutes. Sending that ceiling regardless of the model
	// made the gateway refuse its own page: every turn on a model declaring a
	// smaller max_output came back as a 400 about a number nobody typed.
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	cfg := baseConfig(config.Endpoint{
		ID: "a", URL: e.srv.URL, Models: []string{"qwen2.5-coder-7b-instruct"},
	})
	cfg.Models = []catalog.Model{{
		Name: "qwen2.5-coder-7b-instruct",
		// Smaller than the page's own cap, which is the whole case.
		ContextLength: 32768, MaxOutput: 1024,
	}}
	cfg.Auth.Enabled = true
	cfg.Auth.Store = filepath.Join(t.TempDir(), "store.db")
	rt := newTestRouter(t, cfg)

	org, err := rt.store.CreateOrg(t.Context(), "acme")
	if err != nil {
		t.Fatal(err)
	}
	team, err := rt.store.CreateTeam(t.Context(), org, "platform")
	if err != nil {
		t.Fatal(err)
	}
	project, err := rt.store.CreateProject(t.Context(), team, "gateway")
	if err != nil {
		t.Fatal(err)
	}
	user := store.User{ID: "usr-1", ProjectID: project}
	waitForEndpoint(t, rt)

	body := `{"model":"qwen2.5-coder-7b-instruct","messages":[{"role":"user","content":"hi"}]}`
	req := httptest.NewRequest("POST", "/chat/send", strings.NewReader(body))
	rec := httptest.NewRecorder()
	rt.chatSend(rec, req, user)

	if rec.Code != 200 {
		t.Fatalf("the page's own turn was refused: %d %s", rec.Code, rec.Body.String())
	}
}

func TestTheChatCapIsKeptWhereTheModelAllowsMore(t *testing.T) {
	// The lower of the two, not the model's: a model with no ceiling, or a
	// generous one, must not let a playground turn run unbounded.
	if got := chatTokenCap(catalog.Model{MaxOutput: 8192}); got != chatMaxTokens {
		t.Errorf("a generous model raised the page's cap to %d", got)
	}
	if got := chatTokenCap(catalog.Model{}); got != chatMaxTokens {
		t.Errorf("a model that declares nothing gave %d", got)
	}
	if got := chatTokenCap(catalog.Model{MaxOutput: 1024}); got != 1024 {
		t.Errorf("a model allowing 1024 got %d", got)
	}
}

func TestTheChatPickerOffersAModelThatLoadsWhenAsked(t *testing.T) {
	// An on-demand model has no endpoint until a request asks for it, and on
	// this page a message is that request. Listing only running models left
	// such a model impossible to start from the chat.
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, user := chatRouter(t, e.srv.URL)
	cpu := []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}
	c, err := catalog.New([]catalog.Model{
		{Name: "qwen2.5-coder-7b-instruct", Aliases: []string{"coder-7b", "coder"}, Source: "http://m/a.gguf", Profiles: cpu},
		{Name: "mistral-nemo-12b-instruct", Aliases: []string{"nemo"}, Source: "http://m/b.gguf", Profiles: cpu,
			Scaling: catalog.Scaling{Min: 0, Max: 1}},
		// Stopped outright: nothing will load it, so offering it would only
		// produce an error.
		{Name: "phi-4", Source: "http://m/c.gguf", Profiles: cpu, Scaling: catalog.Scaling{Min: 0, Max: 0}},
	})
	if err != nil {
		t.Fatal(err)
	}
	rt.SetCatalog(c)
	waitForEndpoint(t, rt)

	got := rt.modelsFor(httptest.NewRequest("GET", "/chat", nil), user)
	byName := map[string]bool{}
	for _, m := range got {
		byName[m.Name] = m.OnDemand
	}
	if onDemand, ok := byName["mistral-nemo-12b-instruct"]; !ok || !onDemand {
		t.Errorf("the on-demand model is not offered as one: %+v", got)
	}
	if _, ok := byName["nemo"]; ok {
		t.Errorf("the alias is offered as a second model: %+v", got)
	}
	if _, ok := byName["phi-4"]; ok {
		t.Errorf("a stopped model is offered: %+v", got)
	}
	if onDemand := byName["qwen2.5-coder-7b-instruct"]; onDemand {
		t.Errorf("a running model is marked as not loaded")
	}
}

func TestAChatTurnCarriesItsTemperatureAndAsksForUsage(t *testing.T) {
	// The page shows how many tokens an answer took and how fast, from the
	// engine's own count, and lets a person turn the temperature — held to
	// the range every engine takes, whatever the page sends.
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, user := chatRouter(t, e.srv.URL)
	waitForEndpoint(t, rt)

	body := `{"model":"qwen2.5-coder-7b-instruct","temperature":5,"messages":[{"role":"system","content":"be brief"},{"role":"user","content":"hi"}]}`
	rec := httptest.NewRecorder()
	rt.chatSend(rec, httptest.NewRequest("POST", "/chat/send", strings.NewReader(body)), user)
	if rec.Code != 200 {
		t.Fatalf("%d %s", rec.Code, rec.Body.String())
	}
	var sent struct {
		Temperature   float64 `json:"temperature"`
		StreamOptions struct {
			IncludeUsage bool `json:"include_usage"`
		} `json:"stream_options"`
		Messages []map[string]string `json:"messages"`
	}
	raw, _ := e.last.Load().([]byte)
	if err := json.Unmarshal(raw, &sent); err != nil {
		t.Fatal(err, string(raw))
	}
	if sent.Temperature != 2 || !sent.StreamOptions.IncludeUsage || sent.Messages[0]["role"] != "system" {
		t.Errorf("sent %s", raw)
	}
}

func TestTheChatPageAsksForTheDocumentsChosenAndOnlyThose(t *testing.T) {
	// The picker's choice becomes the header; a header the page sent itself
	// is not trusted. On SQLite, asking is enough to tell: it is refused as
	// knowledge bases being off, before the model is asked.
	e := newFakeEngine(t, "qwen2.5-coder-7b-instruct")
	rt, user := chatRouter(t, e.srv.URL)
	waitForEndpoint(t, rt)
	send := func(knowledge, spoofed string) *httptest.ResponseRecorder {
		body := `{"model":"qwen2.5-coder-7b-instruct","knowledge":"` + knowledge + `","messages":[{"role":"user","content":"hi"}]}`
		req := httptest.NewRequest("POST", "/chat/send", strings.NewReader(body))
		if spoofed != "" {
			req.Header.Set(knowledgeHeader, spoofed)
		}
		rec := httptest.NewRecorder()
		rt.chatSend(rec, req, user)
		return rec
	}
	if rec := send("politikalar", ""); rec.Code != http.StatusNotImplemented || !strings.Contains(rec.Body.String(), "SQLite") {
		t.Errorf("a collection chosen: %d %s", rec.Code, rec.Body.String())
	}
	if rec := send("", "politikalar"); rec.Code != http.StatusOK {
		t.Errorf("a header the page set itself was used: %d %s", rec.Code, rec.Body.String())
	}
	if rec := send("none", ""); rec.Code != http.StatusOK {
		t.Errorf("none: %d %s", rec.Code, rec.Body.String())
	}
}