llmgw/internal/router/quality_test.go

package router

import (
	"encoding/json"
	"fmt"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"
	"testing/fstest"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/bench/eval"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// answeringEngine is an engine that answers a few questions right, one
// wrong, and grades as a judge would.
func answeringEngine(t *testing.T, served string) (*httptest.Server, *atomic.Int32) {
	t.Helper()
	var calls atomic.Int32
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintf(w, "# TYPE vllm:num_requests_running gauge\nvllm:num_requests_running 0\n")
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		var req struct {
			Model    string `json:"model"`
			Messages []struct {
				Content string `json:"content"`
			} `json:"messages"`
			Stream bool `json:"stream"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Stream || len(req.Messages) == 0 {
			http.Error(w, "bad request", http.StatusBadRequest)
			return
		}
		q := req.Messages[len(req.Messages)-1].Content
		answer := "I am not sure."
		switch {
		case served == "grader" && strings.Contains(q, "Answer:\nParis"):
			answer = `{"score": 1, "reason": "names Paris"}`
		case served == "grader":
			answer = `{"score": 0, "reason": "does not name Paris"}`
		case strings.Contains(q, "17 + 25"):
			answer = "<think>17 and 25</think>42"
		case strings.Contains(q, "6 * 7"):
			answer = "6 * 7 = 43"
		case strings.Contains(q, "JSON"):
			answer = "```json\n{\"ok\": true}\n```"
		case strings.Contains(q, "capital"):
			answer = "Paris"
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(map[string]any{"model": req.Model, "choices": []any{map[string]any{
			"message": map[string]any{"role": "assistant", "content": answer}, "finish_reason": "stop"}},
			"usage": map[string]any{"prompt_tokens": 5, "completion_tokens": 3}})
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv, &calls
}

const qualitySuite = `
name: smoke
threshold: 0.7
judge: grader
cases:
  - id: add
    prompt: What is 17 + 25?
    number: 42
  - id: mul
    prompt: What is 6 * 7?
    number: 42
  - id: json
    prompt: Give a JSON object with ok true.
    json_schema: {type: object, required: [ok], properties: {ok: {const: true}}}
  - id: capital
    prompt: What is the capital of France?
    judge: It names Paris.
    weight: 2
`

func TestASuiteIsAskedThroughTheRouterAndScoredAndKept(t *testing.T) {
	model, modelCalls := answeringEngine(t, "qwen")
	grader, graderCalls := answeringEngine(t, "grader")
	rt := newTestRouter(t, baseConfig(
		config.Endpoint{ID: "a", URL: model.URL, Models: []string{"qwen"}},
		config.Endpoint{ID: "j", URL: grader.URL, Models: []string{"grader"}},
	))
	db, err := store.OpenSQLite(t.Context(), filepath.Join(t.TempDir(), "store.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = db.Close() })
	svc := eval.NewService(eval.Deps{Call: rt.Replay, Store: db, Log: slog.New(slog.DiscardHandler),
		Builtin: fstest.MapFS{"smoke.yaml": {Data: []byte(qualitySuite)}}})

	live, err := svc.Start("qwen", eval.Request{Suite: "smoke", Concurrency: 2}, "tester")
	if err != nil {
		t.Fatal(err)
	}
	deadline := time.Now().Add(10 * time.Second)
	for live.State == eval.Running && time.Now().Before(deadline) {
		time.Sleep(20 * time.Millisecond)
		live, _ = svc.Get(t.Context(), live.ID)
	}
	// add 1, mul 0, json 1, capital 1 at weight 2: 4 of 5.
	if live.State != eval.Done || live.Score != 0.8 || !live.Passed || live.Below != 1 || live.Errors != 0 {
		t.Fatalf("%+v", live)
	}
	if modelCalls.Load() != 4 || graderCalls.Load() != 1 {
		t.Errorf("the model was asked %d times and the judge %d", modelCalls.Load(), graderCalls.Load())
	}
	kept, err := db.EvalRunByID(t.Context(), live.ID)
	if err != nil || len(kept.Results) != 4 || kept.Results[0].Answer != "42" ||
		!strings.Contains(kept.Results[1].Detail, "read 43") || kept.Results[3].Detail != "the judge: names Paris" {
		t.Fatalf("kept %+v %v", kept, err)
	}

	// A model nothing serves scores nothing, and says why.
	run, err := svc.Run(t.Context(), "nothing", eval.Request{Suite: "smoke"}, "tester", "")
	if err != nil || run.Passed || run.Errors != 4 || !strings.Contains(run.Error, "no case was answered") {
		t.Fatalf("%+v %v", run, err)
	}
}