A chat model of the cluster reads each chunk and writes questions a person might ask of it with the answers it holds, streamed as NDJSON a chunk at a time, for people to correct before an adapter is trained on it. A chunk that gave no pairs says so on its own line. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
92 lines
3.3 KiB
Go
92 lines
3.3 KiB
Go
package adminapi
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||
"gitea.mynodes.xyz/ilker/llmgw/internal/store"
|
||
)
|
||
|
||
func TestPairsAreReadOutOfWhatAModelWrote(t *testing.T) {
|
||
completion := `{"choices":[{"message":{"content":"<think>let me see</think>\n{\"q\": \"Babalık izni kaç gün?\", \"a\": \"10 iş günü.\"},\nnot json\n{\"q\": \"\", \"a\": \"empty question\"}\n{\"q\": \"Ne zaman?\", \"a\": \"Doğumdan sonra.\"}"}}]}`
|
||
got := datasetPairs([]byte(completion))
|
||
if len(got) != 2 || got[0].Q != "Babalık izni kaç gün?" || got[1].A != "Doğumdan sonra." {
|
||
t.Errorf("pairs %+v", got)
|
||
}
|
||
if datasetPairs([]byte(`{"choices":[]}`)) != nil {
|
||
t.Error("pairs from no choice")
|
||
}
|
||
}
|
||
|
||
func TestADatasetIsWrittenChunkByChunk(t *testing.T) {
|
||
dsn := os.Getenv("LLMGW_TEST_POSTGRES")
|
||
if dsn == "" {
|
||
t.Skip("set LLMGW_TEST_POSTGRES to a pgvector database to run this")
|
||
}
|
||
ctx := context.Background()
|
||
s, err := store.OpenPostgres(ctx, dsn)
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
defer s.Close()
|
||
if s.KnowledgeAvailable() != nil {
|
||
t.Skip("no pgvector")
|
||
}
|
||
models := append(kbModels(), catalog.Model{Name: "writer", Source: "/w",
|
||
Profiles: []catalog.Profile{{Name: "cpu", GPUClass: "h100-80g", TP: 1}}})
|
||
a, _ := setup(t, models)
|
||
a.SetKB(&KBAdmin{Store: s})
|
||
asked := 0
|
||
a.SetReplayer(func(_ context.Context, path, model string, body []byte) (int, []byte, time.Duration, error) {
|
||
asked++
|
||
if path != "/v1/chat/completions" || model != "writer" || !strings.Contains(string(body), "Write in the passage") {
|
||
t.Errorf("asked %s %s %s", path, model, body)
|
||
}
|
||
if strings.Contains(string(body), "bozuk") {
|
||
return 200, []byte(`{"choices":[{"message":{"content":"sorry"}}]}`), 0, nil
|
||
}
|
||
return 200, []byte(`{"choices":[{"message":{"content":"{\"q\":\"Soru?\",\"a\":\"Cevap.\"}"}}]}`), 0, nil
|
||
})
|
||
org, _ := s.CreateOrg(ctx, "ds-"+strings.ToLower(store.NewID("")))
|
||
c, err := s.CreateCollection(ctx, store.KBCollection{OwnerScope: store.ScopeOrg, OwnerID: org, Name: "d", EmbedModel: "bge-m3", Dims: 3})
|
||
if err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
for _, d := range []string{"a.md", "b.md"} {
|
||
body := "iyi metin"
|
||
if d == "b.md" {
|
||
body = "bozuk metin"
|
||
}
|
||
if _, _, err := s.PutDocument(ctx, c, store.KBDocument{Name: d, Digest: d}, []store.KBChunk{{Body: body, Embedding: []float32{1, 0, 0}}}); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
}
|
||
mux := http.NewServeMux()
|
||
a.Register(mux)
|
||
rec := httptest.NewRecorder()
|
||
mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/kb/collections/"+c.ID+"/dataset", strings.NewReader(`{"model":"writer","per_chunk":3}`)))
|
||
var lines []KBDatasetLine
|
||
for _, l := range strings.Split(strings.TrimSpace(rec.Body.String()), "\n") {
|
||
var x KBDatasetLine
|
||
_ = json.Unmarshal([]byte(l), &x)
|
||
lines = append(lines, x)
|
||
}
|
||
if rec.Code != 200 || asked != 2 || len(lines) != 2 || lines[0].Q != "Soru?" || lines[0].Doc != "a.md" ||
|
||
lines[1].Doc != "b.md" || lines[1].Error == "" {
|
||
t.Errorf("%d asked %d: %+v", rec.Code, asked, lines)
|
||
}
|
||
// An embedding model writes no questions.
|
||
rec = httptest.NewRecorder()
|
||
mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/kb/collections/"+c.ID+"/dataset", strings.NewReader(`{"model":"bge-m3"}`)))
|
||
if rec.Code != http.StatusBadRequest {
|
||
t.Errorf("an embedding model: %d", rec.Code)
|
||
}
|
||
}
|