From 5dbc7e4094790b44721ddb3a53834182774adec7 Mon Sep 17 00:00:00 2001 From: Kyle Isom Date: Wed, 23 Sep 2026 03:03:11 -0700 Subject: [PATCH] p2.2: cap embed inputs at 3000 runes, retry too-large 500 Link-heavy chunks estimated under EmbedMaxTok still overflow Dixie's 2048-token embed slot. Cap every embed input at 3000 runes. A 500 whose body contains 'too large' halves inputs and retries once. Fallback logs the truncated upstream body. --- internal/digest/digest.go | 71 +++++++++++++++++++++++++++------- internal/digest/digest_test.go | 45 +++++++++++++++++++++ 2 files changed, 103 insertions(+), 13 deletions(-) diff --git a/internal/digest/digest.go b/internal/digest/digest.go index ff14fec..d97b764 100644 --- a/internal/digest/digest.go +++ b/internal/digest/digest.go @@ -31,7 +31,7 @@ const ( SmallDoc = 6000 ChunkTokens = 1200 EmbedBatch = 16 - EmbedMaxTok = 1536 + EmbedMaxRunes = 3000 DefaultPermit = 2 DefaultQDepth = 8 DefaultQWait = 20 * time.Second @@ -209,6 +209,29 @@ func (e *Engine) embedModel() string { return "honcho-embed" } +func (e *Engine) logFallback(reason string, err error) { + e.noteFallback(reason) + if e == nil || e.Log == nil { + return + } + if err != nil { + e.Log.Warn("digest fallback", "reason", reason, "err", truncateErr(err, 512)) + return + } + e.Log.Warn("digest fallback", "reason", reason) +} + +func truncateErr(err error, n int) string { + if err == nil { + return "" + } + s := err.Error() + if len(s) <= n { + return s + } + return s[:n] +} + func (e *Engine) noteFallback(reason string) { if e == nil { return @@ -241,7 +264,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL string) Re } ok, reason := e.Gate.Acquire(ctx) if !ok { - e.noteFallback(reason) + e.logFallback(reason, nil) r.Fallback = reason r.DigestMS = time.Since(start).Milliseconds() return r @@ -260,7 +283,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL string) Re body, nUsed, nTotal, err := e.selectText(dctx, cleaned, question) r.ChunksUsed, r.ChunksTotal = nUsed, nTotal if err != nil { - e.noteFallback("embed") + e.logFallback("embed", err) r.Fallback = "embed" r.DigestMS = time.Since(start).Milliseconds() return r @@ -277,13 +300,13 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL string) Re if dctx.Err() != nil { reason = "timeout" } - e.noteFallback(reason) + e.logFallback(reason, err) r.Fallback = reason return r } text = strings.TrimSpace(text) if text == "" { - e.noteFallback("empty") + e.logFallback("empty", nil) r.Fallback = "empty" return r } @@ -358,11 +381,7 @@ func (e *Engine) selectText(ctx context.Context, cleaned, question string) (stri func capEmbed(chunks []string) []string { out := make([]string, len(chunks)) for i, c := range chunks { - if EstTokens(c) > EmbedMaxTok { - out[i] = trimTokens(c, EmbedMaxTok) - } else { - out[i] = c - } + out[i] = trimRunes(c, EmbedMaxRunes) } return out } @@ -371,6 +390,7 @@ func (e *Engine) embed(ctx context.Context, inputs []string) ([][]float64, error if e.EmbedURL == "" { return nil, fmt.Errorf("no embed url") } + inputs = capEmbed(inputs) var all [][]float64 for i := 0; i < len(inputs); i += EmbedBatch { j := i + EmbedBatch @@ -387,6 +407,14 @@ func (e *Engine) embed(ctx context.Context, inputs []string) ([][]float64, error } func (e *Engine) embedBatch(ctx context.Context, inputs []string) ([][]float64, error) { + vecs, err := e.embedOnce(ctx, inputs) + if err != nil && strings.Contains(strings.ToLower(err.Error()), "too large") { + return e.embedOnce(ctx, halveInputs(inputs)) + } + return vecs, err +} + +func (e *Engine) embedOnce(ctx context.Context, inputs []string) ([][]float64, error) { payload, _ := json.Marshal(map[string]any{"input": inputs, "model": e.embedModel()}) req, err := http.NewRequestWithContext(ctx, http.MethodPost, e.EmbedURL, bytes.NewReader(payload)) if err != nil { @@ -399,7 +427,8 @@ func (e *Engine) embedBatch(ctx context.Context, inputs []string) ([][]float64, } defer resp.Body.Close() if resp.StatusCode != 200 { - return nil, fmt.Errorf("embed %d", resp.StatusCode) + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return nil, fmt.Errorf("embed %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) } var env struct { Data []struct { @@ -455,7 +484,8 @@ func (e *Engine) chat(ctx context.Context, user string) (string, error) { } defer resp.Body.Close() if resp.StatusCode != 200 { - return "", fmt.Errorf("chat %d", resp.StatusCode) + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 512)) + return "", fmt.Errorf("chat %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) } return readSSE(ctx, resp.Body) } @@ -553,14 +583,29 @@ func countBudget(chunks []string, budget int) int { func CapBudget(s string) string { return trimTokens(s, SmallDoc) } func trimTokens(s string, max int) string { + return trimRunes(s, max*4) +} + +func trimRunes(s string, n int) string { r := []rune(s) - n := max * 4 if n >= len(r) { return s } + if n < 0 { + return "" + } return string(r[:n]) } +func halveInputs(inputs []string) []string { + out := make([]string, len(inputs)) + for i, s := range inputs { + n := utf8.RuneCountInString(s) + out[i] = trimRunes(s, n/2) + } + return out +} + func cosine(a, b []float64) float64 { if len(a) == 0 || len(b) == 0 { return 0 diff --git a/internal/digest/digest_test.go b/internal/digest/digest_test.go index b4c592c..efe7bcf 100644 --- a/internal/digest/digest_test.go +++ b/internal/digest/digest_test.go @@ -11,6 +11,7 @@ import ( "sync/atomic" "testing" "time" + "unicode/utf8" ) func sseHandler(t *testing.T, reply string, saw *atomic.Pointer[chatReq], disconnected *atomic.Bool) http.Handler { @@ -289,3 +290,47 @@ func TestEmbedUsesConfiguredModel(t *testing.T) { t.Fatalf("default model %q", got) } } + +func TestEmbedTooLargeRetriesAndCompletes(t *testing.T) { + var n500 atomic.Int32 + var saw atomic.Pointer[chatReq] + var disc atomic.Bool + mux := http.NewServeMux() + mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) { + var body struct { + Input []string `json:"input"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + for _, s := range body.Input { + if utf8.RuneCountInString(s) > 1500 { + n500.Add(1) + w.WriteHeader(500) + _, _ = w.Write([]byte("input (2064 tokens) is too large to process. (current batch size: 2048)")) + return + } + } + var data []map[string]any + for i := range body.Input { + data = append(data, map[string]any{"index": i, "embedding": []float64{1, float64(i)}}) + } + _ = json.NewEncoder(w).Encode(map[string]any{"data": data}) + }) + mux.Handle("/chat", sseHandler(t, "ok from doc", &saw, &disc)) + ts := httptest.NewServer(mux) + defer ts.Close() + e := &Engine{ + ChatURL: ts.URL + "/chat", EmbedURL: ts.URL + "/embed", + Gate: NewGate(2, 8, time.Second), Client: ts.Client(), + } + cleaned := strings.Repeat("https://example.com/path ", 4000) + r := e.Run(context.Background(), cleaned, "what is this?", "https://ex.test/pypi") + if r.Fallback != "" { + t.Fatalf("fallback %s text %q", r.Fallback, r.Text) + } + if n500.Load() < 1 { + t.Fatal("want a too-large 500 then retry") + } + if !strings.Contains(r.Text, "ok") { + t.Fatalf("text %q", r.Text) + } +}