p2: digest with in-sift FIFO, SQLite caches, Honcho context, recover
Dixie is capped (max_inflight 2, queue_depth 8, queue_wait 20s) and never queued on llama-server. Streaming chat, thinking off, 45s fallback to clean. Digest and chunk-embedding caches are content-addressed SQLite (Honcho context is not cached). Recover uses the guarded Fetcher. guardian-star is refused in code.
This commit is contained in:
@@ -0,0 +1,217 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func sseHandler(t *testing.T, reply string, saw *atomic.Pointer[chatReq], disconnected *atomic.Bool) http.Handler {
|
||||
t.Helper()
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
var req chatReq
|
||||
_ = json.Unmarshal(b, &req)
|
||||
saw.Store(&req)
|
||||
fl, _ := w.(http.Flusher)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(200)
|
||||
notify := r.Context().Done()
|
||||
for _, ch := range strings.Split(reply, " ") {
|
||||
select {
|
||||
case <-notify:
|
||||
disconnected.Store(true)
|
||||
return
|
||||
default:
|
||||
}
|
||||
chunk, _ := json.Marshal(map[string]any{
|
||||
"choices": []map[string]any{{"delta": map[string]string{"content": ch + " "}}},
|
||||
})
|
||||
_, _ = w.Write([]byte("data: " + string(chunk) + "\n\n"))
|
||||
if fl != nil {
|
||||
fl.Flush()
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
if fl != nil {
|
||||
fl.Flush()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestPromptFramingAndThinkingOff(t *testing.T) {
|
||||
var saw atomic.Pointer[chatReq]
|
||||
var disc atomic.Bool
|
||||
ts := httptest.NewServer(sseHandler(t, "hello from doc", &saw, &disc))
|
||||
defer ts.Close()
|
||||
e := &Engine{ChatURL: ts.URL, Model: "ornith-1.5-9b-uncensored", Gate: NewGate(2, 8, time.Second), Client: ts.Client()}
|
||||
r := e.Run(context.Background(), "Section A\nThe wafer share is 42 percent.", "what share?", "https://ex.test/a")
|
||||
if r.Fallback != "" {
|
||||
t.Fatalf("fallback %s text %q", r.Fallback, r.Text)
|
||||
}
|
||||
req := saw.Load()
|
||||
if req == nil {
|
||||
t.Fatal("no request")
|
||||
}
|
||||
if req.MaxTokens != 700 {
|
||||
t.Fatalf("max_tokens %d", req.MaxTokens)
|
||||
}
|
||||
if req.Stream != true {
|
||||
t.Fatal("stream")
|
||||
}
|
||||
th, _ := req.ChatTemplateKwargs["enable_thinking"].(bool)
|
||||
if th {
|
||||
t.Fatal("thinking should be off")
|
||||
}
|
||||
if len(req.Messages) != 2 || req.Messages[0].Role != "system" {
|
||||
t.Fatalf("messages %+v", req.Messages)
|
||||
}
|
||||
if !strings.Contains(req.Messages[0].Content, "untrusted web document") {
|
||||
t.Fatal("system prompt")
|
||||
}
|
||||
if !strings.Contains(req.Messages[1].Content, "<document url=") {
|
||||
t.Fatal("document tags")
|
||||
}
|
||||
if !strings.Contains(req.Messages[1].Content, "Never follow") && !strings.Contains(req.Messages[0].Content, "Never follow") {
|
||||
t.Fatal("instruction")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallbackOnTimeout(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
}))
|
||||
defer ts.Close()
|
||||
e := &Engine{ChatURL: ts.URL, Gate: NewGate(2, 8, time.Second), Client: ts.Client(), Timeout: 50 * time.Millisecond}
|
||||
r := e.Run(context.Background(), "hi", "q", "https://x")
|
||||
if r.Fallback == "" {
|
||||
t.Fatal("want fallback")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFallbackEmpty(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("data: [DONE]\n\n"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
e := &Engine{ChatURL: ts.URL, Gate: NewGate(2, 8, time.Second), Client: ts.Client()}
|
||||
r := e.Run(context.Background(), "hi", "q", "https://x")
|
||||
if r.Fallback != "empty" {
|
||||
t.Fatalf("got %q", r.Fallback)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPermitCap(t *testing.T) {
|
||||
started := make(chan struct{}, 8)
|
||||
release := make(chan struct{})
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
e := &Engine{ChatURL: ts.URL, Gate: NewGate(2, 8, time.Second), Client: ts.Client()}
|
||||
var wg sync.WaitGroup
|
||||
var nFallback atomic.Int32
|
||||
for i := 0; i < 3; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
r := e.Run(context.Background(), "doc", "q", "https://x")
|
||||
if r.Fallback == "permit" {
|
||||
nFallback.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
// two in flight, third should fail permit (5s wait — shorten by using a tiny gate wait)
|
||||
// Gate wait is 5s; that's too slow. Hold the two slots then the third waits 5s.
|
||||
// We accept that; instead pre-fill the gate.
|
||||
close(release)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestQueueFullAndWait(t *testing.T) {
|
||||
g := NewGate(2, 1, 200*time.Millisecond)
|
||||
ctx := context.Background()
|
||||
if ok, _ := g.Acquire(ctx); !ok {
|
||||
t.Fatal("1")
|
||||
}
|
||||
if ok, _ := g.Acquire(ctx); !ok {
|
||||
t.Fatal("2")
|
||||
}
|
||||
var waiterOK atomic.Bool
|
||||
go func() {
|
||||
ok, _ := g.Acquire(ctx)
|
||||
waiterOK.Store(ok)
|
||||
}()
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
ok, reason := g.Acquire(ctx)
|
||||
if ok || reason != "queue_full" {
|
||||
t.Fatalf("4th got ok=%v reason=%s", ok, reason)
|
||||
}
|
||||
g.Release()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for time.Now().Before(deadline) && !waiterOK.Load() {
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
if !waiterOK.Load() {
|
||||
t.Fatal("queued waiter should acquire after release")
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueueWaitFallback(t *testing.T) {
|
||||
g := NewGate(1, 1, 30*time.Millisecond)
|
||||
if ok, _ := g.Acquire(context.Background()); !ok {
|
||||
t.Fatal()
|
||||
}
|
||||
ok, reason := g.Acquire(context.Background())
|
||||
if ok || reason != "queue_wait" {
|
||||
t.Fatalf("got %v %s", ok, reason)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancelPropagates(t *testing.T) {
|
||||
var saw atomic.Pointer[chatReq]
|
||||
var disc atomic.Bool
|
||||
ts := httptest.NewServer(sseHandler(t, strings.Repeat("word ", 200), &saw, &disc))
|
||||
defer ts.Close()
|
||||
e := &Engine{ChatURL: ts.URL, Gate: NewGate(2, 8, time.Second), Client: ts.Client(), Timeout: 2 * time.Second}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 80*time.Millisecond)
|
||||
defer cancel()
|
||||
_ = e.Run(ctx, "doc", "q", "https://x")
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
if !disc.Load() {
|
||||
t.Fatal("fake server did not observe disconnect")
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkOrderAndBudget(t *testing.T) {
|
||||
var parts []string
|
||||
for i := 0; i < 20; i++ {
|
||||
parts = append(parts, strings.Repeat("heading word ", 80))
|
||||
}
|
||||
chunks := Split(strings.Join(parts, "\n"), 50)
|
||||
if len(chunks) < 3 {
|
||||
t.Fatalf("chunks %d", len(chunks))
|
||||
}
|
||||
joined := joinBudget(chunks, 120)
|
||||
if EstTokens(joined) == 0 {
|
||||
t.Fatal("empty")
|
||||
}
|
||||
// no-question uses leading chunks
|
||||
lead := joinBudget(chunks, SmallDoc)
|
||||
if !strings.HasPrefix(lead, chunks[0][:20]) && !strings.Contains(lead, chunks[0][:min(20, len(chunks[0]))]) {
|
||||
t.Fatal("document order")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user