package llm import ( "context" "encoding/json" "net/http" "net/http/httptest" "strings" "testing" "time" ) func TestComplete(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method != "POST" { t.Errorf("method = %q, want POST", r.Method) } if !strings.HasSuffix(r.URL.Path, "/v1/chat/completions") { t.Errorf("path = %q, want /v1/chat/completions", r.URL.Path) } var reqBody struct { Messages []struct { Role string `json:"role"` Content string `json:"content"` } `json:"messages"` Grammar string `json:"grammar"` MaxTokens int `json:"max_tokens"` } if err := json.NewDecoder(r.Body).Decode(&reqBody); err != nil { t.Fatalf("decode request body: %v", err) } if len(reqBody.Messages) < 2 { t.Fatalf("expected at least 2 messages, got %d", len(reqBody.Messages)) } if reqBody.Messages[0].Role != "system" || reqBody.Messages[1].Content != "hi" { t.Errorf("unexpected messages: %+v", reqBody.Messages) } if reqBody.Grammar != `root ::= "x"` { t.Errorf("grammar = %q, want root ::= \"x\"", reqBody.Grammar) } w.Header().Set("Content-Type", "application/json") w.Write([]byte(`{"choices":[{"message":{"content":"ok"}}]}`)) })) defer srv.Close() c := New(srv.URL, 5*time.Second) got, err := c.Complete(context.Background(), Req{ System: "be helpful", User: "hi", Grammar: `root ::= "x"`, MaxTokens: 42, }) if err != nil { t.Fatalf("Complete: %v", err) } if got != "ok" { t.Errorf("got %q, want %q", got, "ok") } } // TestSetBaseURL — a model swap re-points every holder of the client rather than // rebuilding the router, the replier and the extractors (Vikunja #250). func TestSetBaseURL(t *testing.T) { var hit string srvA := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hit = "A" w.Write([]byte(`{"choices":[{"message":{"content":"a"}}]}`)) })) defer srvA.Close() srvB := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hit = "B" w.Write([]byte(`{"choices":[{"message":{"content":"b"}}]}`)) })) defer srvB.Close() c := New(srvA.URL, 5*time.Second) if _, err := c.Complete(context.Background(), Req{User: "x"}); err != nil { t.Fatalf("Complete against A: %v", err) } if hit != "A" { t.Fatalf("first request went to %q; want A", hit) } c.SetBaseURL(srvB.URL) if got := c.BaseURL(); got != srvB.URL { t.Errorf("BaseURL = %q; want %q", got, srvB.URL) } if _, err := c.Complete(context.Background(), Req{User: "x"}); err != nil { t.Fatalf("Complete against B: %v", err) } if hit != "B" { t.Errorf("request after the swap went to %q; want B", hit) } }