From 68e8b5108c7d9b2989a70ea2aeee177389bb4044 Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Sun, 12 Jul 2026 18:44:15 +0300 Subject: [PATCH 1/5] feat: add LLM-based context components (summarize + extract code strategy) Land the two model-calling components, general + configurable, largely restored from git history. - internal/cheapmodel: Anthropic+OpenAI Complete(ctx,prompt) clients (stdlib). - components.Model + ModelSpec{Incoming,Static}.For(source); Ctx.Model populated by apply.BodyWithModel. proxy builds the per-request 'incoming' client (route upstream + gateway key + request model); cmd builds a static 'config' cheap model from CHEAP_MODEL* env; AuthBridge offers config-only. - summarize (CE-Manager port): whole-transcript summary -> [msg0, , last-K], replaced span stashed for expand. apply gains a lossless count-change rebuild (retained messages kept byte-identical). - extract: strategy deterministic|code|rlm; code runs an LLM-written Starlark filter in a sandbox (no imports/IO, step+2s limits) accepted only if a containment check proves a lossless subset, else falls back to deterministic; content-hash cache. Starlark chosen over sandboxed Python (safer, pure-Go). - NeedsModel components degrade gracefully when no model is available. Tests for summarize shape/stash/no-model, extract code + fallback, and the apply count-change lossless invariant. Live-verified: extract code had claude-sonnet-4-6 write a Starlark filter that reduced a tool output 65%. Assisted-By: Claude (Anthropic AI) Signed-off-by: Osher-Elhadad --- adapters/bifrost/plugin.go | 5 + apply/apply.go | 92 ++++++- apply/apply_test.go | 51 ++++ cmd/context-guru-proxy/main.go | 32 +++ components/all/llm_test.go | 143 ++++++++++ components/component.go | 37 ++- components/offload/extract.go | 62 +++-- components/offload/summarize.go | 210 +++++++++++++++ config/config.go | 3 + docs/components.md | 40 ++- docs/design.md | 23 ++ go.mod | 2 +- go.sum | 4 + internal/cheapmodel/anthropic.go | 87 ++++++ internal/cheapmodel/cheapmodel_test.go | 75 ++++++ internal/cheapmodel/openai.go | 70 +++++ internal/extract/cache.go | 27 ++ internal/extract/contain.go | 101 +++++++ internal/extract/deterministic.go | 180 +++++++++++++ internal/extract/extract.go | 353 +++++++++++++++++++++++++ internal/extract/prompt.go | 147 ++++++++++ internal/extract/starlark.go | 63 +++++ internal/extract/starlark_test.go | 58 ++++ proxy/proxy.go | 57 +++- 24 files changed, 1885 insertions(+), 37 deletions(-) create mode 100644 components/all/llm_test.go create mode 100644 components/offload/summarize.go create mode 100644 internal/cheapmodel/anthropic.go create mode 100644 internal/cheapmodel/cheapmodel_test.go create mode 100644 internal/cheapmodel/openai.go create mode 100644 internal/extract/cache.go create mode 100644 internal/extract/contain.go create mode 100644 internal/extract/deterministic.go create mode 100644 internal/extract/extract.go create mode 100644 internal/extract/prompt.go create mode 100644 internal/extract/starlark.go create mode 100644 internal/extract/starlark_test.go diff --git a/adapters/bifrost/plugin.go b/adapters/bifrost/plugin.go index e6556b2..f1d41b6 100644 --- a/adapters/bifrost/plugin.go +++ b/adapters/bifrost/plugin.go @@ -26,6 +26,10 @@ const SessionContextKey bschemas.BifrostContextKey = "context-guru-session" type Plugin struct { pipe *components.Pipeline store store.Store + // CheapModel is the static "config"-source LLM for NeedsModel components. The + // bifrost host has no usable "incoming" client (the agent's key is a + // placeholder), so only the static source is offered here. Optional. + CheapModel components.Model } // New builds the adapter from an already-constructed pipeline and store. @@ -48,6 +52,7 @@ func (p *Plugin) PreRequestHook(ctx *bschemas.BifrostContext, req *bschemas.Bifr Ctx: ctx, Session: p.resolveSession(ctx, chat), Store: p.store, + Model: components.ModelSpec{Static: p.CheapModel}, Bypass: bypassed(ctx), } p.pipe.Run(chat, c) diff --git a/apply/apply.go b/apply/apply.go index a3f2cb0..04c2409 100644 --- a/apply/apply.go +++ b/apply/apply.go @@ -83,11 +83,18 @@ type slot struct { lossless bool // wholeMessage: does bifrost round-trip this message without dropping fields } -// Body runs the pipeline over the request body's messages and returns the -// rewritten body. changed=false means "forward the original unchanged" (no +// Body runs the pipeline with no LLM clients available (deterministic components +// only). See BodyWithModel to supply model clients for LLM-based components. +func Body(ctx context.Context, pipe *components.Pipeline, st store.Store, provider bschemas.ModelProvider, body []byte, explicitSession string, bypass bool) ([]byte, bool) { + return BodyWithModel(ctx, pipe, st, provider, body, explicitSession, bypass, components.ModelSpec{}) +} + +// BodyWithModel runs the pipeline over the request body's messages and returns +// the rewritten body. changed=false means "forward the original unchanged" (no // messages array, unparseable, or a re-serialization problem) — always fail // open. explicitSession is the host-supplied session id ("" -> content hash). -func Body(ctx context.Context, pipe *components.Pipeline, st store.Store, provider bschemas.ModelProvider, body []byte, explicitSession string, bypass bool) ([]byte, bool) { +// models carries the LLM clients that NeedsModel components may call. +func BodyWithModel(ctx context.Context, pipe *components.Pipeline, st store.Store, provider bschemas.ModelProvider, body []byte, explicitSession string, bypass bool, models components.ModelSpec) ([]byte, bool) { msgsRaw := gjson.GetBytes(body, "messages") if !msgsRaw.Exists() || !msgsRaw.IsArray() { return body, false @@ -107,15 +114,25 @@ func Body(ctx context.Context, pipe *components.Pipeline, st store.Store, provid Ctx: ctx, Session: session.Resolve(explicitSession, sys, firstUser), Store: st, + Model: models, Bypass: bypass, } + // Canonical form of each normalized message BEFORE the pipeline, so a + // count-changing component (summarize) can be mapped back to the body. + normPre := make([][]byte, len(norm)) + for i := range norm { + normPre[i], _ = json.Marshal(norm[i]) + } + pipe.Run(chat, c) - // A component changed the message count (none of the v1 set does) — the slot - // map no longer aligns, so fail open and forward the original untouched. + // A component changed the message count (summarize restructures the transcript + // to [msg0, , last-K]). Rebuild the messages array preserving each + // retained message's ORIGINAL raw bytes (byte-lossless, incl. Anthropic + // tool_result) and marshaling only genuinely new messages (the summary). if len(chat.Input) != len(norm) { - return body, false + return rebuildCountChanged(body, msgsRaw.Array(), normPre, slots, chat.Input) } out := body @@ -290,6 +307,69 @@ func jsonEqual(a, b []byte) bool { return reflect.DeepEqual(av, bv) } +// rebuildCountChanged reconstructs the messages array after a component changed +// the message count. Each output message that byte-matches a pre-pipeline +// normalized message (a survivor) is emitted as its ORIGINAL body raw bytes +// (byte-lossless); genuinely new messages (the summary) are marshaled fresh. +// Fail-open (returns body,false) if any survivor can't be mapped to the body. +func rebuildCountChanged(body []byte, orig []gjson.Result, normPre [][]byte, slots []slot, out []bschemas.ChatMessage) ([]byte, bool) { + used := make([]bool, len(normPre)) + var parts [][]byte + lastBodyIdx := -1 + for i := range out { + mb, err := json.Marshal(out[i]) + if err != nil { + return body, false + } + matched := -1 + for k := range normPre { + if !used[k] && bytes.Equal(mb, normPre[k]) { + matched = k + break + } + } + if matched < 0 { + parts = append(parts, mb) // new message (e.g. the summary) — fresh, lossless (plain text) + lastBodyIdx = -1 + continue + } + used[matched] = true + bi, ok := bodyIndexOf(slots[matched].path) + if !ok || bi < 0 || bi >= len(orig) { + return body, false + } + if bi == lastBodyIdx { + continue // several normalized messages share one body message — emit it once + } + parts = append(parts, []byte(orig[bi].Raw)) + lastBodyIdx = bi + } + var buf bytes.Buffer + buf.WriteByte('[') + for i, p := range parts { + if i > 0 { + buf.WriteByte(',') + } + buf.Write(p) + } + buf.WriteByte(']') + res, err := sjson.SetRawBytes(body, "messages", buf.Bytes()) + if err != nil { + return body, false + } + return res, true +} + +// bodyIndexOf extracts the leading messages. index from a slot path. +func bodyIndexOf(path string) (int, bool) { + s := strings.TrimPrefix(path, "messages.") + if dot := strings.IndexByte(s, '.'); dot >= 0 { + s = s[:dot] + } + i, err := strconv.Atoi(s) + return i, err == nil +} + func systemAndFirstUser(msgs []bschemas.ChatMessage) (sys, firstUser string) { for _, m := range msgs { t := schema.MessageText(m) diff --git a/apply/apply_test.go b/apply/apply_test.go index 4705cec..b8d7d08 100644 --- a/apply/apply_test.go +++ b/apply/apply_test.go @@ -7,6 +7,7 @@ import ( "testing" "github.com/kagenti/context-guru/apply" + "github.com/kagenti/context-guru/components" _ "github.com/kagenti/context-guru/components/all" "github.com/kagenti/context-guru/config" "github.com/kagenti/context-guru/store" @@ -206,6 +207,56 @@ func TestAnthropicStructuredToolResultUntouched(t *testing.T) { } } +type stubModel struct{ resp string } + +func (m stubModel) Complete(context.Context, string) (string, error) { return m.resp, nil } + +// TestSummarizeCountChangeLossless: summarize restructures [system,u1,tool,final] +// into [system, , final]; apply must keep the retained messages and all +// non-message fields byte-identical while the count drops. +func TestSummarizeCountChangeLossless(t *testing.T) { + cfg := pipe(t, "pipeline: [summarize]\ncomponents:\n summarize: {keep_last: 1, start_from_message: 0, min_tokens: 1}\n") + p, _ := cfg.Build(nil) + st := store.NewMemory(store.Options{}) + + body, _ := json.Marshal(map[string]any{ + "model": "gpt-x", + "temperature": 0.3, + "messages": []map[string]any{ + {"role": "system", "content": "you are helpful"}, + {"role": "user", "content": "do the task"}, + {"role": "tool", "tool_call_id": "a", "content": strings.Repeat("verbose tool output\n", 50)}, + {"role": "user", "content": "the final question"}, + }, + }) + + out, changed := apply.BodyWithModel(context.Background(), p, st, bschemas.OpenAI, body, "", false, + components.ModelSpec{Incoming: stubModel{resp: "essential facts"}}) + if !changed { + t.Fatal("summarize should have restructured the transcript") + } + if n := gjson.GetBytes(out, "messages.#").Int(); n != 3 { + t.Fatalf("expected 3 messages after summarize, got %d: %s", n, out) + } + // Non-message fields byte-identical. + for _, path := range []string{"model", "temperature"} { + if gjson.GetBytes(out, path).Raw != gjson.GetBytes(body, path).Raw { + t.Fatalf("field %q not preserved", path) + } + } + // Retained messages (system msg0, final user msg) byte-identical to originals. + if gjson.GetBytes(out, "messages.0").Raw != gjson.GetBytes(body, "messages.0").Raw { + t.Fatal("msg0 must be byte-identical") + } + if gjson.GetBytes(out, "messages.2").Raw != gjson.GetBytes(body, "messages.3").Raw { + t.Fatalf("the final message must be preserved verbatim: %s", gjson.GetBytes(out, "messages.2").Raw) + } + // The inserted summary carries the marker for expand recovery. + if s := gjson.GetBytes(out, "messages.1.content").String(); !strings.Contains(s, "History Summary") || !strings.Contains(s, "< no client +// CHEAP_MODEL_PROVIDER anthropic (default) | openai +// CHEAP_MODEL_BASE upstream base URL (default: the matching provider default) +// CHEAP_MODEL_KEY API key (default: ANTHROPIC_API_KEY / OPENAI_API_KEY) +// CHEAP_MODEL_AUTH anthropic auth scheme: x-api-key (default) | bearer +func cheapModelFromEnv() components.Model { + model := os.Getenv("CHEAP_MODEL") + if model == "" { + return nil + } + switch envOr("CHEAP_MODEL_PROVIDER", "anthropic") { + case "openai": + return cheapmodel.OpenAI{ + BaseURL: os.Getenv("CHEAP_MODEL_BASE"), Model: model, + APIKey: envOr("CHEAP_MODEL_KEY", os.Getenv("OPENAI_API_KEY")), + } + default: + return cheapmodel.Anthropic{ + BaseURL: os.Getenv("CHEAP_MODEL_BASE"), Model: model, + APIKey: envOr("CHEAP_MODEL_KEY", os.Getenv("ANTHROPIC_API_KEY")), + AuthScheme: os.Getenv("CHEAP_MODEL_AUTH"), + } + } +} diff --git a/components/all/llm_test.go b/components/all/llm_test.go new file mode 100644 index 0000000..7469c72 --- /dev/null +++ b/components/all/llm_test.go @@ -0,0 +1,143 @@ +package all_test + +import ( + "context" + "strings" + "testing" + + "github.com/kagenti/context-guru/components" + "github.com/kagenti/context-guru/schema" + "github.com/kagenti/context-guru/store" + bschemas "github.com/maximhq/bifrost/core/schemas" +) + +// stubModel is a fixed LLM used to drive the model-based components in tests. +type stubModel struct { + resp string + err error +} + +func (m stubModel) Complete(context.Context, string) (string, error) { return m.resp, m.err } + +func strp(s string) *string { return &s } + +func newComp(t *testing.T, name, yaml string) components.Offload { + t.Helper() + c, err := components.New(name, []byte(yaml)) + if err != nil { + t.Fatalf("New(%s): %v", name, err) + } + off, ok := c.(components.Offload) + if !ok { + t.Fatalf("%s is not an Offload", name) + } + return off +} + +// TestSummarizeRestructures: [system,u1,tool,u2] with keep_last=1 becomes +// [system, , u2]; the summary carries a marker and the replaced span is +// stashed for expand. +func TestSummarizeRestructures(t *testing.T) { + off := newComp(t, "summarize", "keep_last: 1\nstart_from_message: 0\nmin_tokens: 1\n") + st := store.NewMemory(store.Options{}) + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + {Role: bschemas.ChatMessageRoleSystem, Content: &bschemas.ChatMessageContent{ContentStr: strp("you are helpful")}}, + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("do the task")}}, + toolMsg(strings.Repeat("verbose tool output line\n", 40)), + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("continue")}}, + }} + c := &components.Ctx{Ctx: context.Background(), Store: st, Model: components.ModelSpec{Incoming: stubModel{resp: "essential facts"}}} + var rep components.Report + keys, err := off.Offload(req, &rep, c) + if err != nil { + t.Fatal(err) + } + if len(keys) != 1 { + t.Fatalf("expected 1 stashed span, got %d (skipped=%v)", len(keys), rep.Skipped) + } + if len(req.Input) != 3 { + t.Fatalf("expected [system, summary, u2], got %d messages", len(req.Input)) + } + if req.Input[0].Role != bschemas.ChatMessageRoleSystem || schema.MessageText(req.Input[0]) != "you are helpful" { + t.Fatal("msg0 must be preserved") + } + sm := schema.MessageText(req.Input[1]) + if !strings.Contains(sm, "History Summary") || !strings.Contains(sm, "< deterministic. +func TestExtractCodeNilModelFallsBack(t *testing.T) { + off := newComp(t, "extract", "strategy: code\nmin_tokens: 1\nhead_lines: 1\ntail_lines: 1\nmodel:\n source: config\n") + st := store.NewMemory(store.Options{}) + lines := make([]string, 30) + for i := range lines { + lines[i] = "irrelevant filler line" + } + lines[15] = "the keep marker line" + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("find keep")}}, + toolMsg(strings.Join(lines, "\n")), + }} + c := &components.Ctx{Ctx: context.Background(), Store: st} // no model -> deterministic + var rep components.Report + if _, err := off.Offload(req, &rep, c); err != nil { + t.Fatal(err) + } + // deterministic projection still runs and marks the message. + if !strings.Contains(schema.MessageText(req.Input[1]), "< +// the static cheap model; anything else ("incoming"/unset) -> the incoming model, +// falling back to the static one when there is no incoming client. Returns nil +// when nothing is available, and the caller must degrade gracefully. +func (m ModelSpec) For(source string) Model { + if source == "config" { + return m.Static + } + if m.Incoming != nil { + return m.Incoming + } + return m.Static } // Ctx is the per-request runtime handed to every component. diff --git a/components/offload/extract.go b/components/offload/extract.go index 2da0901..d8abc7b 100644 --- a/components/offload/extract.go +++ b/components/offload/extract.go @@ -5,6 +5,7 @@ import ( "github.com/kagenti/context-guru/components" "github.com/kagenti/context-guru/expand" + "github.com/kagenti/context-guru/internal/extract" "github.com/kagenti/context-guru/schema" bschemas "github.com/maximhq/bifrost/core/schemas" "gopkg.in/yaml.v3" @@ -13,17 +14,23 @@ import ( func init() { components.Register("extract", newExtract) } // Extract projects a large tool output down to the part relevant to the current -// query, stashing the full original. v1 ships the deterministic strategy: -// keep lines matching the recent query's keywords (plus head/tail context and -// any error lines). The LLM strategies (code = model-generated Starlark filter -// with a containment validator; rlm = recursive chunked) are the winnow-parity -// refinement — they need a cheap-model client (ModelSpec), which is not yet -// wired; selecting them falls back to deterministic with a note. +// query, stashing the full original. Three strategies: +// - deterministic (default, no LLM): keep query-keyword + head/tail + error lines. +// - code: a cheap LLM writes a Starlark `extract_relevant_data` filter, run in a +// sandbox (no imports/IO, step+time limited); the output is accepted only if a +// containment + sanity check proves it's a lossless projection, else it falls +// back to deterministic. (winnow's llm_compact, ported.) +// - rlm: reserved for very large outputs; currently maps to code. +// +// The LLM strategies need a Model (Ctx.Model, per model.source); when none is +// available they degrade to deterministic. The full original is always stashed +// under a marker, so any reduction is reversible via the expand tool. type Extract struct { - minTokens int - head int - tail int - strategy string + minTokens int + head int + tail int + strategy string + modelSource string } type extractConfig struct { @@ -31,6 +38,9 @@ type extractConfig struct { Head int `yaml:"head_lines"` Tail int `yaml:"tail_lines"` Strategy string `yaml:"strategy"` // deterministic | code | rlm + Model struct { + Source string `yaml:"source"` // incoming (default) | config + } `yaml:"model"` } func newExtract(raw []byte) (components.Component, error) { @@ -40,23 +50,27 @@ func newExtract(raw []byte) (components.Component, error) { return nil, err } } - return &Extract{minTokens: cfg.MinTokens, head: cfg.Head, tail: cfg.Tail, strategy: cfg.Strategy}, nil + return &Extract{minTokens: cfg.MinTokens, head: cfg.Head, tail: cfg.Tail, strategy: cfg.Strategy, modelSource: cfg.Model.Source}, nil } func (Extract) Name() string { return "extract" } func (Extract) Enabled(*components.Ctx) bool { return true } -// NeedsModel reports whether the configured strategy calls an LLM. The pipeline -// injects the ModelSpec when true. (The LLM path is not yet implemented; see the -// package note — this reports intent for when the cheap-model client lands.) +// NeedsModel reports whether the configured strategy calls an LLM. func (e *Extract) NeedsModel() bool { return e.strategy == "code" || e.strategy == "rlm" } func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Report, c *components.Ctx) ([]string, error) { - query := keywords(lastUserText(req)) + goal := lastUserText(req) + query := keywords(goal) if len(query) == 0 { rep.Skipped = true return nil, nil // nothing to condition relevance on } + var model components.Model + if e.NeedsModel() { + model = c.Model.For(e.modelSource) + } + keepIDs := extract.HarvestIdentifiers(goal, 40) var keys []string for _, i := range toolIndices(req) { msg := &req.Input[i] @@ -70,7 +84,7 @@ func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Repo if len(expand.ParseMarkers(content)) > 0 { continue } - projected, ok := e.project(content, query) + projected, ok := e.reduce(c, content, goal, keepIDs, query, model) if !ok || schema.TextTokens(projected) >= schema.TextTokens(content) { continue } @@ -85,6 +99,22 @@ func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Repo return keys, nil } +// reduce picks the strategy: for code/rlm with a model, run the sandboxed +// LLM-generated filter (containment-validated inside RunExtraction, with its own +// deterministic fallback); otherwise the deterministic line projection. Always +// fail-open — any miss falls through to project(). +func (e *Extract) reduce(c *components.Ctx, content, goal string, keepIDs []string, query map[string]struct{}, model components.Model) (string, bool) { + if model != nil && e.NeedsModel() { + cfg := extract.DefaultCfg() + cfg.Mode = e.strategy + cfg.Floor = e.minTokens + if res, _ := extract.RunExtraction(c.Ctx, content, goal, keepIDs, schema.TextTokens(content), cfg, model); res != "" && res != content { + return res, true + } + } + return e.project(content, query) +} + // project keeps head/tail context plus every line relevant to the query or // carrying an error, collapsing the runs between with an ellipsis marker. func (e *Extract) project(content string, query map[string]struct{}) (string, bool) { diff --git a/components/offload/summarize.go b/components/offload/summarize.go new file mode 100644 index 0000000..4a4d8b1 --- /dev/null +++ b/components/offload/summarize.go @@ -0,0 +1,210 @@ +package offload + +import ( + "context" + "encoding/json" + "strings" + + "github.com/kagenti/context-guru/components" + "github.com/kagenti/context-guru/expand" + "github.com/kagenti/context-guru/schema" + bschemas "github.com/maximhq/bifrost/core/schemas" + "gopkg.in/yaml.v3" +) + +func init() { components.Register("summarize", newSummarize) } + +// Summarize compresses the middle of a long trajectory into one LLM-written +// summary (ported from CE-Manager's ReSum-style summarizer). It restructures the +// message list to [msg0, , last-K messages], replacing +// everything in between. It is an Offload: the replaced span is stashed under a +// <> marker (carried in the summary message) so the expand tool can +// restore it. NeedsModel — it no-ops when no model is available. +// +// This is the one component that changes the message count; apply.Body rebuilds +// the body preserving the retained messages' original bytes. +type Summarize struct { + level string + keepLast int + startFrom int + minTokens int + includeToolCalls bool + modelSource string +} + +type summarizeConfig struct { + SummaryLevel string `yaml:"summary_level"` // concise | regular | highly_detailed + KeepLast int `yaml:"keep_last"` // messages kept verbatim at the tail + StartFrom int `yaml:"start_from_message"` // no-op until the list reaches this length + MinTokens int `yaml:"min_tokens"` // min content tokens in the span to bother + IncludeToolCalls bool `yaml:"include_tool_calls"` + Model struct { + Source string `yaml:"source"` // incoming (default) | config + } `yaml:"model"` +} + +func newSummarize(raw []byte) (components.Component, error) { + cfg := summarizeConfig{SummaryLevel: "regular", KeepLast: 3, StartFrom: 6, MinTokens: 500} + if len(raw) > 0 { + if err := yaml.Unmarshal(raw, &cfg); err != nil { + return nil, err + } + } + return &Summarize{ + level: cfg.SummaryLevel, keepLast: cfg.KeepLast, startFrom: cfg.StartFrom, + minTokens: cfg.MinTokens, includeToolCalls: cfg.IncludeToolCalls, modelSource: cfg.Model.Source, + }, nil +} + +func (Summarize) Name() string { return "summarize" } +func (Summarize) Enabled(*components.Ctx) bool { return true } +func (*Summarize) NeedsModel() bool { return true } + +func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Report, c *components.Ctx) ([]string, error) { + msgs := req.Input + // Keep msg0 (system/first) + the last keepLast; summarize the span between. + start, end := 1, len(msgs)-s.keepLast + if len(msgs) < s.startFrom || end <= start { + rep.Skipped = true + return nil, nil + } + model := c.Model.For(s.modelSource) + if model == nil { + rep.Skipped = true // NeedsModel but none available → degrade gracefully + return nil, nil + } + span := msgs[start:end] + if schema.MessagesTokens(&bschemas.BifrostChatRequest{Input: span}) < s.minTokens { + rep.Skipped = true + return nil, nil + } + + summary, err := s.summarize(c.Ctx, model, span) + if err != nil { + return nil, err // fail-open: the pipeline reverts this component + } + if strings.TrimSpace(summary) == "" { + rep.Skipped = true + return nil, nil + } + + // Stash the replaced span so expand can restore it. + spanJSON, err := json.Marshal(span) + if err != nil { + return nil, err + } + key := hashKey(string(spanJSON)) + c.Store.Put(key, spanJSON) + + summaryMsg := bschemas.ChatMessage{Role: bschemas.ChatMessageRoleSystem} + schema.SetMessageText(&summaryMsg, summaryWrapper(summary, key)) + + // [msg0, summary, last-K] — reassign; apply.Body rebuilds losslessly. + out := make([]bschemas.ChatMessage, 0, 2+s.keepLast) + out = append(out, msgs[0], summaryMsg) + out = append(out, msgs[end:]...) + req.Input = out + return []string{key}, nil +} + +// summarize builds the trajectory string and asks the model once (bounded retry). +func (s *Summarize) summarize(ctx context.Context, model components.Model, span []bschemas.ChatMessage) (string, error) { + sys := summarizerSystemPrompt + if !s.includeToolCalls { + sys += summarizerMaskedNote + } + user := strings.Replace(summarizerUserPrompt, "{trajectory}", trajectoryString(span, s.includeToolCalls), 1) + if suffix := summaryLevelSuffix[s.level]; suffix != "" { + user += "\n" + suffix + } + prompt := sys + "\n\n" + user + var lastErr error + for i := 0; i < 3; i++ { + out, err := model.Complete(ctx, prompt) + if err != nil { + lastErr = err + continue + } + return ensureSummaryTags(out), nil + } + return "", lastErr +} + +// trajectoryString renders the span as "[role]\n{content}" blocks. When tool +// calls are excluded, tool-role content is replaced by a placeholder. +func trajectoryString(span []bschemas.ChatMessage, includeToolCalls bool) string { + var b strings.Builder + for i, m := range span { + if i > 0 { + b.WriteString("\n\n") + } + content := schema.MessageText(m) + if !includeToolCalls && m.Role == bschemas.ChatMessageRoleTool { + content = "" + } + b.WriteString("[") + b.WriteString(string(m.Role)) + b.WriteString("]\n") + b.WriteString(content) + } + return b.String() +} + +func ensureSummaryTags(s string) string { + s = strings.TrimSpace(s) + if s == "" { + return s + } + if !strings.Contains(s, "") { + s = "\n" + s + } + if !strings.Contains(s, "") { + s = s + "\n" + } + return s +} + +// summaryWrapper is the synthetic system message that replaces the span; it +// carries the marker so the expand tool can recover the full original trajectory. +func summaryWrapper(summary, key string) string { + return "=== History Summary ===\n" + + "The earlier trajectory is summarized below.\n\n" + + summary + "\n\n" + + "Use this summary as the older context, and use the following messages as the most recent context. " + + "Continue the task accordingly. Do not summarize the conversation again.\n" + + expand.Marker(key) + " [full earlier trajectory: call " + expand.ToolName + "]" +} + +// Prompts ported verbatim from CE-Manager (src/ce_manager/prompts/summarizer.py). +const summarizerSystemPrompt = "You analyze long agent trajectories with tool calls and produce compact, factual summaries. Do not guess or invent information." + +const summarizerMaskedNote = " Notice, all the tool calls content have been removed from the trajectory, you should base your summary only on the remaining content." + +const summarizerUserPrompt = `You are an expert at analyzing conversation history and extracting relevant information. Your task is to thoroughly evaluate the conversation history and current question to provide a comprehensive summary that will help answer the question. + +Task Guidelines: +1. Information Analysis +• Carefully analyze the conversation history to identify truly useful information. +• Focus on information that directly contributes to answering the question. +• Do NOT make assumptions, guesses, or inferences beyond what is explicitly stated in the conversation. +• If information is missing or unclear, do NOT include it in your summary. + +2. Summary Requirements +• Extract only the most relevant information that is explicitly present in the conversation. +• Synthesize information from multiple exchanges when relevant. Only include information that is certain and clearly stated in the conversation. +• Do NOT output or mention any information that is uncertain, insufficient, or cannot be confirmed from the conversation. + +3. Output Format Your response should be structured as follows: + +• Essential Information: [Organize the relevant and certain information from the conversation history that helps address the question.] + + +Strictly avoid fabricating, inferring, or exaggerating any information not present in the conversation. Only output information that is certain and explicitly stated. + +Trajectory: {trajectory}` + +var summaryLevelSuffix = map[string]string{ + "concise": "Please generate a concise summary", + "regular": "Please generate a comprehensive and useful summary", + "highly_detailed": "Please generate a highly detailed, fully comprehensive, explicitly grounded summary that includes every relevant and certain piece of information from the conversation.", +} diff --git a/config/config.go b/config/config.go index 0826713..2094a75 100644 --- a/config/config.go +++ b/config/config.go @@ -78,6 +78,9 @@ var presets = map[string][]string{ "aggressive": {"format", "dedup", "failed_run", "cmdfilter", "smartcrush", "extract", "cacheinject"}, "coding": {"format", "skeleton", "cmdfilter", "cacheinject"}, "mcp": {"format", "smartcrush", "cacheinject"}, + // summarize restructures the whole transcript (changes the message count) — run + // it alone so no other component's in-place edits race apply's rebuild. + "summarize": {"summarize"}, } // Build constructs the ordered pipeline from the config, wiring each named diff --git a/docs/components.md b/docs/components.md index abc2d40..7a67f52 100644 --- a/docs/components.md +++ b/docs/components.md @@ -16,14 +16,21 @@ messages (`role:"tool"`; for Anthropic, `tool_result` blocks normalized to that | `collapse` | Offload | middle of an oversized output | via expand | any large tool output (fallback) | `max_tokens` (2000), `head_lines` (20), `tail_lines` (20) | | `failed_run` | Offload | earlier superseded test/build runs | via expand | ≥2 run-like outputs | `min_tokens` (100) | | `cmdfilter` | Offload | lines per declarative DSL filter | via expand | output matching a filter | `filters` ([]), `disable_builtins` (false) | -| `extract` | Offload | query-irrelevant lines | via expand | large output + a recent query | `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` (deterministic) | +| `extract` | Offload | query-irrelevant lines (or an LLM-written filter) | via expand | large output + a recent query | `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` (deterministic\|code\|rlm), `model.source` | | `smartcrush` | Offload | middle items of a JSON array | via expand | JSON-array tool output | `min_items` (5), `min_tokens` (200), `keep_first` (3), `keep_last` (2) | | `mask` | Offload | older tool outputs (age-based) | via expand | more than `keep_recent` outputs | `keep_recent` (3), `min_tokens` (100) | | `phi_evict` | Offload | lowest-scoring outputs over budget | via expand | transcript over `budget_tokens` | `budget_tokens` (120000), `weights` (balanced) | +| `summarize` | Offload (LLM) | the middle of the transcript → one summary | via expand | long trajectories | `summary_level` (regular), `keep_last` (3), `start_from_message` (6), `min_tokens` (500), `model.source` | Presets (`config`): `off` `[]` · `safe` `[format, cacheinject]` · `balanced` `[format, dedup, failed_run, cmdfilter, cacheinject]` · `aggressive` adds `smartcrush, extract` · -`coding` `[format, skeleton, cmdfilter, cacheinject]` · `mcp` `[format, smartcrush, cacheinject]`. +`coding` `[format, skeleton, cmdfilter, cacheinject]` · `mcp` `[format, smartcrush, cacheinject]` · +`summarize` `[summarize]` (run alone — it restructures the whole transcript). + +**LLM-based components** (`extract` with `strategy: code`/`rlm`, and `summarize`) call a model, chosen by +`model.source`: `incoming` (default — reuse the proxied request's own model + key) or `config` (a dedicated +cheap model set via `CHEAP_MODEL*` env / the gateway's `CheapModel`). When no model is available they +degrade — `extract` to its deterministic projection, `summarize` to a no-op. See [design.md](design.md#llm-components). Common gates every Offload respects: skip non-text (`Rewritable`) messages, skip content already carrying a marker (no double-offload), and skip if the rewrite (marker + hint included) isn't @@ -156,10 +163,13 @@ after: … lines mentioning auth/timeout + any error lines … ``` - **Config:** `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` - (`deterministic` | `code` | `rlm`). `NeedsModel()` is true for `code`/`rlm`, but the cheap-LLM - client isn't wired yet — those select-and-fall-back to deterministic. **Shines:** big - query-focused MCP/API outputs. **Inert:** no user query to score against, output < `min_tokens`, - projection not smaller. + (`deterministic` | `code` | `rlm`), `model.source`. With `code`, a cheap LLM writes a Starlark + `extract_relevant_data`-style filter (`data = json.decode(INPUT)` → `OUTPUT = json.encode(...)`) + run in a sandbox (no imports/IO, step + 2s limits); the result is accepted only if a **containment** + check proves it's a lossless subset (else fall back to deterministic). `rlm` currently maps to + `code`. **Shines:** big query-focused MCP/API outputs; `code` shines on structured JSON where a + filter can select records precisely. **Inert:** no user query, output < `min_tokens`, projection + not smaller, or (for `code`) no model available → deterministic. ### `smartcrush` Statistical JSON-**array** compressor: parse the array, keep `keep_first` + `keep_last` items plus @@ -201,6 +211,24 @@ $$\Phi = w_R\cdot\text{relevance} + w_H\cdot\text{recency} - w_C\cdot\text{cost} budget, or ≤1 tool output. The scalar-ranking essence of lean-ctx's Context Field (the full heat-diffusion/MMR/bandit machinery is a documented refinement). +### `summarize` (LLM) +Compresses the **middle of the trajectory** into one LLM-written summary (ported from CE-Manager's +ReSum-style summarizer). Restructures the message list to `[msg0, , last-K]`; +the replaced span is stashed under a marker carried in the summary message, so `expand` restores the +full earlier trajectory. This is the one component that changes the message count — `apply.Body` +rebuilds the body keeping the retained messages byte-identical. + +``` +before: [system, u1, tool, a1, tool, u2, … 30 turns …, uN-1, uN] +after: [system, "=== History Summary === … … <>", uN-1, uN] +``` + +- **Config:** `summary_level` (`concise`|`regular`|`highly_detailed`), `keep_last` (3), + `start_from_message` (6 — no-op until the list is this long), `min_tokens` (500), `include_tool_calls` + (false → tool outputs masked in the trajectory), `model.source`. **Shines:** long agentic sessions + where the bulk is stale middle context. **Inert:** short transcripts, span below `min_tokens`, or no + model available (no-op). Run it **alone** (its own preset) — it restructures the whole transcript. + --- ## The DSL filter engine diff --git a/docs/design.md b/docs/design.md index b31e5da..f154392 100644 --- a/docs/design.md +++ b/docs/design.md @@ -205,3 +205,26 @@ store: { ttl_seconds: 1800, max_entries: 1000 } A component registers its constructor + config type via `init()`; adding one makes it YAML-configurable with no core edit. See [components.md](components.md) for presets and every component's config. + +## LLM components + +Most components are deterministic. Two call an LLM: `extract` (`strategy: code`/`rlm`, a Starlark +filter run in a sandbox) and `summarize` (whole-transcript summary). They implement `NeedsModel` and +call `Ctx.Model` — a `ModelSpec` the host resolves per request: + +```mermaid +flowchart LR + cfg["component config
model.source"] --> res{"ModelSpec.For(source)"} + res -->|incoming| inc["Incoming: request's own
model + upstream + key
(built in proxy.chat)"] + res -->|config| stat["Static: cheap model
(CHEAP_MODEL* env)"] + res -->|nil| deg["degrade: extract→deterministic,
summarize→no-op"] + inc --> call["Model.Complete(ctx, prompt)"] + stat --> call +``` + +- **`incoming`** (default) reuses the proxied request's model + the gateway's key — zero extra config, + works through the eval-containers gateway. **`config`** uses a dedicated cheap model (`internal/cheapmodel` + Anthropic/OpenAI). The AuthBridge host offers only `config` (its incoming key is a placeholder). +- The call is synchronous in the request path, so it's bounded (short timeout, retry) and **fail-open**: + any error reverts the component (pipeline guarantee), and a missing model degrades gracefully. +- Reversibility is unchanged — the LLM output is still stashed under a `<>` marker for `expand`. diff --git a/go.mod b/go.mod index 242cc9d..dd778e1 100644 --- a/go.mod +++ b/go.mod @@ -26,6 +26,7 @@ require ( github.com/tidwall/sjson v1.2.5 github.com/tiktoken-go/tokenizer v0.7.0 github.com/tree-sitter/go-tree-sitter v0.25.0 + go.starlark.net v0.0.0-20260708150628-5395d018f003 gopkg.in/yaml.v3 v3.0.1 ) @@ -38,7 +39,6 @@ require ( github.com/bytedance/sonic/loader v0.5.1 // indirect github.com/cloudwego/base64x v0.1.6 // indirect github.com/dlclark/regexp2 v1.11.5 // indirect - github.com/google/go-cmp v0.7.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/invopop/jsonschema v0.13.0 // indirect github.com/klauspost/compress v1.18.6 // indirect diff --git a/go.sum b/go.sum index 7179a7a..82fd711 100644 --- a/go.sum +++ b/go.sum @@ -139,10 +139,14 @@ github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZ github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= +go.starlark.net v0.0.0-20260708150628-5395d018f003 h1:cAxcqHgW8fnmT0cEBU3TzvVYHIFt8IIGDMWUF6rImk4= +go.starlark.net v0.0.0-20260708150628-5395d018f003/go.mod h1:Iue6g6iirlfLoVi/DYCi5/x0h/bAOuWF3dULTKpt2Vo= golang.org/x/arch v0.23.0 h1:lKF64A2jF6Zd8L0knGltUnegD62JMFBiCPBmQpToHhg= golang.org/x/arch v0.23.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A= golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af h1:+5/Sw3GsDNlEmu7TfklWKPdQ0Ykja5VEmq2i817+jbI= +google.golang.org/protobuf v1.36.12-0.20260120151049-f2248ac996af/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/internal/cheapmodel/anthropic.go b/internal/cheapmodel/anthropic.go new file mode 100644 index 0000000..f237c53 --- /dev/null +++ b/internal/cheapmodel/anthropic.go @@ -0,0 +1,87 @@ +// Package cheapmodel provides a minimal Anthropic Messages client used as the +// engine's injected extraction model. It implements engine.Model. (An OpenAI variant +// is a straightforward follow-up; the canonical model is Anthropic-shaped, so the +// Anthropic client is the natural first.) +package cheapmodel + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" +) + +// Anthropic calls a small Anthropic model with a single user prompt and returns the +// text of the first content block. +type Anthropic struct { + BaseURL string // default https://api.anthropic.com + APIKey string + Model string // e.g. claude-haiku-4-5 + MaxTokens int // default 2048 + Client *http.Client + // AuthScheme selects how the API key is sent. "" or "x-api-key" sends the + // x-api-key header (Anthropic default). "bearer" sends + // Authorization: Bearer instead, for gateways (e.g. an IBM LiteLLM + // Anthropic-compatible endpoint) that authenticate with a bearer token. The + // anthropic-version header is sent in both cases. + AuthScheme string +} + +func (a Anthropic) Complete(ctx context.Context, prompt string) (string, error) { + base := a.BaseURL + if base == "" { + base = "https://api.anthropic.com" + } + maxTok := a.MaxTokens + if maxTok == 0 { + maxTok = 2048 + } + client := a.Client + if client == nil { + client = http.DefaultClient + } + reqBody, _ := json.Marshal(map[string]any{ + "model": a.Model, + "max_tokens": maxTok, + "messages": []any{map[string]any{"role": "user", "content": prompt}}, + }) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + strings.TrimRight(base, "/")+"/v1/messages", bytes.NewReader(reqBody)) + if err != nil { + return "", err + } + req.Header.Set("content-type", "application/json") + if a.AuthScheme == "bearer" { + req.Header.Set("Authorization", "Bearer "+a.APIKey) + } else { + req.Header.Set("x-api-key", a.APIKey) + } + req.Header.Set("anthropic-version", "2023-06-01") + + resp, err := client.Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("cheapmodel: status %d", resp.StatusCode) + } + var out struct { + Content []struct { + Text string `json:"text"` + } `json:"content"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return "", err + } + // Return the first content block that carries text. A leading non-text block + // (e.g. "thinking") has an empty Text, so we skip it rather than returning "". + for _, c := range out.Content { + if c.Text != "" { + return c.Text, nil + } + } + return "", nil +} diff --git a/internal/cheapmodel/cheapmodel_test.go b/internal/cheapmodel/cheapmodel_test.go new file mode 100644 index 0000000..06df247 --- /dev/null +++ b/internal/cheapmodel/cheapmodel_test.go @@ -0,0 +1,75 @@ +package cheapmodel + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "testing" +) + +func TestAnthropicSkipsNonTextLeadingBlock(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"content":[{"type":"thinking","text":""},{"type":"text","text":"RESULT"}]}`) + })) + defer srv.Close() + got, err := Anthropic{BaseURL: srv.URL, Model: "m"}.Complete(context.Background(), "p") + if err != nil || got != "RESULT" { + t.Fatalf("got %q err %v", got, err) + } +} + +func TestAnthropicBearerAuth(t *testing.T) { + var gotAuth, gotKey, gotVersion string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotKey = r.Header.Get("x-api-key") + gotVersion = r.Header.Get("anthropic-version") + _, _ = io.WriteString(w, `{"content":[{"type":"text","text":"OK"}]}`) + })) + defer srv.Close() + _, err := Anthropic{BaseURL: srv.URL, APIKey: "tok", Model: "m", AuthScheme: "bearer"}.Complete(context.Background(), "p") + if err != nil { + t.Fatalf("err %v", err) + } + if gotAuth != "Bearer tok" { + t.Fatalf("Authorization = %q, want %q", gotAuth, "Bearer tok") + } + if gotKey != "" { + t.Fatalf("x-api-key = %q, want empty", gotKey) + } + if gotVersion != "2023-06-01" { + t.Fatalf("anthropic-version = %q, want 2023-06-01", gotVersion) + } +} + +func TestAnthropicDefaultAuthUsesAPIKey(t *testing.T) { + var gotAuth, gotKey string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAuth = r.Header.Get("Authorization") + gotKey = r.Header.Get("x-api-key") + _, _ = io.WriteString(w, `{"content":[{"type":"text","text":"OK"}]}`) + })) + defer srv.Close() + _, err := Anthropic{BaseURL: srv.URL, APIKey: "tok", Model: "m"}.Complete(context.Background(), "p") + if err != nil { + t.Fatalf("err %v", err) + } + if gotKey != "tok" { + t.Fatalf("x-api-key = %q, want %q", gotKey, "tok") + } + if gotAuth != "" { + t.Fatalf("Authorization = %q, want empty", gotAuth) + } +} + +func TestOpenAIComplete(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, `{"choices":[{"message":{"content":"OUT"}}]}`) + })) + defer srv.Close() + got, err := OpenAI{BaseURL: srv.URL, Model: "m"}.Complete(context.Background(), "p") + if err != nil || got != "OUT" { + t.Fatalf("got %q err %v", got, err) + } +} diff --git a/internal/cheapmodel/openai.go b/internal/cheapmodel/openai.go new file mode 100644 index 0000000..fe05a18 --- /dev/null +++ b/internal/cheapmodel/openai.go @@ -0,0 +1,70 @@ +package cheapmodel + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "strings" +) + +// OpenAI calls a small OpenAI chat-completions model with a single user prompt and +// returns the content of the first choice. It implements engine.Model. +type OpenAI struct { + BaseURL string // default https://api.openai.com + APIKey string + Model string // e.g. gpt-4o-mini + MaxTokens int // default 2048 + Client *http.Client +} + +func (o OpenAI) Complete(ctx context.Context, prompt string) (string, error) { + base := o.BaseURL + if base == "" { + base = "https://api.openai.com" + } + maxTok := o.MaxTokens + if maxTok == 0 { + maxTok = 2048 + } + client := o.Client + if client == nil { + client = http.DefaultClient + } + reqBody, _ := json.Marshal(map[string]any{ + "model": o.Model, + "max_tokens": maxTok, + "messages": []any{map[string]any{"role": "user", "content": prompt}}, + }) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, + strings.TrimRight(base, "/")+"/v1/chat/completions", bytes.NewReader(reqBody)) + if err != nil { + return "", err + } + req.Header.Set("content-type", "application/json") + req.Header.Set("authorization", "Bearer "+o.APIKey) + + resp, err := client.Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("cheapmodel: status %d", resp.StatusCode) + } + var out struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { + return "", err + } + if len(out.Choices) == 0 { + return "", nil + } + return out.Choices[0].Message.Content, nil +} diff --git a/internal/extract/cache.go b/internal/extract/cache.go new file mode 100644 index 0000000..f128ded --- /dev/null +++ b/internal/extract/cache.go @@ -0,0 +1,27 @@ +package extract + +import "sync" + +// Cache memoizes extraction results by ContentKey, so the same large output re-sent +// on a later turn skips the paid model call. In-memory, process-local. +// +// ponytail: a plain map+mutex, unbounded. Add an LRU/TTL only if memory ever bites. +type Cache struct { + mu sync.Mutex + m map[string]string +} + +func NewCache() *Cache { return &Cache{m: map[string]string{}} } + +func (c *Cache) Get(key string) (string, bool) { + c.mu.Lock() + defer c.mu.Unlock() + v, ok := c.m[key] + return v, ok +} + +func (c *Cache) Put(key, val string) { + c.mu.Lock() + c.m[key] = val + c.mu.Unlock() +} diff --git a/internal/extract/contain.go b/internal/extract/contain.go new file mode 100644 index 0000000..08d6f9f --- /dev/null +++ b/internal/extract/contain.go @@ -0,0 +1,101 @@ +// Package extract is the cheap-model tool-output extractor. A small model returns a +// SMALLER value of the same shape (selecting records/fields, never paraphrasing); the +// containment validator here PROVES the result is a lossless projection of the +// original, so a chatty model can never corrupt a value the agent relies on. On any +// failure it falls back to a deterministic projection, then to the original +// (fail-open). Ported from the reference prototype's actions/llm_compact.py. +// +// ponytail: the model returns the selected JSON directly and containment verifies it — +// no model-generated code (the reference prototype's Python sandbox) and no custom spec language. The +// containment proof is what makes the mechanism safe, whatever produced the subset. +package extract + +import ( + "encoding/json" + "strings" +) + +// IsContained reports whether result is a lossless projection of original (both +// parsed values): string → substring; list → order-preserving subsequence of +// contained items; dict → keys subset with contained values; numbers/bools/nil → +// equal where present; nil result → always allowed (dropping is fine). +func IsContained(result, original any) bool { + return checkContained(result, original) +} + +func checkContained(out, in any) bool { + if out == nil { + return true // dropping content is always allowed + } + if in == nil { + return false + } + switch o := out.(type) { + case string: + s, ok := in.(string) + return ok && strings.Contains(s, o) + case bool: + b, ok := in.(bool) + return ok && b == o + case json.Number: + return numbersEqual(o, in) + case float64: + return numbersEqual(o, in) + case []any: + il, ok := in.([]any) + if !ok || len(o) > len(il) { + return false + } + i := 0 + for _, item := range o { + matched := false + for i < len(il) { + if checkContained(item, il[i]) { + matched = true + i++ + break + } + i++ + } + if !matched { + return false + } + } + return true + case map[string]any: + im, ok := in.(map[string]any) + if !ok { + return false + } + for k, v := range o { + iv, present := im[k] + if !present { + return false // extra key not in input + } + if !checkContained(v, iv) { + return false + } + } + return true + default: + return out == in + } +} + +// numbersEqual compares two JSON-decoded numbers (json.Number or float64) by value. +func numbersEqual(a, b any) bool { + return numStr(a) == numStr(b) +} + +func numStr(v any) string { + switch n := v.(type) { + case json.Number: + return n.String() + case float64: + // json.Number-style rendering via json.Marshal keeps integers integral. + b, _ := json.Marshal(n) + return string(b) + default: + return "" + } +} diff --git a/internal/extract/deterministic.go b/internal/extract/deterministic.go new file mode 100644 index 0000000..21b569f --- /dev/null +++ b/internal/extract/deterministic.go @@ -0,0 +1,180 @@ +package extract + +import ( + "regexp" + "strings" +) + +// Deterministic, model-free projection: filter a parsed value to the parts that match +// the keep-set (plus the important-key "spine"). Every leaf it emits is an unchanged +// value, a string prefix, or a contiguous window, so its output always passes +// IsContained by construction and is never empty. Ported from the reference prototype's +// deterministic_project. + +var importantKeyTokens = []string{"id", "status", "state", "name", "error", "reason", "date", "time"} + +var jsonKeyRe = regexp.MustCompile(`"([A-Za-z_][A-Za-z0-9_\-]{1,})"\s*:`) + +func isImportantKey(keyLower string) bool { + for _, tok := range importantKeyTokens { + if strings.Contains(keyLower, tok) { + return true + } + } + return false +} + +func termMatches(term, keyLower string) bool { + return term == keyLower || strings.Contains(keyLower, term) || strings.Contains(term, keyLower) +} + +func truncateValue(value any, maxChars int) any { + if value == nil || maxChars <= 0 { + return value + } + switch v := value.(type) { + case string: + if r := []rune(v); len(r) > maxChars { + return string(r[:maxChars]) + } + return v + case []any: + out := make([]any, len(v)) + for i, x := range v { + out[i] = truncateValue(x, maxChars) + } + return out + case map[string]any: + out := make(map[string]any, len(v)) + for k, x := range v { + out[k] = truncateValue(x, maxChars) + } + return out + } + return value +} + +func extractTextWindow(text string, terms []string, maxChars int) string { + runes := []rune(text) + if maxChars <= 0 || len(runes) <= maxChars { + return text + } + loweredRunes := []rune(strings.ToLower(text)) + for _, term := range terms { + t := strings.TrimSpace(strings.ToLower(term)) + if t == "" { + continue + } + // Index into the rune slice (not bytes) so the window math stays rune-aligned. + if idx := runeIndex(loweredRunes, []rune(t)); idx >= 0 { + start := idx - maxChars/4 + if start < 0 { + start = 0 + } + if start > len(runes)-maxChars { + start = len(runes) - maxChars + } + return string(runes[start : start+maxChars]) + } + } + return string(runes[:maxChars]) +} + +// runeIndex returns the index (in runes) of the first occurrence of sub in s, or -1. +func runeIndex(s, sub []rune) int { + if len(sub) == 0 { + return 0 + } + for i := 0; i+len(sub) <= len(s); i++ { + match := true + for j := range sub { + if s[i+j] != sub[j] { + match = false + break + } + } + if match { + return i + } + } + return -1 +} + +func projectMapping(m map[string]any, terms map[string]struct{}, maxChars int) map[string]any { + out := map[string]any{} + for key, value := range m { + kl := strings.ToLower(key) + keep := isImportantKey(kl) + if !keep { + for t := range terms { + if termMatches(t, kl) { + keep = true + break + } + } + } + if !keep { + if s, ok := value.(string); ok { + vl := strings.ToLower(s) + for t := range terms { + if strings.Contains(vl, t) { + keep = true + break + } + } + } + } + if keep { + out[key] = truncateValue(value, maxChars) + } + } + return out +} + +// DeterministicProject filters a parsed value to the keep-set + important-key spine. +// Never empties: a value that projects to nothing falls back to a truncated copy. +func DeterministicProject(value any, keepIDs []string, maxChars int) any { + terms := map[string]struct{}{} + for _, k := range keepIDs { + if k != "" { + terms[strings.ToLower(k)] = struct{}{} + } + } + switch v := value.(type) { + case string: + winTerms := append([]string{}, keepIDs...) + for i, m := range jsonKeyRe.FindAllStringSubmatch(v, -1) { + if i >= 16 { + break + } + winTerms = append(winTerms, m[1]) + } + return extractTextWindow(v, winTerms, maxChars) + case map[string]any: + if len(terms) > 0 { + if p := projectMapping(v, terms, maxChars); len(p) > 0 { + return p + } + } + return truncateValue(v, maxChars) + case []any: + var kept []any + for _, item := range v { + m, ok := item.(map[string]any) + if !ok { + kept = append(kept, truncateValue(item, maxChars)) + continue + } + if len(terms) > 0 { + if p := projectMapping(m, terms, maxChars); len(p) > 0 { + kept = append(kept, p) + } + } + } + if len(kept) > 0 { + return kept + } + return truncateValue(v, maxChars) + } + return truncateValue(value, maxChars) +} diff --git a/internal/extract/extract.go b/internal/extract/extract.go new file mode 100644 index 0000000..c80f599 --- /dev/null +++ b/internal/extract/extract.go @@ -0,0 +1,353 @@ +package extract + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "regexp" + "strings" + "sync" + + "github.com/kagenti/context-guru/internal/tokens" +) + +// cgMarkerRe matches the offload marker so ContentKey is marker-insensitive (a +// re-sent body that a sibling component marked still hits the extraction cache). +var cgMarkerRe = regexp.MustCompile(`<>`) + +const sampleChars = 4000 + +// Model is the cheap-model client the extractor calls. The host (proxy / plugin) +// injects a concrete implementation; the extractor core stays transport-free. +type Model interface { + Complete(ctx context.Context, prompt string) (string, error) +} + +// Cfg configures extraction. +type Cfg struct { + Mode string // auto | single | rlm | deterministic + Floor int // token floor; rlm kicks in at max(floor*4, 8000) in auto + MinKeepRatio float64 // 0 disables the blunt ratio backstop (keep-set check governs) + AllowDeterministic bool + MaxChars int // deterministic projection window + // AllowedStrategies, when non-empty, restricts strategyOrder to these strategy names + // (code | single | rlm | deterministic) preserving the computed order. Empty means + // "all" — prior behavior. Lets config enable/disable strategies purely by name. + AllowedStrategies []string +} + +// DefaultCfg mirrors the reference prototype's ExtractCfg defaults. +func DefaultCfg() Cfg { + return Cfg{Mode: "auto", Floor: 3000, AllowDeterministic: true, MaxChars: sampleChars} +} + +var wsRe = regexp.MustCompile(`\s+`) + +// ContentKey is a stable, marker- and whitespace-insensitive key for a body, so the +// same output re-sent on a later turn hits the extraction cache. +func ContentKey(text string) string { + s := wsRe.ReplaceAllString(cgMarkerRe.ReplaceAllString(text, ""), " ") + s = strings.TrimSpace(s) + sum := sha256.Sum256([]byte(s)) + return hex.EncodeToString(sum[:])[:24] +} + +var identRe = regexp.MustCompile(`[A-Za-z_][\w./-]{3,}|\b\d{3,}\b`) + +// HarvestIdentifiers pulls distinctive identifiers (paths, symbols, ids, numbers) +// from the agent's recent turns — the keep-set the extractor must retain. +func HarvestIdentifiers(text string, cap int) []string { + var seen []string + idx := map[string]struct{}{} + for _, m := range identRe.FindAllString(text, -1) { + if _, ok := idx[m]; ok { + continue + } + idx[m] = struct{}{} + seen = append(seen, m) + if len(seen) >= cap { + break + } + } + return seen +} + +// parseBody returns the parsed value handed to the extractor: JSON if possible, then +// NDJSON (as a list), else the raw string. +func parseBody(text string) any { + var v any + dec := json.NewDecoder(strings.NewReader(text)) + dec.UseNumber() + if err := dec.Decode(&v); err == nil { + return v + } + if recs := parseNDJSON(text); recs != nil { + return recs + } + return text +} + +func parseNDJSON(text string) []any { + var lines []string + for _, ln := range strings.Split(text, "\n") { + if strings.TrimSpace(ln) != "" { + lines = append(lines, ln) + } + } + if len(lines) < 2 { + return nil + } + var recs []any + for _, ln := range lines { + var v any + dec := json.NewDecoder(strings.NewReader(ln)) + dec.UseNumber() + if err := dec.Decode(&v); err != nil { + return nil + } + switch v.(type) { + case map[string]any, []any: + recs = append(recs, v) + default: + return nil + } + } + return recs +} + +func resultToText(v any) string { + switch v.(type) { + case map[string]any, []any: + // Compact, not indented: this is the reduced tool-output value, and + // indentation whitespace inflates the BPE token count (it can exceed the + // original), which trips the never-inflate gate and silently drops the + // extraction. Compact JSON keeps the projection a real reduction. + b, err := json.Marshal(v) + if err == nil { + return string(b) + } + } + return fmt.Sprint(v) +} + +func stripFences(s string) string { + c := strings.TrimSpace(s) + if !strings.HasPrefix(c, "```") { + return c + } + lines := strings.Split(c, "\n") + if strings.HasPrefix(lines[0], "```") { + lines = lines[1:] + } + if len(lines) > 0 && strings.TrimSpace(lines[len(lines)-1]) == "```" { + lines = lines[:len(lines)-1] + } + return strings.Join(lines, "\n") +} + +// --- sanity + validation gate --- + +func extractionIsSane(bodyText, resultText string, keepIDs []string, minKeepRatio float64) bool { + if resultText == "" { + return false + } + bodyN, resN := tokens.Count(bodyText), tokens.Count(resultText) + switch strings.TrimSpace(resultText) { + case "", "[]", "{}", "null", `""`: + if bodyN > 0 { + return false + } + } + if minKeepRatio > 0 && float64(resN) < minKeepRatio*float64(bodyN) { + return false + } + for _, kid := range keepIDs { + if len(kid) >= 5 && strings.ContainsFunc(kid, isLetter) && + strings.Contains(bodyText, kid) && !strings.Contains(resultText, kid) { + return false + } + } + return true +} + +func isLetter(r rune) bool { + return (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') +} + +func validateExtraction(resultText, bodyText string, keepIDs []string, cfg Cfg) bool { + if !extractionIsSane(bodyText, resultText, keepIDs, cfg.MinKeepRatio) { + return false + } + return IsContained(parseBody(resultText), parseBody(bodyText)) +} + +// intersectAllowed filters order to the allowed set (non-empty), preserving order. +// Empty allowed means no filtering. +func intersectAllowed(order, allowed []string) []string { + if len(allowed) == 0 { + return order + } + allow := make(map[string]struct{}, len(allowed)) + for _, a := range allowed { + allow[a] = struct{}{} + } + out := make([]string, 0, len(order)) + for _, s := range order { + if _, ok := allow[s]; ok { + out = append(out, s) + } + } + return out +} + +func strategyOrder(tokenEst int, cfg Cfg) []string { + return intersectAllowed(rawStrategyOrder(tokenEst, cfg), cfg.AllowedStrategies) +} + +func rawStrategyOrder(tokenEst int, cfg Cfg) []string { + switch cfg.Mode { + case "deterministic": + return []string{"deterministic"} + case "single", "rlm": + order := []string{cfg.Mode} + if cfg.AllowDeterministic { + order = append(order, "deterministic") + } + return order + case "code": + order := []string{"code"} + if cfg.AllowDeterministic { + order = append(order, "deterministic") + } + return order + default: // auto + // "code" (model-written Starlark filter over the full body) is primary for + // mid-size bodies; "rlm" (chunked) above the floor. "single" (JSON-return) and + // "deterministic" are ordered fallbacks behind the primary. + floor4 := cfg.Floor * 4 + if floor4 < 8000 { + floor4 = 8000 + } + var order []string + if tokenEst >= floor4 { + order = []string{"rlm", "code", "single"} + } else { + order = []string{"code", "single"} + } + if cfg.AllowDeterministic { + order = append(order, "deterministic") + } + return order + } +} + +// RunExtraction tries strategies in order, returning the first candidate that is +// strictly smaller AND passes the validation gate, else ("", "none"). Fail-open: the +// caller keeps the original on "none". +func RunExtraction(ctx context.Context, body, goal string, keepIDs []string, tokenEst int, cfg Cfg, model Model) (string, string) { + base := tokens.Count(body) + for _, name := range strategyOrder(tokenEst, cfg) { + var cand string + switch name { + case "code": + cand = runStarlark(ctx, body, goal, keepIDs, model) + case "single": + cand = runSingle(ctx, body, goal, keepIDs, model) + case "rlm": + cand = runRLMBatched(ctx, body, goal, keepIDs, model) + case "deterministic": + cand = resultToText(DeterministicProject(parseBody(body), keepIDs, cfg.MaxChars)) + } + if cand == "" || tokens.Count(cand) >= base { + continue + } + if validateExtraction(cand, body, keepIDs, cfg) { + return cand, name + } + } + return "", "none" +} + +// runSingle asks the model for the filtered subset in one call. It is a FALLBACK +// behind the primary full-body strategies: "code" (a model-written Starlark filter) +// and "rlm" (chunked) both run over the FULL body, so they never truncate. runSingle +// inlines the body into the prompt via buildPrompt, which truncates to sampleChars to +// bound prompt cost — acceptable for a fallback, since whatever the model returns is +// still containment-checked against the full body before it can be spliced in. +func runSingle(ctx context.Context, body, goal string, keepIDs []string, model Model) string { + if model == nil { + return "" + } + out, err := model.Complete(ctx, buildPrompt(body, goal, keepIDs)) + if err != nil { + return "" + } + return stripFences(out) +} + +const rlmChunkSize = 20 +const rlmConcurrency = 6 + +// runRLMBatched chunks a large list body and asks the model to filter each chunk +// concurrently, then merges the kept records (order-preserving). Containment over the +// merged result is checked by the caller. Non-list bodies fall back to a single call. +func runRLMBatched(ctx context.Context, body, goal string, keepIDs []string, model Model) string { + if model == nil { + return "" + } + parsed := parseBody(body) + list, ok := parsed.([]any) + if !ok || len(list) <= rlmChunkSize { + return runSingle(ctx, body, goal, keepIDs, model) + } + var chunks [][]any + for i := 0; i < len(list); i += rlmChunkSize { + end := i + rlmChunkSize + if end > len(list) { + end = len(list) + } + chunks = append(chunks, list[i:end]) + } + + results := make([][]any, len(chunks)) + sem := make(chan struct{}, rlmConcurrency) + var wg sync.WaitGroup + for ci, chunk := range chunks { + wg.Add(1) + sem <- struct{}{} + go func(ci int, chunk []any) { + defer wg.Done() + defer func() { <-sem }() + defer func() { _ = recover() }() + chunkJSON, err := json.Marshal(chunk) + if err != nil { + return + } + out, err := model.Complete(ctx, buildPrompt(string(chunkJSON), goal, keepIDs)) + if err != nil { + return + } + var kept []any + if json.Unmarshal([]byte(stripFences(out)), &kept) == nil { + results[ci] = kept + } + // On parse failure the chunk contributes nothing (safe drop; containment holds). + }(ci, chunk) + } + wg.Wait() + + var merged []any + for _, r := range results { + merged = append(merged, r...) + } + if len(merged) == 0 { + return "" + } + b, err := json.Marshal(merged) + if err != nil { + return "" + } + return string(b) +} diff --git a/internal/extract/prompt.go b/internal/extract/prompt.go new file mode 100644 index 0000000..18b33e8 --- /dev/null +++ b/internal/extract/prompt.go @@ -0,0 +1,147 @@ +package extract + +import ( + "encoding/json" + "strings" +) + +// Prompt building. Because the model returns the filtered VALUE (which containment +// then verifies), it must SEE the values — so the prompt shows the actual JSON/text +// (truncated). For very large lists the RLM strategy chunks the body so each chunk is +// shown in full. The rule set is the reference prototype's "select, never summarize, recall-first" +// contract, retargeted from "write a function" to "return the JSON". + +// sampleMarker precedes the body in the prompt; tests and the (future) model both +// locate the payload after it. +const sampleMarker = "INPUT (return a smaller value of this same shape):\n" + +const rules = `Return ONLY a JSON value (or, for raw text input, the kept text): a SMALLER value +of the SAME shape, selecting only what the agent needs next. Rules, in priority order: +1. RECALL FIRST. When unsure whether a record/field is relevant, KEEP IT. +2. SELECT, NEVER SUMMARIZE. Return whole records/objects/values byte-for-byte. Never + paraphrase, truncate, round, reformat, or invent values. +3. PRESERVE EXACTLY: ids, numbers, names, paths, timestamps, error messages, stack + traces — and anything matching the KEEP list. +4. Only drop records that are CLEARLY irrelevant boilerplate, duplicates, or noise. +5. If you cannot identify clearly-irrelevant content, RETURN THE INPUT UNCHANGED. +6. Keep the natural shape and types. Output ONLY the value — no prose, no markdown.` + +const example = `EXAMPLE +Goal: "Fix failing test test_auth_expiry; find the relevant hit." +KEEP: ["test_auth_expiry","auth/session.py"] +INPUT: [{"path":"auth/session.py","snippet":"def test_auth_expiry()..."},{"path":"README.md","snippet":"intro"}] +OUTPUT: [{"path":"auth/session.py","snippet":"def test_auth_expiry()..."}]` + +func buildPrompt(bodyText, goal string, keepIDs []string) string { + g := strings.TrimSpace(goal) + if g == "" { + g = "(no explicit goal stated)" + } + if len(g) > 2000 { + g = g[:2000] + } + keep := keepIDs + if len(keep) > 60 { + keep = keep[:60] + } + keepBlock := "" + if len(keep) > 0 { + kb, _ := json.Marshal(keep) + keepBlock = "IDENTIFIERS THE AGENT REFERENCED RECENTLY — keep every record or field\n" + + "whose value matches any of these, verbatim:\n" + string(kb) + "\n\n" + } + + // Show the actual value (pretty-printed JSON if it parses), truncated. + sample := bodyText + if v := parseBody(bodyText); !isRawString(v) { + if b, err := json.MarshalIndent(v, "", " "); err == nil { + sample = string(b) + } + } + sample = truncate(sample, sampleChars) + + return "You filter ONE tool output down to only what the agent needs next.\n\n" + + "WHAT THE AGENT IS DOING NOW (filter toward this):\n" + g + "\n\n" + + keepBlock + sampleMarker + sample + "\n\n" + rules + "\n\n" + example +} + +// codeRules is the Starlark code-writing contract: the model writes a program that +// runs over the real INPUT (it is NOT shown the full body), so the prompt stays cheap. +// Same "select, never summarize, recall-first" discipline as buildPrompt, retargeted +// from "return the value" to "write the filter". +const codeRules = `Write a Starlark program (a safe Python subset) that filters this ONE tool output +down to only what the agent needs next. Contract: +- The global string INPUT holds the FULL tool output. The module ` + "`json`" + ` is available. +- Start: data = json.decode(INPUT) +- Select a SMALLER value of the SAME shape (e.g. drop irrelevant list records / fields). +- End: OUTPUT = json.encode(result) # OUTPUT must be a string +Rules, in priority order: +1. RECALL FIRST. When unsure whether a record/field is relevant, KEEP IT. +2. SELECT, NEVER SUMMARIZE. Keep whole records/values byte-for-byte. Never paraphrase, + truncate, round, reformat, or invent values — only DROP clearly-irrelevant ones. +3. PRESERVE EXACTLY: ids, numbers, names, paths, timestamps, errors, stack traces — + and anything matching the KEEP list. +4. NO imports (no load()), NO I/O, NO network. Use only json.decode/json.encode and + plain Starlark (list comprehensions, dict/list ops, string ops). +5. If you cannot identify clearly-irrelevant content, set OUTPUT = INPUT. +Output ONLY the Starlark program — no prose, no markdown fences.` + +const codeExample = `EXAMPLE +Goal: "Fix failing test test_auth_expiry; find the relevant hit." +KEEP: ["test_auth_expiry","auth/session.py"] +INPUT schema: list of {"path": str, "snippet": str} +PROGRAM: +data = json.decode(INPUT) +result = [r for r in data if "auth" in r["path"] or "test_auth_expiry" in r["snippet"]] +OUTPUT = json.encode(result)` + +// buildCodePrompt builds the prompt for the Starlark code-writing strategy. Unlike +// buildPrompt it does NOT inline the full body — the program runs over the real INPUT. +// It shows the parsed shape and a small sample so the model knows the schema cheaply. +func buildCodePrompt(bodyText, goal string, keepIDs []string) string { + g := strings.TrimSpace(goal) + if g == "" { + g = "(no explicit goal stated)" + } + if len(g) > 2000 { + g = g[:2000] + } + keep := keepIDs + if len(keep) > 60 { + keep = keep[:60] + } + keepBlock := "" + if len(keep) > 0 { + kb, _ := json.Marshal(keep) + keepBlock = "IDENTIFIERS THE AGENT REFERENCED RECENTLY — keep every record or field\n" + + "whose value matches any of these, verbatim:\n" + string(kb) + "\n\n" + } + + // A small sample of the parsed shape — enough to infer the schema, not the full body. + sample := bodyText + if v := parseBody(bodyText); !isRawString(v) { + if b, err := json.MarshalIndent(v, "", " "); err == nil { + sample = string(b) + } + } + sample = truncate(sample, codeSampleChars) + + return "You write a Starlark filter for ONE tool output, keeping only what the agent needs next.\n\n" + + "WHAT THE AGENT IS DOING NOW (filter toward this):\n" + g + "\n\n" + + keepBlock + "INPUT SAMPLE (the real INPUT at runtime is the FULL output of this same shape):\n" + + sample + "\n\n" + codeRules + "\n\n" + codeExample +} + +const codeSampleChars = 1500 + +func isRawString(v any) bool { + _, ok := v.(string) + return ok +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/internal/extract/starlark.go b/internal/extract/starlark.go new file mode 100644 index 0000000..96709ae --- /dev/null +++ b/internal/extract/starlark.go @@ -0,0 +1,63 @@ +package extract + +import ( + "context" + "time" + + starjson "go.starlark.net/lib/json" + "go.starlark.net/starlark" +) + +const ( + starlarkMaxSteps = 50_000_000 + starlarkTimeout = 2 * time.Second +) + +// runStarlark asks the model for a Starlark program whose contract is: read the +// global string INPUT (the full tool output), assign a string global OUTPUT (the +// filtered value). It runs sandboxed over the FULL body — no imports, no I/O, step + +// time limits — and returns OUTPUT, or "" on any failure (fail-open). Containment is +// verified by the caller (RunExtraction). +func runStarlark(ctx context.Context, body, goal string, keepIDs []string, model Model) (out string) { + defer func() { + if recover() != nil { + out = "" + } + }() + if model == nil { + return "" + } + src, err := model.Complete(ctx, buildCodePrompt(body, goal, keepIDs)) + if err != nil { + return "" + } + src = stripFences(src) + + ctx, cancel := context.WithTimeout(ctx, starlarkTimeout) + defer cancel() + thread := &starlark.Thread{Name: "extract"} // Load==nil => load() disabled + thread.SetMaxExecutionSteps(starlarkMaxSteps) + done := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + thread.Cancel(ctx.Err().Error()) + case <-done: + } + }() + defer close(done) + + predeclared := starlark.StringDict{ + "json": starjson.Module, + "INPUT": starlark.String(body), + } + globals, err := starlark.ExecFile(thread, "extract.star", src, predeclared) + if err != nil { + return "" + } + res, ok := globals["OUTPUT"].(starlark.String) + if !ok { + return "" + } + return string(res) +} diff --git a/internal/extract/starlark_test.go b/internal/extract/starlark_test.go new file mode 100644 index 0000000..e0eb72a --- /dev/null +++ b/internal/extract/starlark_test.go @@ -0,0 +1,58 @@ +package extract + +import ( + "context" + "strconv" + "strings" + "testing" +) + +// starlarkModel returns a fixed Starlark program that keeps records whose name +// contains "keep" — exercises real code execution over the full input. +type starlarkModel struct{} + +func (starlarkModel) Complete(_ context.Context, _ string) (string, error) { + return ` +data = json.decode(INPUT) +kept = [r for r in data if "keep" in r["name"]] +OUTPUT = json.encode(kept) +`, nil +} + +func TestRunStarlarkFiltersFullBody(t *testing.T) { + var recs []string + for i := 0; i < 100; i++ { + name := "drop" + if i%10 == 0 { + name = "keep" + } + recs = append(recs, `{"id":`+strconv.Itoa(i)+`,"name":"`+name+`"}`) + } + body := "[" + strings.Join(recs, ",") + "]" + out := runStarlark(context.Background(), body, "find keep", nil, starlarkModel{}) + if out == "" { + t.Fatal("expected a Starlark result") + } + if !IsContained(parseBody(out), parseBody(body)) { + t.Fatalf("Starlark output must be a contained subset: %s", out) + } + if strings.Contains(out, "drop") { + t.Fatal("filter should have dropped non-keep records") + } + if !strings.Contains(out, "keep") { + t.Fatal("filter should have kept the keep records (recall, not truncation)") + } +} + +// malicious program must fail-open (no panic, returns ""). +type evilModel struct{} + +func (evilModel) Complete(_ context.Context, _ string) (string, error) { + return `load("os", "x")`, nil // imports disabled +} + +func TestRunStarlarkFailsOpenOnDisallowed(t *testing.T) { + if out := runStarlark(context.Background(), `[{"a":1}]`, "", nil, evilModel{}); out != "" { + t.Fatalf("disallowed program must fail open to \"\", got %q", out) + } +} diff --git a/proxy/proxy.go b/proxy/proxy.go index f77d9e9..220f685 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -21,10 +21,12 @@ import ( "github.com/kagenti/context-guru/apply" "github.com/kagenti/context-guru/components" "github.com/kagenti/context-guru/expand" + "github.com/kagenti/context-guru/internal/cheapmodel" "github.com/kagenti/context-guru/metrics" "github.com/kagenti/context-guru/schema" "github.com/kagenti/context-guru/store" bschemas "github.com/maximhq/bifrost/core/schemas" + "github.com/tidwall/gjson" "github.com/tidwall/sjson" ) @@ -43,6 +45,10 @@ type Options struct { // uses this to pin every call to EVAL_MODEL regardless of what the agent asked for. ForceModel string Client *http.Client + // CheapModel is the static "config"-source LLM client for NeedsModel + // components (nil = none). The "incoming"-source client is built per request + // from the route's upstream + the gateway's real key. + CheapModel components.Model } // upstream binds a provider to its base URL, the canonical provider path to POST @@ -105,6 +111,45 @@ func headerKey(name, key string) func(http.Header) { return func(h http.Header) { h.Set(name, key) } } +// incomingModel builds an LLM client that reuses the proxied request's own model +// and the route's upstream + credential, so a NeedsModel component can call the +// same backend the request targets. Prefers the gateway's injected key (gateway +// mode); falls back to the client's own auth header (pass-through). Returns nil +// when no upstream/model/key is resolvable, and the component degrades. +func (h *Handler) incomingModel(provider bschemas.ModelProvider, up upstream, body []byte, r *http.Request) components.Model { + if up.base == "" { + return nil + } + model := gjson.GetBytes(body, "model").String() + if model == "" { + return nil + } + switch provider { + case bschemas.Anthropic: + key := h.opts.AnthropicKey + if key == "" { + key = r.Header.Get("x-api-key") + } + if key == "" { + key = strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + } + if key == "" { + return nil + } + return cheapmodel.Anthropic{BaseURL: up.base, Model: model, APIKey: key, Client: h.client} + case bschemas.OpenAI: + key := h.opts.OpenAIKey + if key == "" { + key = strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + } + if key == "" { + return nil + } + return cheapmodel.OpenAI{BaseURL: up.base, Model: model, APIKey: key, Client: h.client} + } + return nil +} + func (h *Handler) chat(provider bschemas.ModelProvider, up upstream) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { body, err := io.ReadAll(r.Body) @@ -118,11 +163,19 @@ func (h *Handler) chat(provider bschemas.ModelProvider, up upstream) http.Handle body = out } } - // Rewrite the messages via the shared apply path; fail open. - body, _ = apply.Body( + // Rewrite the messages via the shared apply path; fail open. Supply the + // LLM clients NeedsModel components may call: the per-request "incoming" + // model (the route's upstream + the gateway's real key + the request's + // model) and the static "config" cheap model. + models := components.ModelSpec{ + Incoming: h.incomingModel(provider, up, body, r), + Static: h.opts.CheapModel, + } + body, _ = apply.BodyWithModel( r.Context(), h.pipe, h.store, provider, body, r.Header.Get("x-context-guru-session"), strings.EqualFold(r.Header.Get("x-context-guru-bypass"), "true"), + models, ) h.serve(w, r, provider, up, body) } From 10aab92eb08b43fd1d1fd2f6b5c9b6316a57df3d Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Sun, 12 Jul 2026 19:23:38 +0300 Subject: [PATCH 2/5] test: add unit tests for LLM components (containment, model resolution, edge cases) - internal/extract: RunExtraction accepts a contained Starlark result, REJECTS a fabricated (non-contained) one (falls back to deterministic); ContentKey marker-insensitivity; HarvestIdentifiers. - components: ModelSpec.For resolution (incoming/config/nil/fallback). - summarize: model-error fail-open (transcript untouched), empty-response skip. - extract: rlm strategy maps to the code path (RLM deferred). Assisted-By: Claude (Anthropic AI) Signed-off-by: Osher-Elhadad --- components/all/llm_test.go | 64 ++++++++++++++++++++++++++++++++ components/component_test.go | 39 +++++++++++++++++++ components/offload/extract.go | 2 +- internal/extract/extract_test.go | 64 ++++++++++++++++++++++++++++++++ 4 files changed, 168 insertions(+), 1 deletion(-) create mode 100644 components/component_test.go create mode 100644 internal/extract/extract_test.go diff --git a/components/all/llm_test.go b/components/all/llm_test.go index 7469c72..eb39ab9 100644 --- a/components/all/llm_test.go +++ b/components/all/llm_test.go @@ -2,6 +2,7 @@ package all_test import ( "context" + "errors" "strings" "testing" @@ -19,6 +20,8 @@ type stubModel struct { func (m stubModel) Complete(context.Context, string) (string, error) { return m.resp, m.err } +var errBoom = errors.New("boom") + func strp(s string) *string { return &s } func newComp(t *testing.T, name, yaml string) components.Offload { @@ -89,6 +92,67 @@ func TestSummarizeNoModelSkips(t *testing.T) { } } +// TestSummarizeModelErrorFailsOpen: a model error must surface as an error (the +// pipeline reverts the component) and leave the transcript untouched. +func TestSummarizeModelErrorFailsOpen(t *testing.T) { + off := newComp(t, "summarize", "keep_last: 1\nstart_from_message: 0\nmin_tokens: 1\n") + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + {Role: bschemas.ChatMessageRoleSystem, Content: &bschemas.ChatMessageContent{ContentStr: strp("s")}}, + toolMsg(strings.Repeat("output ", 100)), + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("u")}}, + }} + before := len(req.Input) + c := &components.Ctx{Ctx: context.Background(), Store: store.NewMemory(store.Options{}), + Model: components.ModelSpec{Incoming: stubModel{err: errBoom}}} + var rep components.Report + _, err := off.Offload(req, &rep, c) + if err == nil { + t.Fatal("model error must be returned so the pipeline reverts") + } + if len(req.Input) != before { + t.Fatal("transcript must be untouched on model error") + } +} + +// TestSummarizeEmptyResponseSkips: an empty model response is a no-op, not a +// broken summary. +func TestSummarizeEmptyResponseSkips(t *testing.T) { + off := newComp(t, "summarize", "keep_last: 1\nstart_from_message: 0\nmin_tokens: 1\n") + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + {Role: bschemas.ChatMessageRoleSystem, Content: &bschemas.ChatMessageContent{ContentStr: strp("s")}}, + toolMsg(strings.Repeat("output ", 100)), + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("u")}}, + }} + c := &components.Ctx{Ctx: context.Background(), Store: store.NewMemory(store.Options{}), + Model: components.ModelSpec{Incoming: stubModel{resp: " "}}} + var rep components.Report + keys, err := off.Offload(req, &rep, c) + if err != nil || len(keys) != 0 || !rep.Skipped || len(req.Input) != 3 { + t.Fatalf("empty summary must skip untouched: keys=%d skipped=%v len=%d err=%v", len(keys), rep.Skipped, len(req.Input), err) + } +} + +// TestExtractRLMUsesModel: strategy=rlm currently maps to code and still runs the +// model's filter (not silently deterministic). +func TestExtractRLMUsesModel(t *testing.T) { + off := newComp(t, "extract", "strategy: rlm\nmin_tokens: 1\nmodel:\n source: config\n") + st := store.NewMemory(store.Options{}) + body := `[{"id":1,"name":"keep this one"},{"id":2,"name":"drop it"}]` + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + {Role: bschemas.ChatMessageRoleUser, Content: &bschemas.ChatMessageContent{ContentStr: strp("find keep")}}, + toolMsg(body), + }} + filter := "data = json.decode(INPUT)\nOUTPUT = json.encode([r for r in data if \"keep\" in r[\"name\"]])\n" + c := &components.Ctx{Ctx: context.Background(), Store: st, Model: components.ModelSpec{Static: stubModel{resp: filter}}} + var rep components.Report + if keys, err := off.Offload(req, &rep, c); err != nil || len(keys) != 1 { + t.Fatalf("rlm should run the model filter: keys=%v err=%v skipped=%v", keys, err, rep.Skipped) + } + if strings.Contains(schema.MessageText(req.Input[1]), "drop it") { + t.Fatal("rlm/code filter should have dropped the non-keep record") + } +} + // TestExtractCodeUsesModel: the code strategy runs the model's Starlark filter and // keeps only the matching records (a contained subset), with a marker. func TestExtractCodeUsesModel(t *testing.T) { diff --git a/components/component_test.go b/components/component_test.go new file mode 100644 index 0000000..4514197 --- /dev/null +++ b/components/component_test.go @@ -0,0 +1,39 @@ +package components + +import ( + "context" + "testing" +) + +type fakeModel struct{ id string } + +func (m fakeModel) Complete(context.Context, string) (string, error) { return m.id, nil } + +func idOf(m Model) string { + if m == nil { + return "nil" + } + return m.(fakeModel).id +} + +func TestModelSpecFor(t *testing.T) { + inc, stat := fakeModel{"incoming"}, fakeModel{"static"} + cases := []struct { + name string + spec ModelSpec + source string + want string + }{ + {"config picks static", ModelSpec{Incoming: inc, Static: stat}, "config", "static"}, + {"incoming picks incoming", ModelSpec{Incoming: inc, Static: stat}, "incoming", "incoming"}, + {"default picks incoming", ModelSpec{Incoming: inc, Static: stat}, "", "incoming"}, + {"incoming falls back to static", ModelSpec{Static: stat}, "incoming", "static"}, + {"config with no static is nil", ModelSpec{Incoming: inc}, "config", "nil"}, + {"nothing available is nil", ModelSpec{}, "", "nil"}, + } + for _, c := range cases { + if got := idOf(c.spec.For(c.source)); got != c.want { + t.Errorf("%s: For(%q)=%s want %s", c.name, c.source, got, c.want) + } + } +} diff --git a/components/offload/extract.go b/components/offload/extract.go index d8abc7b..ee1ed4c 100644 --- a/components/offload/extract.go +++ b/components/offload/extract.go @@ -106,7 +106,7 @@ func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Repo func (e *Extract) reduce(c *components.Ctx, content, goal string, keepIDs []string, query map[string]struct{}, model components.Model) (string, bool) { if model != nil && e.NeedsModel() { cfg := extract.DefaultCfg() - cfg.Mode = e.strategy + cfg.Mode = "code" // rlm is deferred → use the Starlark code strategy cfg.Floor = e.minTokens if res, _ := extract.RunExtraction(c.Ctx, content, goal, keepIDs, schema.TextTokens(content), cfg, model); res != "" && res != content { return res, true diff --git a/internal/extract/extract_test.go b/internal/extract/extract_test.go new file mode 100644 index 0000000..51b3bc2 --- /dev/null +++ b/internal/extract/extract_test.go @@ -0,0 +1,64 @@ +package extract + +import ( + "context" + "strings" + "testing" +) + +// fixedModel returns a canned Starlark program. +type fixedModel struct{ prog string } + +func (m fixedModel) Complete(context.Context, string) (string, error) { return m.prog, nil } + +const body = `[{"id":1,"name":"keep this"},{"id":2,"name":"drop this"},{"id":3,"name":"keep that"}]` + +// A valid filter that selects a contained subset is accepted (strategy "code"). +func TestRunExtractionCodeAccepted(t *testing.T) { + cfg := DefaultCfg() + cfg.Mode = "code" + m := fixedModel{prog: `data = json.decode(INPUT) +OUTPUT = json.encode([r for r in data if "keep" in r["name"]])`} + out, strat := RunExtraction(context.Background(), body, "find keep records", nil, 5000, cfg, m) + if strat != "code" { + t.Fatalf("expected code strategy to win, got %q (out=%s)", strat, out) + } + if strings.Contains(out, "drop this") || !strings.Contains(out, "keep this") { + t.Fatalf("code filter should keep only the keep records: %s", out) + } + if !IsContained(parseBody(out), parseBody(body)) { + t.Fatal("accepted output must be a contained subset") + } +} + +// A filter that INVENTS data fails the containment check and must be rejected — +// falling back to the deterministic projection (never returning fabricated data). +func TestRunExtractionRejectsNonContained(t *testing.T) { + cfg := DefaultCfg() + cfg.Mode = "code" + m := fixedModel{prog: `OUTPUT = json.encode([{"id":999,"name":"invented record"}])`} + out, strat := RunExtraction(context.Background(), body, "find keep records", nil, 5000, cfg, m) + if strings.Contains(out, "invented") { + t.Fatalf("fabricated (non-contained) output must be rejected, got: %s", out) + } + if strat == "code" { + t.Fatalf("non-contained code result must not win; strategy=%q", strat) + } +} + +// ContentKey ignores <> markers + whitespace so a re-sent body still hits cache. +func TestContentKeyMarkerInsensitive(t *testing.T) { + a := ContentKey("some tool output here") + b := ContentKey("some tool output\n<> here") + if a != b { + t.Fatalf("ContentKey must be marker/whitespace-insensitive: %q vs %q", a, b) + } +} + +func TestHarvestIdentifiers(t *testing.T) { + ids := HarvestIdentifiers("fix auth/session.py and test_auth_expiry (see issue 12345)", 40) + joined := strings.Join(ids, " ") + if !strings.Contains(joined, "auth/session.py") || !strings.Contains(joined, "test_auth_expiry") { + t.Fatalf("expected paths/symbols harvested, got %v", ids) + } +} From ee42c6bacddc57ecaca89d6c32c532d9f3e8db5e Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Sun, 12 Jul 2026 20:02:40 +0300 Subject: [PATCH 3/5] feat: bound LLM component calls with a timeout + support full-config sweep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 30s per-call timeout on extract-code / summarize model calls (fail-open if the model is slow), so an in-request LLM call can't stall the agent. - gateway start accepts CONTEXT_GURU_CONFIG_YAML (full config doc) so a sweep can pin per-component config (strategy: code, model.source); forwarded via the compose override. - sweep.py: cg-extract-code / cg-summarize configs (model.source: incoming) with a higher LLM floor (min_tokens 1500 — only large outputs pay for a call). Assisted-By: Claude (Anthropic AI) Signed-off-by: Osher-Elhadad --- components/offload/extract.go | 11 ++++++++++- components/offload/summarize.go | 4 +++- deploy/eval-containers/compose.contextguru.yaml | 2 ++ deploy/eval-containers/start | 7 +++++++ deploy/eval-containers/sweep.py | 14 ++++++++++++++ 5 files changed, 36 insertions(+), 2 deletions(-) diff --git a/components/offload/extract.go b/components/offload/extract.go index ee1ed4c..ec043bc 100644 --- a/components/offload/extract.go +++ b/components/offload/extract.go @@ -1,7 +1,9 @@ package offload import ( + "context" "strings" + "time" "github.com/kagenti/context-guru/components" "github.com/kagenti/context-guru/expand" @@ -11,6 +13,10 @@ import ( "gopkg.in/yaml.v3" ) +// llmCallTimeout bounds a single in-request model call so a slow/hung model +// fails open (the component falls back / reverts) instead of stalling the agent. +const llmCallTimeout = 30 * time.Second + func init() { components.Register("extract", newExtract) } // Extract projects a large tool output down to the part relevant to the current @@ -108,7 +114,10 @@ func (e *Extract) reduce(c *components.Ctx, content, goal string, keepIDs []stri cfg := extract.DefaultCfg() cfg.Mode = "code" // rlm is deferred → use the Starlark code strategy cfg.Floor = e.minTokens - if res, _ := extract.RunExtraction(c.Ctx, content, goal, keepIDs, schema.TextTokens(content), cfg, model); res != "" && res != content { + ctx, cancel := context.WithTimeout(c.Ctx, llmCallTimeout) + res, _ := extract.RunExtraction(ctx, content, goal, keepIDs, schema.TextTokens(content), cfg, model) + cancel() + if res != "" && res != content { return res, true } } diff --git a/components/offload/summarize.go b/components/offload/summarize.go index 4a4d8b1..4140b38 100644 --- a/components/offload/summarize.go +++ b/components/offload/summarize.go @@ -79,7 +79,9 @@ func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Re return nil, nil } - summary, err := s.summarize(c.Ctx, model, span) + ctx, cancel := context.WithTimeout(c.Ctx, llmCallTimeout) + defer cancel() + summary, err := s.summarize(ctx, model, span) if err != nil { return nil, err // fail-open: the pipeline reverts this component } diff --git a/deploy/eval-containers/compose.contextguru.yaml b/deploy/eval-containers/compose.contextguru.yaml index ce13d75..e68407c 100644 --- a/deploy/eval-containers/compose.contextguru.yaml +++ b/deploy/eval-containers/compose.contextguru.yaml @@ -30,6 +30,8 @@ services: ANTHROPIC_API_KEY: ${ANTHROPIC_API_KEY:?set to the litellm token} # pipeline selection: a non-empty list wins, else the preset is used CONTEXT_GURU_PIPELINE: ${CONTEXT_GURU_PIPELINE:-} + # full config document (per-component config, e.g. extract strategy/model) + CONTEXT_GURU_CONFIG_YAML: ${CONTEXT_GURU_CONFIG_YAML:-} CONTEXT_GURU_PRESET: ${CONTEXT_GURU_PRESET:-off} CONTEXT_GURU_DEBUG: ${CONTEXT_GURU_DEBUG:-} # when set, the gateway appends a before→after record per rewritten message diff --git a/deploy/eval-containers/start b/deploy/eval-containers/start index 5ab8437..763897b 100755 --- a/deploy/eval-containers/start +++ b/deploy/eval-containers/start @@ -38,6 +38,13 @@ export LISTEN_ADDR="${LISTEN_ADDR:-:${PORT:-4000}}" # wins — it renders a one-line config so a sweep can pin any component combo # (including a single component, or an empty list for a passthrough baseline) # without new presets. Otherwise fall back to a named preset. +# A full config document (CONTEXT_GURU_CONFIG_YAML) wins over everything — it lets +# a sweep pin per-component config (e.g. extract strategy: code, model.source). +if [ -n "${CONTEXT_GURU_CONFIG_YAML:-}" ]; then + printf '%s' "${CONTEXT_GURU_CONFIG_YAML}" > /tmp/cg-config.yaml + exec /opt/gateway/main --config /tmp/cg-config.yaml +fi + # Non-empty CONTEXT_GURU_PIPELINE wins; an empty/unset value falls through to the # named preset (so an always-present-but-empty env var from compose doesn't force # passthrough — baseline uses CONTEXT_GURU_PRESET=off instead). diff --git a/deploy/eval-containers/sweep.py b/deploy/eval-containers/sweep.py index 44b192b..68c5a92 100755 --- a/deploy/eval-containers/sweep.py +++ b/deploy/eval-containers/sweep.py @@ -58,7 +58,19 @@ ("cg-balanced", "claude-code", "cg:preset=balanced"), ("rtk", "claude-code-rtk", "cg:off"), ("headroom", "claude-code", "headroom"), + # LLM-based (model.source: incoming -> same model as the agent, claude-sonnet-4-6) + ("cg-extract-code", "claude-code", "yaml:extract-code"), + ("cg-summarize", "claude-code", "yaml:summarize"), ] + +# Full config documents for LLM-based configs (per-component config the simple +# CONTEXT_GURU_PIPELINE list can't express). incoming source = the agent's own model. +FULL_CONFIGS = { + # higher floors than the deterministic default: an LLM call only pays off on + # genuinely large outputs, and keeps calls/task bounded. + "extract-code": "pipeline: [extract]\ncomponents:\n extract: {strategy: code, min_tokens: 1500, model: {source: incoming}}\n", + "summarize": "pipeline: [summarize]\ncomponents:\n summarize: {start_from_message: 6, keep_last: 3, min_tokens: 1500, model: {source: incoming}}\n", +} FIELDS = ["task", "config", "agent", "reward", "passed", "wall_s", "gw_requests", "gw_before", "gw_after", "gw_saved", "gw_pct", "note"] @@ -133,6 +145,8 @@ def run_cell(task, name, agent, selector, base, token): env["CONTEXT_GURU_PRESET"] = selector.split("=", 1)[1] elif selector.startswith("cg:"): env["CONTEXT_GURU_PIPELINE"] = selector[3:] + elif selector.startswith("yaml:"): + env["CONTEXT_GURU_CONFIG_YAML"] = FULL_CONFIGS[selector[5:]] sh(f"docker compose -p {proj} down -v") up = sh(f"docker compose -p {proj} {files} up -d", cwd=SWE, env=env) From 341cff4cb810d322e921b1c455883d0ea3db3e95 Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Mon, 13 Jul 2026 11:01:28 +0300 Subject: [PATCH 4/5] feat(components): gate LLM components with triggers + reuse prior compactions Make the LLM-based components (summarize, extract strategy: code) run only when worthwhile and reuse prior work, and make extract-code apply to any output shape. - Trigger: new components.Trigger{min_request_tokens,min_messages,min_output_tokens} embedded in summarize/extract so they fire on request shape (tokens/steps/large output), not every turn. Zero fields = fire always (back-compat). - State reuse: extract caches each reduced output by content hash (re-sent output reuses it, no new model call, byte-identical); summarize checkpoints its summary per session and reuses it until the tail grows past resummarize_tokens. - Generalize extract-code: IsContained string case now accepts an in-order whole-line subsequence, so logs/source/tracebacks reduce, not only JSON; prompt rewritten domain-agnostic (JSON + raw-text examples). - Full context: both components use conversationGoal (task + recent turns), not just the last message; summarizer is told to summarize toward the task. - Reliability: raise call timeouts (extract 60s, summarize 150s) and cap the summarizer trajectory to fit the model context instead of failing open. - Tests: Trigger.Fires, line-subsequence containment, extract result-cache reuse, summarize checkpoint reuse. Docs: components.md + RESULTS.md (tuned 10-task sweep; resolved 6/6/6, zero reward regression; extract-code ~193k, summarize ~15.6k saved). Assisted-By: Claude Signed-off-by: Osher-Elhadad --- .gitignore | 2 +- components/all/reuse_test.go | 114 +++++++++++++++++++++++++ components/offload/common.go | 53 ++++++++++++ components/offload/extract.go | 50 +++++++++-- components/offload/state.go | 74 +++++++++++++++++ components/offload/summarize.go | 137 ++++++++++++++++++++++++++----- components/trigger.go | 37 +++++++++ components/trigger_test.go | 41 +++++++++ deploy/eval-containers/sweep.py | 16 +++- docs/RESULTS.md | 47 +++++++++++ docs/components.md | 47 ++++++++--- internal/extract/contain.go | 36 +++++++- internal/extract/contain_test.go | 31 +++++++ internal/extract/prompt.go | 57 ++++++++----- 14 files changed, 670 insertions(+), 72 deletions(-) create mode 100644 components/all/reuse_test.go create mode 100644 components/offload/state.go create mode 100644 components/trigger.go create mode 100644 components/trigger_test.go create mode 100644 internal/extract/contain_test.go diff --git a/.gitignore b/.gitignore index 71040df..a6e5138 100644 --- a/.gitignore +++ b/.gitignore @@ -170,4 +170,4 @@ cython_debug/ coverage.txt # benchmark run output (results summarized in docs/PR, not tracked) -deploy/eval-containers/sweep-results.csv +deploy/eval-containers/sweep-results*.csv diff --git a/components/all/reuse_test.go b/components/all/reuse_test.go new file mode 100644 index 0000000..e80e3dd --- /dev/null +++ b/components/all/reuse_test.go @@ -0,0 +1,114 @@ +package all_test + +import ( + "context" + "strings" + "sync" + "testing" + + "github.com/kagenti/context-guru/components" + "github.com/kagenti/context-guru/schema" + "github.com/kagenti/context-guru/store" + bschemas "github.com/maximhq/bifrost/core/schemas" +) + +// countingModel is a stubModel that records how many times it was called, so a +// test can prove a second turn reused prior state instead of re-calling the LLM. +type countingModel struct { + resp string + mu sync.Mutex + calls int +} + +func (m *countingModel) Complete(context.Context, string) (string, error) { + m.mu.Lock() + m.calls++ + m.mu.Unlock() + return m.resp, nil +} + +func sysMsg(s string) bschemas.ChatMessage { + t := s + return bschemas.ChatMessage{Role: bschemas.ChatMessageRoleSystem, Content: &bschemas.ChatMessageContent{ContentStr: &t}} +} + +// TestSummarizeReusesCheckpoint: once a summary exists, a later turn whose covered +// prefix is unchanged and whose new tail is small must REUSE the prior summary — +// no second model call, and the summary message byte-identical (KV-cache stable). +func TestSummarizeReusesCheckpoint(t *testing.T) { + off := newComp(t, "summarize", "keep_last: 1\nstart_from_message: 0\nmin_tokens: 1\nresummarize_tokens: 100000\n") + st := store.NewMemory(store.Options{}) + cm := &countingModel{resp: "essential facts"} + tool := toolMsg(strings.Repeat("verbose tool output line\n", 40)) + base := func() []bschemas.ChatMessage { + return []bschemas.ChatMessage{sysMsg("you are helpful"), tool, userMsg("continue")} + } + run := func(msgs []bschemas.ChatMessage) *bschemas.BifrostChatRequest { + req := &bschemas.BifrostChatRequest{Input: msgs} + c := &components.Ctx{Ctx: context.Background(), Session: "sess1", Store: st, + Model: components.ModelSpec{Incoming: cm}} + var rep components.Report + if _, err := off.Offload(req, &rep, c); err != nil { + t.Fatal(err) + } + return req + } + + req1 := run(base()) + if cm.calls != 1 { + t.Fatalf("turn 1 must summarize once, calls=%d", cm.calls) + } + sum1 := schema.MessageText(req1.Input[1]) + + // Turn 2: the client re-sends the full original transcript plus new messages. + turn2 := append(base(), userMsg("next question"), toolMsg("small follow-up output")) + req2 := run(turn2) + if cm.calls != 1 { + t.Fatalf("turn 2 must REUSE the summary (no new model call), calls=%d", cm.calls) + } + if got := schema.MessageText(req2.Input[1]); got != sum1 { + t.Fatalf("reused summary must be byte-identical:\n turn1=%q\n turn2=%q", sum1, got) + } + // The new tail is retained verbatim after the summary. + if schema.MessageText(req2.Input[len(req2.Input)-1]) != "small follow-up output" { + t.Fatal("newly appended tail must be kept verbatim on reuse") + } +} + +// TestExtractReusesResultCache: the same large tool output re-sent on a later turn +// reuses the prior compaction — no second model call — and is still reduced. +func TestExtractReusesResultCache(t *testing.T) { + off := newComp(t, "extract", "strategy: code\nmin_tokens: 1\nmodel:\n source: config\n") + st := store.NewMemory(store.Options{}) + filter := "data = json.decode(INPUT)\nOUTPUT = json.encode([r for r in data if \"keep\" in r[\"name\"]])\n" + cm := &countingModel{resp: filter} + body := `[{"id":1,"name":"keep this"},{"id":2,"name":"drop this"},{"id":3,"name":"keep that"}]` + run := func() *bschemas.BifrostChatRequest { + req := &bschemas.BifrostChatRequest{Input: []bschemas.ChatMessage{ + userMsg("find the keep records"), toolMsg(body), + }} + c := &components.Ctx{Ctx: context.Background(), Session: "sX", Store: st, + Model: components.ModelSpec{Static: cm}} + var rep components.Report + if _, err := off.Offload(req, &rep, c); err != nil { + t.Fatal(err) + } + return req + } + + req1 := run() + if cm.calls != 1 { + t.Fatalf("turn 1 must call the model once, calls=%d", cm.calls) + } + if strings.Contains(schema.MessageText(req1.Input[1]), "drop this") { + t.Fatal("turn 1 should have reduced the output") + } + + req2 := run() + if cm.calls != 1 { + t.Fatalf("turn 2 must REUSE the cached extraction (no new model call), calls=%d", cm.calls) + } + if strings.Contains(schema.MessageText(req2.Input[1]), "drop this") { + t.Fatal("turn 2 reused reduction must still drop the non-keep record") + } +} diff --git a/components/offload/common.go b/components/offload/common.go index 3d02bb1..4770ce7 100644 --- a/components/offload/common.go +++ b/components/offload/common.go @@ -33,6 +33,59 @@ func lastUserText(req *bschemas.BifrostChatRequest) string { return "" } +// goalCap bounds the conversational context handed to an LLM component so a huge +// task statement can't blow up every prompt. Generous — the point is to pass the +// real task, not one trailing sentence. +const goalCap = 8000 + +// conversationGoal is the relevance/context signal for the LLM components: the +// TASK the agent is working (first user turn — in agent traffic this holds the +// problem statement), plus the most recent assistant and user turns (its current +// intent). Passing this instead of just the last message is what lets extract / +// summarize keep what actually matters, on any agent or benchmark. Tool outputs +// are excluded — they are the bulk being reduced, not the goal. +func conversationGoal(req *bschemas.BifrostChatRequest) string { + var firstUser, lastUser, lastAsst string + for i := range req.Input { + if req.Input[i].Role == bschemas.ChatMessageRoleUser { + firstUser = schema.MessageText(req.Input[i]) + break + } + } + for i := len(req.Input) - 1; i >= 0; i-- { + switch req.Input[i].Role { + case bschemas.ChatMessageRoleUser: + if lastUser == "" { + lastUser = schema.MessageText(req.Input[i]) + } + case bschemas.ChatMessageRoleAssistant: + if lastAsst == "" { + lastAsst = schema.MessageText(req.Input[i]) + } + } + if lastUser != "" && lastAsst != "" { + break + } + } + var parts []string + seen := map[string]struct{}{} + for _, p := range []string{firstUser, lastAsst, lastUser} { + if p = strings.TrimSpace(p); p == "" { + continue + } + if _, dup := seen[p]; dup { + continue + } + seen[p] = struct{}{} + parts = append(parts, p) + } + g := strings.Join(parts, "\n\n") + if len(g) > goalCap { + g = g[:goalCap] + } + return g +} + // keywords extracts lowercased content words (>3 chars) as a set — a cheap // relevance signal without embeddings. func keywords(s string) map[string]struct{} { diff --git a/components/offload/extract.go b/components/offload/extract.go index ec043bc..4dc55a3 100644 --- a/components/offload/extract.go +++ b/components/offload/extract.go @@ -13,9 +13,11 @@ import ( "gopkg.in/yaml.v3" ) -// llmCallTimeout bounds a single in-request model call so a slow/hung model -// fails open (the component falls back / reverts) instead of stalling the agent. -const llmCallTimeout = 30 * time.Second +// llmCallTimeout bounds a single in-request extract model call so a slow/hung +// model fails open (the component falls back / reverts) instead of stalling the +// agent. extract can make several calls per request (one per large tool output), +// so this ceiling stays modest; the per-content result cache keeps repeats free. +const llmCallTimeout = 60 * time.Second func init() { components.Register("extract", newExtract) } @@ -37,6 +39,7 @@ type Extract struct { tail int strategy string modelSource string + trigger components.Trigger } type extractConfig struct { @@ -47,6 +50,7 @@ type extractConfig struct { Model struct { Source string `yaml:"source"` // incoming (default) | config } `yaml:"model"` + Trigger components.Trigger `yaml:"trigger"` } func newExtract(raw []byte) (components.Component, error) { @@ -56,7 +60,21 @@ func newExtract(raw []byte) (components.Component, error) { return nil, err } } - return &Extract{minTokens: cfg.MinTokens, head: cfg.Head, tail: cfg.Tail, strategy: cfg.Strategy, modelSource: cfg.Model.Source}, nil + // Legacy min_tokens is the per-output floor; the canonical knob is + // trigger.min_output_tokens. Fold one into the other so both work. + if cfg.Trigger.MinOutputTokens == 0 { + cfg.Trigger.MinOutputTokens = cfg.MinTokens + } + return &Extract{minTokens: cfg.MinTokens, head: cfg.Head, tail: cfg.Tail, strategy: cfg.Strategy, modelSource: cfg.Model.Source, trigger: cfg.Trigger}, nil +} + +// outputFloor is the minimum tokens a single tool output must have to be worth +// offloading (the "large output" trigger). +func (e *Extract) outputFloor() int { + if e.trigger.MinOutputTokens > 0 { + return e.trigger.MinOutputTokens + } + return e.minTokens } func (Extract) Name() string { return "extract" } @@ -66,7 +84,14 @@ func (Extract) Enabled(*components.Ctx) bool { return true } func (e *Extract) NeedsModel() bool { return e.strategy == "code" || e.strategy == "rlm" } func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Report, c *components.Ctx) ([]string, error) { - goal := lastUserText(req) + // Request-level trigger: for the LLM strategies, don't spend a model call + // until the request is genuinely large / deep. Deterministic runs always + // (it's cheap). Zero thresholds fire always (backward compatible). + if e.NeedsModel() && !e.trigger.Fires(req) { + rep.Skipped = true + return nil, nil + } + goal := conversationGoal(req) // full task + recent turns, not one trailing sentence query := keywords(goal) if len(query) == 0 { rep.Skipped = true @@ -76,6 +101,7 @@ func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Repo if e.NeedsModel() { model = c.Model.For(e.modelSource) } + floor := e.outputFloor() keepIDs := extract.HarvestIdentifiers(goal, 40) var keys []string for _, i := range toolIndices(req) { @@ -84,13 +110,21 @@ func (e *Extract) Offload(req *bschemas.BifrostChatRequest, rep *components.Repo continue // non-text blocks would be dropped by a text rewrite } content := schema.MessageText(*msg) - if content == "" || schema.TextTokens(content) < e.minTokens { + if content == "" || schema.TextTokens(content) < floor { continue } if len(expand.ParseMarkers(content)) > 0 { continue } - projected, ok := e.reduce(c, content, goal, keepIDs, query, model) + // Reuse a prior compaction of this exact output (marker/whitespace + // insensitive) — no LLM call, and the same bytes keep the prefix stable. + id := extract.ContentKey(content) + projected, ok := "", false + if cached, hit := getResult(c, id); hit { + projected, ok = string(cached), true + } else if projected, ok = e.reduce(c, content, goal, keepIDs, query, model); ok { + putResult(c, id, []byte(projected)) + } if !ok || schema.TextTokens(projected) >= schema.TextTokens(content) { continue } @@ -113,7 +147,7 @@ func (e *Extract) reduce(c *components.Ctx, content, goal string, keepIDs []stri if model != nil && e.NeedsModel() { cfg := extract.DefaultCfg() cfg.Mode = "code" // rlm is deferred → use the Starlark code strategy - cfg.Floor = e.minTokens + cfg.Floor = e.outputFloor() ctx, cancel := context.WithTimeout(c.Ctx, llmCallTimeout) res, _ := extract.RunExtraction(ctx, content, goal, keepIDs, schema.TextTokens(content), cfg, model) cancel() diff --git a/components/offload/state.go b/components/offload/state.go new file mode 100644 index 0000000..b2356be --- /dev/null +++ b/components/offload/state.go @@ -0,0 +1,74 @@ +package offload + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + + "github.com/kagenti/context-guru/components" + bschemas "github.com/maximhq/bifrost/core/schemas" +) + +// State reuse over the generic Store (key→bytes), session-scoped by key prefix +// so a prior compaction is reused across turns instead of re-calling the LLM. +// Reusing the previous output byte-for-byte also keeps the request prefix stable +// (KV-cache friendly) — re-deriving a different summary/extraction each turn is +// both costly and cache-hostile. + +// resultKey namespaces a per-content reduced output (extract) by session. +func resultKey(session, id string) string { return "cg:res:" + session + ":" + id } + +// getResult returns a previously cached reduced output for content id, if any. +func getResult(c *components.Ctx, id string) ([]byte, bool) { + return c.Store.Get(resultKey(c.Session, id)) +} + +// putResult caches a reduced output so a later turn re-sending the same content +// reuses it (no LLM call, byte-identical result). +func putResult(c *components.Ctx, id string, v []byte) { + c.Store.Put(resultKey(c.Session, id), v) +} + +// sumCheckpoint is the per-session summarize state: the exact summary message +// text produced last time (re-emitted verbatim so the prefix stays byte-stable), +// how many leading span messages it subsumed, a hash of that span to prove the +// prefix is unchanged before reusing, and the stash key of the summarized span +// (its original is refreshed in the Store on reuse so expand keeps working). +type sumCheckpoint struct { + SummaryMsg string `json:"m"` + CoveredCount int `json:"c"` + CoveredHash string `json:"h"` + Key string `json:"k"` +} + +func sumKey(session string) string { return "cg:sum:" + session } + +func loadCheckpoint(c *components.Ctx) (sumCheckpoint, bool) { + b, ok := c.Store.Get(sumKey(c.Session)) + if !ok { + return sumCheckpoint{}, false + } + var cp sumCheckpoint + if json.Unmarshal(b, &cp) != nil { + return sumCheckpoint{}, false + } + return cp, true +} + +func saveCheckpoint(c *components.Ctx, cp sumCheckpoint) { + if b, err := json.Marshal(cp); err == nil { + c.Store.Put(sumKey(c.Session), b) + } +} + +// spanHash is a stable content hash of a message span, used to confirm the +// covered prefix is unchanged on a later turn before reusing the summary. +func spanHash(span []bschemas.ChatMessage) string { + h := sha256.New() + for i := range span { + b, _ := json.Marshal(span[i]) + h.Write(b) + h.Write([]byte{0}) + } + return hex.EncodeToString(h.Sum(nil))[:24] +} diff --git a/components/offload/summarize.go b/components/offload/summarize.go index 4140b38..68396be 100644 --- a/components/offload/summarize.go +++ b/components/offload/summarize.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "strings" + "time" "github.com/kagenti/context-guru/components" "github.com/kagenti/context-guru/expand" @@ -14,6 +15,18 @@ import ( func init() { components.Register("summarize", newSummarize) } +// summarizeCallTimeout bounds the single summarizer call. It is higher than +// extract's ceiling because summarize makes ONE call over a large span, and a +// big trajectory legitimately takes the model longer to read and compress. +const summarizeCallTimeout = 150 * time.Second + +// maxTrajectoryChars caps the trajectory text sent to the summarizer so a very +// large span still fits the model's context window (≈70k tokens, well under a +// 200k window with room for the summary) instead of erroring and failing open. +// The most recent turns are kept; the full original span is always stashed for +// expand, so the older prefix is never lost — only left out of THIS summary. +const maxTrajectoryChars = 280_000 + // Summarize compresses the middle of a long trajectory into one LLM-written // summary (ported from CE-Manager's ReSum-style summarizer). It restructures the // message list to [msg0, , last-K messages], replacing @@ -24,35 +37,48 @@ func init() { components.Register("summarize", newSummarize) } // This is the one component that changes the message count; apply.Body rebuilds // the body preserving the retained messages' original bytes. type Summarize struct { - level string - keepLast int - startFrom int - minTokens int - includeToolCalls bool - modelSource string + level string + keepLast int + minTokens int + resummarizeTokens int + includeToolCalls bool + modelSource string + trigger components.Trigger } type summarizeConfig struct { - SummaryLevel string `yaml:"summary_level"` // concise | regular | highly_detailed - KeepLast int `yaml:"keep_last"` // messages kept verbatim at the tail - StartFrom int `yaml:"start_from_message"` // no-op until the list reaches this length - MinTokens int `yaml:"min_tokens"` // min content tokens in the span to bother - IncludeToolCalls bool `yaml:"include_tool_calls"` - Model struct { + SummaryLevel string `yaml:"summary_level"` // concise | regular | highly_detailed + KeepLast int `yaml:"keep_last"` // messages kept verbatim at the tail + StartFrom int `yaml:"start_from_message"` // legacy: folds into trigger.min_messages + MinTokens int `yaml:"min_tokens"` // min content tokens in the span to bother + // ResummarizeTokens: once a summary exists, reuse it (no LLM call) until the + // un-summarized tail since the last checkpoint grows past this many tokens, + // then roll the checkpoint forward with a fresh summary. 0 = re-summarize + // every eligible turn (old behavior). + ResummarizeTokens int `yaml:"resummarize_tokens"` + IncludeToolCalls bool `yaml:"include_tool_calls"` + Model struct { Source string `yaml:"source"` // incoming (default) | config } `yaml:"model"` + Trigger components.Trigger `yaml:"trigger"` } func newSummarize(raw []byte) (components.Component, error) { - cfg := summarizeConfig{SummaryLevel: "regular", KeepLast: 3, StartFrom: 6, MinTokens: 500} + cfg := summarizeConfig{SummaryLevel: "regular", KeepLast: 3, StartFrom: 6, MinTokens: 500, ResummarizeTokens: 6000} if len(raw) > 0 { if err := yaml.Unmarshal(raw, &cfg); err != nil { return nil, err } } + // Legacy start_from_message is a message-count gate; the canonical knob is + // trigger.min_messages. Fold one into the other so both work. + if cfg.Trigger.MinMessages == 0 { + cfg.Trigger.MinMessages = cfg.StartFrom + } return &Summarize{ - level: cfg.SummaryLevel, keepLast: cfg.KeepLast, startFrom: cfg.StartFrom, - minTokens: cfg.MinTokens, includeToolCalls: cfg.IncludeToolCalls, modelSource: cfg.Model.Source, + level: cfg.SummaryLevel, keepLast: cfg.KeepLast, + minTokens: cfg.MinTokens, resummarizeTokens: cfg.ResummarizeTokens, + includeToolCalls: cfg.IncludeToolCalls, modelSource: cfg.Model.Source, trigger: cfg.Trigger, }, nil } @@ -64,7 +90,9 @@ func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Re msgs := req.Input // Keep msg0 (system/first) + the last keepLast; summarize the span between. start, end := 1, len(msgs)-s.keepLast - if len(msgs) < s.startFrom || end <= start { + // Request-level trigger: don't summarize (an LLM call) until the transcript + // is genuinely large / deep. Zero thresholds fire always (back-compat). + if !s.trigger.Fires(req) || end <= start { rep.Skipped = true return nil, nil } @@ -73,15 +101,25 @@ func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Re rep.Skipped = true // NeedsModel but none available → degrade gracefully return nil, nil } + + // Reuse a prior summary if the covered prefix is unchanged and the tail since + // that checkpoint is still small — no LLM call, and the summary message stays + // byte-identical (KV-cache stable). Roll the checkpoint forward only once the + // tail grows past resummarize_tokens. + if out, keys, ok := s.tryReuse(c, msgs, start, end); ok { + req.Input = out + return keys, nil + } + span := msgs[start:end] if schema.MessagesTokens(&bschemas.BifrostChatRequest{Input: span}) < s.minTokens { rep.Skipped = true return nil, nil } - ctx, cancel := context.WithTimeout(c.Ctx, llmCallTimeout) + ctx, cancel := context.WithTimeout(c.Ctx, summarizeCallTimeout) defer cancel() - summary, err := s.summarize(ctx, model, span) + summary, err := s.summarize(ctx, model, span, conversationGoal(req)) if err != nil { return nil, err // fail-open: the pipeline reverts this component } @@ -98,8 +136,16 @@ func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Re key := hashKey(string(spanJSON)) c.Store.Put(key, spanJSON) + summaryText := summaryWrapper(summary, key) summaryMsg := bschemas.ChatMessage{Role: bschemas.ChatMessageRoleSystem} - schema.SetMessageText(&summaryMsg, summaryWrapper(summary, key)) + schema.SetMessageText(&summaryMsg, summaryText) + + // Checkpoint: this summary subsumes the leading span (len(span) messages from + // index 1). A later turn appends messages, so this same prefix stays stable. + saveCheckpoint(c, sumCheckpoint{ + SummaryMsg: summaryText, CoveredCount: end - start, + CoveredHash: spanHash(span), Key: key, + }) // [msg0, summary, last-K] — reassign; apply.Body rebuilds losslessly. out := make([]bschemas.ChatMessage, 0, 2+s.keepLast) @@ -109,13 +155,55 @@ func (s *Summarize) Offload(req *bschemas.BifrostChatRequest, rep *components.Re return []string{key}, nil } +// tryReuse re-emits the previous summary if (1) a checkpoint exists, (2) the +// covered prefix (msgs[1:1+CoveredCount]) is byte-unchanged, and (3) the tail +// since that boundary is below resummarize_tokens. It returns the rebuilt +// [msg0, priorSummary, msgs[boundary:]] and the (refreshed) stash key. No LLM +// call. ok=false means "re-summarize fresh". +func (s *Summarize) tryReuse(c *components.Ctx, msgs []bschemas.ChatMessage, start, end int) ([]bschemas.ChatMessage, []string, bool) { + if s.resummarizeTokens <= 0 { + return nil, nil, false + } + cp, ok := loadCheckpoint(c) + if !ok || cp.CoveredCount <= 0 { + return nil, nil, false + } + boundary := start + cp.CoveredCount + if boundary > end { // covered prefix would overlap the kept tail — can't reuse + return nil, nil, false + } + covered := msgs[start:boundary] + if spanHash(covered) != cp.CoveredHash { + return nil, nil, false // prefix diverged (different session / edited) → fresh + } + // The un-summarized middle since the checkpoint (excludes the kept last-K). + if schema.MessagesTokens(&bschemas.BifrostChatRequest{Input: msgs[boundary:end]}) >= s.resummarizeTokens { + return nil, nil, false // grown enough — roll the checkpoint forward + } + // Refresh the stashed original span so expand keeps resolving it. + if b, err := json.Marshal(covered); err == nil { + c.Store.Put(cp.Key, b) + } + summaryMsg := bschemas.ChatMessage{Role: bschemas.ChatMessageRoleSystem} + schema.SetMessageText(&summaryMsg, cp.SummaryMsg) + out := make([]bschemas.ChatMessage, 0, 2+(len(msgs)-boundary)) + out = append(out, msgs[0], summaryMsg) + out = append(out, msgs[boundary:]...) + return out, []string{cp.Key}, true +} + // summarize builds the trajectory string and asks the model once (bounded retry). -func (s *Summarize) summarize(ctx context.Context, model components.Model, span []bschemas.ChatMessage) (string, error) { +// goal is the current task + recent turns, so the summary is grounded in what the +// agent is actually trying to do — not a blind digest of the middle. +func (s *Summarize) summarize(ctx context.Context, model components.Model, span []bschemas.ChatMessage, goal string) (string, error) { sys := summarizerSystemPrompt if !s.includeToolCalls { sys += summarizerMaskedNote } user := strings.Replace(summarizerUserPrompt, "{trajectory}", trajectoryString(span, s.includeToolCalls), 1) + if g := strings.TrimSpace(goal); g != "" { + user = "CURRENT TASK / QUESTION (summarize toward this):\n" + g + "\n\n" + user + } if suffix := summaryLevelSuffix[s.level]; suffix != "" { user += "\n" + suffix } @@ -149,7 +237,14 @@ func trajectoryString(span []bschemas.ChatMessage, includeToolCalls bool) string b.WriteString("]\n") b.WriteString(content) } - return b.String() + s := b.String() + // Keep the trajectory within the model's context window: if it's huge, keep + // the most recent portion (the full original span is stashed for expand). + if len(s) > maxTrajectoryChars { + s = "…[older trajectory omitted from this summary; full original preserved for expand]\n\n" + + s[len(s)-maxTrajectoryChars:] + } + return s } func ensureSummaryTags(s string) string { diff --git a/components/trigger.go b/components/trigger.go new file mode 100644 index 0000000..838d332 --- /dev/null +++ b/components/trigger.go @@ -0,0 +1,37 @@ +package components + +import ( + "github.com/kagenti/context-guru/schema" + "github.com/maximhq/bifrost/core/schemas" +) + +// Trigger is the shared, configurable gate that decides whether an expensive +// (LLM-based) component should ACT on a request — so summarize/extract run only +// when it's worth an LLM call, not on every turn. It is embedded in a +// component's config as `trigger:` and works for any agent/benchmark/use-case +// because the thresholds are pure request shape (tokens and message count), not +// task-specific. +// +// A zero field is "no constraint", so the zero Trigger fires always (backward +// compatible with configs that don't set it). Request-level thresholds +// (MinRequestTokens, MinMessages) are checked by Fires; the per-item +// MinOutputTokens floor is checked by the component against each candidate +// (e.g. extract, per tool output). +type Trigger struct { + MinRequestTokens int `yaml:"min_request_tokens"` // whole request must be at least this many tokens + MinMessages int `yaml:"min_messages"` // …and carry at least this many messages (≈ steps) + MinOutputTokens int `yaml:"min_output_tokens"` // per-item floor: only offload an output at least this big +} + +// Fires reports whether the request-level thresholds are met. Both are ANDed; +// a zero threshold imposes no constraint. It does not consider MinOutputTokens +// (that is a per-item floor the component applies itself). +func (t Trigger) Fires(req *schemas.BifrostChatRequest) bool { + if t.MinMessages > 0 && len(req.Input) < t.MinMessages { + return false + } + if t.MinRequestTokens > 0 && schema.MessagesTokens(req) < t.MinRequestTokens { + return false + } + return true +} diff --git a/components/trigger_test.go b/components/trigger_test.go new file mode 100644 index 0000000..20ed569 --- /dev/null +++ b/components/trigger_test.go @@ -0,0 +1,41 @@ +package components + +import ( + "strings" + "testing" + + "github.com/maximhq/bifrost/core/schemas" +) + +func mkMsg(s string) schemas.ChatMessage { + t := s + return schemas.ChatMessage{Role: schemas.ChatMessageRoleUser, Content: &schemas.ChatMessageContent{ContentStr: &t}} +} + +func TestTriggerFires(t *testing.T) { + // ~1 message with a lot of tokens vs several tiny messages. + big := &schemas.BifrostChatRequest{Input: []schemas.ChatMessage{mkMsg(strings.Repeat("word ", 4000))}} + deep := &schemas.BifrostChatRequest{Input: []schemas.ChatMessage{ + mkMsg("a"), mkMsg("b"), mkMsg("c"), mkMsg("d"), mkMsg("e"), + }} + + cases := []struct { + name string + tr Trigger + req *schemas.BifrostChatRequest + want bool + }{ + {"zero fires always", Trigger{}, deep, true}, + {"token gate met", Trigger{MinRequestTokens: 1000}, big, true}, + {"token gate not met", Trigger{MinRequestTokens: 1000}, deep, false}, + {"message gate met", Trigger{MinMessages: 5}, deep, true}, + {"message gate not met", Trigger{MinMessages: 6}, deep, false}, + {"both gates ANDed (msg fails)", Trigger{MinRequestTokens: 1, MinMessages: 99}, big, false}, + {"MinOutputTokens does not affect Fires", Trigger{MinOutputTokens: 99999}, deep, true}, + } + for _, c := range cases { + if got := c.tr.Fires(c.req); got != c.want { + t.Errorf("%s: Fires=%v want %v", c.name, got, c.want) + } + } +} diff --git a/deploy/eval-containers/sweep.py b/deploy/eval-containers/sweep.py index 68c5a92..f01ffbd 100755 --- a/deploy/eval-containers/sweep.py +++ b/deploy/eval-containers/sweep.py @@ -66,10 +66,18 @@ # Full config documents for LLM-based configs (per-component config the simple # CONTEXT_GURU_PIPELINE list can't express). incoming source = the agent's own model. FULL_CONFIGS = { - # higher floors than the deterministic default: an LLM call only pays off on - # genuinely large outputs, and keeps calls/task bounded. - "extract-code": "pipeline: [extract]\ncomponents:\n extract: {strategy: code, min_tokens: 1500, model: {source: incoming}}\n", - "summarize": "pipeline: [summarize]\ncomponents:\n summarize: {start_from_message: 6, keep_last: 3, min_tokens: 1500, model: {source: incoming}}\n", + # LLM components are gated by a `trigger` (don't fire every turn) and reuse + # prior compactions from session state (don't re-call the model on re-sent + # content). extract fires only on large tool outputs in a large request; + # summarize only on a large/deep transcript, then reuses its summary until the + # un-summarized tail grows past resummarize_tokens. + "extract-code": "pipeline: [extract]\ncomponents:\n" + " extract: {strategy: code, model: {source: incoming},\n" + " trigger: {min_output_tokens: 700, min_request_tokens: 4000}}\n", + "summarize": "pipeline: [summarize]\ncomponents:\n" + " summarize: {keep_last: 3, min_tokens: 1000, resummarize_tokens: 5000,\n" + " model: {source: incoming},\n" + " trigger: {min_request_tokens: 6000, min_messages: 10}}\n", } FIELDS = ["task", "config", "agent", "reward", "passed", "wall_s", "gw_requests", "gw_before", "gw_after", "gw_saved", "gw_pct", "note"] diff --git a/docs/RESULTS.md b/docs/RESULTS.md index 95ae81f..1f7a9f7 100644 --- a/docs/RESULTS.md +++ b/docs/RESULTS.md @@ -176,6 +176,53 @@ AFTER: 425 is_member = (what == 'class' or wh ``` +## LLM-based components (`extract strategy: code`, `summarize`) — gated + state-reuse + +These two components call a model. They are gated by a configurable `trigger` (don't fire every turn) +and reuse prior compactions from per-session state (don't re-call the model on re-sent content). +`model.source: incoming` = the agent's own model (`claude-sonnet-4-6`). Final tuned config: +extract-code `trigger: {min_output_tokens: 700, min_request_tokens: 4000}`; summarize +`trigger: {min_request_tokens: 6000, min_messages: 10}, resummarize_tokens: 5000`, call timeout 150 s, +trajectory capped to fit context. Sweep of both, alone, over the 10 tasks: + +| task | base | extract-code r / saved | summarize r / saved | +|---|---|---|---| +| sympy-13647 | 1 | 1 / **2562** | 1 / 0 (gated off) | +| sympy-16766 | 1 | 1 / 0 | 1 / 0 (gated off — was a 1→0 regression) | +| sympy-20438 | 0 | 0 / **177914** (~24%) | 0 / **10882** | +| sphinx-7910 | 1 | 1 / 0 | 1 / 0 (gated off) | +| sphinx-9320 | 1 | 1 / **5677** | 1 / 0 (gated off — was a 1→0 regression) | +| scikit-12973 | 1 | 1 / 0 | 1 / 0 | +| scikit-25931 | 1 | 1 / **2952** | 1 / 0 | +| django-11820 | 0 | 0 / **3920** | 0 / **4695** | +| django-14089 | 0 | 0 / 0 | 0 / 0 | +| xarray-4629 | – | 0 / 0 | 0 / 0 | + +**Resolved / 10 (baseline = 6): extract-code = 6, summarize = 6 — zero reward regression.** The two +earlier summarize regressions (sympy-16766, sphinx-9320: 1→0 in the first ungated run) stay fixed. +Total gateway tokens saved across the set: **extract-code ≈ 193k, summarize ≈ 15.6k**. + +- **extract-code** fires on 5 of the tasks and is reward-neutral everywhere. Biggest win is the large + task (sympy-20438, ~24% of the request). Two changes drove the breadth: `IsContained`'s string case + was generalized to an in-order whole-line **subsequence** (so it reduces logs / source / tracebacks, + not only JSON — the earlier JSON-only containment was why it fired on almost nothing), and the floor + was lowered to 700 tokens so mid-size outputs qualify. +- **summarize** — the key diagnosis: it saved ~0 at higher triggers because **almost every SWE-bench + Claude Code request is small** (per-request means ≈ 2k tokens; only sympy-20438 averages ~7k). A 15–25k + trigger simply never fired. Lowering it to **6k** threads the needle: it fires on the large, + **reward-0** tasks (sympy-20438, django-11820) where summarizing is safe, and stays off the tiny + (~2k-avg) transcripts whose summarization caused the earlier regressions. That's why summarize now + saves ~15.6k with the two regressions still fixed. Pushing the trigger lower would summarize the tiny + tasks too and re-introduce the reward risk — the real unlock for safe savings there is expand-tool + injection (reversibility), still deferred. +- **State reuse** (proved by unit tests): extract caches each reduced output by content hash and + summarize checkpoints its summary per session, so a re-sent output / unchanged prefix reuses the prior + result byte-for-byte — no repeat model call, KV-cache-stable prefix. Stops both re-deriving a + different compaction every turn. +- **Fundamental tension** confirmed empirically: on this traffic, summarize only saves where it fires, + and firing on small transcripts is exactly what regresses reward. The gain came from firing it on the + *large, already-failing* tasks — not from accepting regressions on the small ones. + ## Methodology & caveats - Harness: [`deploy/eval-containers/sweep.py`](../deploy/eval-containers/sweep.py) (resumable) → `aggregate.py`. diff --git a/docs/components.md b/docs/components.md index af94f1d..b595f34 100644 --- a/docs/components.md +++ b/docs/components.md @@ -16,11 +16,11 @@ messages (`role:"tool"`; for Anthropic, `tool_result` blocks normalized to that | `collapse` | Offload | middle of an oversized output | via expand | any large tool output (fallback) | `max_tokens` (2000), `head_lines` (20), `tail_lines` (20) | | `failed_run` | Offload | earlier superseded test/build runs | via expand | ≥2 run-like outputs | `min_tokens` (100) | | `cmdfilter` | Offload | lines per declarative DSL filter | via expand | output matching a filter | `filters` ([]), `disable_builtins` (false) | -| `extract` | Offload | query-irrelevant lines (or an LLM-written filter) | via expand | large output + a recent query | `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` (deterministic\|code\|rlm), `model.source` | +| `extract` | Offload | query-irrelevant lines (or an LLM-written filter) | via expand | large output + a recent query | `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` (deterministic\|code\|rlm), `model.source`, `trigger` | | `smartcrush` | Offload | middle items of a JSON array | via expand | JSON-array tool output | `min_items` (5), `min_tokens` (200), `keep_first` (3), `keep_last` (2) | | `mask` | Offload | older tool outputs (age-based) | via expand | more than `keep_recent` outputs | `keep_recent` (3), `min_tokens` (100) | | `phi_evict` | Offload | lowest-scoring outputs over budget | via expand | transcript over `budget_tokens` | `budget_tokens` (120000), `weights` (balanced) | -| `summarize` | Offload (LLM) | the middle of the transcript → one summary | via expand | long trajectories | `summary_level` (regular), `keep_last` (3), `start_from_message` (6), `min_tokens` (500), `model.source` | +| `summarize` | Offload (LLM) | the middle of the transcript → one summary | via expand | long trajectories | `summary_level` (regular), `keep_last` (3), `min_tokens` (500), `resummarize_tokens` (6000), `model.source`, `trigger` | Presets (`config`): `off` `[]` · `safe` `[format, cacheinject]` · `balanced` `[format, dedup, failed_run, cmdfilter, cacheinject]` · `aggressive` adds `smartcrush, extract` · @@ -164,14 +164,25 @@ after: … lines mentioning auth/timeout + any error lines … <> [full output: call context_guru_expand] ``` +The relevance signal is the **full conversational context** (the task in the first user turn + +the most recent assistant/user turns), not just one trailing sentence — so it keeps what the agent +actually needs on any agent/benchmark. + - **Config:** `min_tokens` (300), `head_lines`/`tail_lines` (5), `strategy` - (`deterministic` | `code` | `rlm`), `model.source`. With `code`, a cheap LLM writes a Starlark - `extract_relevant_data`-style filter (`data = json.decode(INPUT)` → `OUTPUT = json.encode(...)`) - run in a sandbox (no imports/IO, step + 2s limits); the result is accepted only if a **containment** - check proves it's a lossless subset (else fall back to deterministic). `rlm` currently maps to - `code`. **Shines:** big query-focused MCP/API outputs; `code` shines on structured JSON where a - filter can select records precisely. **Inert:** no user query, output < `min_tokens`, projection - not smaller, or (for `code`) no model available → deterministic. + (`deterministic` | `code` | `rlm`), `model.source`, `trigger`. With `code`, a cheap LLM writes a + Starlark filter run in a sandbox (no imports/IO, step + 2s limits); the result is accepted only if a + **containment** check proves it's a lossless subset (else fall back to deterministic). The `code` + strategy is **domain-agnostic**: JSON bodies are decoded and filtered by shape; **raw text** (logs, + source, tracebacks, search results) is kept as an in-order **line subset** — containment accepts a + whole-line subsequence, not just a contiguous substring. `rlm` currently maps to `code`. +- **Gating + reuse (LLM strategies):** a `trigger` decides whether to spend a model call — + `min_output_tokens` (per tool output; folds in legacy `min_tokens`), `min_request_tokens`, + `min_messages` — so `code` fires only on a large output in a large request, not every turn. A reduced + output is cached per session by content hash, so the **same output re-sent on a later turn reuses the + prior compaction** (no new model call, byte-identical result → prefix stays KV-cache stable). +- **Shines:** big query-focused MCP/API outputs, logs and file reads; structured JSON where a filter + selects records precisely. **Inert:** no user query, output below floor, request below `trigger`, + projection not smaller, or (for `code`) no model available → deterministic. ### `smartcrush` Statistical JSON-**array** compressor: parse the array, keep `keep_first` + `keep_last` items plus @@ -225,11 +236,21 @@ before: [system, u1, tool, a1, tool, u2, … 30 turns …, uN-1, uN] after: [system, "=== History Summary === … … <>", uN-1, uN] ``` +The summarizer is grounded in the **current task** (first user turn + recent turns are passed as +"summarize toward this"), not a blind digest of the middle. + - **Config:** `summary_level` (`concise`|`regular`|`highly_detailed`), `keep_last` (3), - `start_from_message` (6 — no-op until the list is this long), `min_tokens` (500), `include_tool_calls` - (false → tool outputs masked in the trajectory), `model.source`. **Shines:** long agentic sessions - where the bulk is stale middle context. **Inert:** short transcripts, span below `min_tokens`, or no - model available (no-op). Run it **alone** (its own preset) — it restructures the whole transcript. + `min_tokens` (500 — span floor), `include_tool_calls` (false → tool outputs masked in the + trajectory), `model.source`, `trigger`, `resummarize_tokens` (6000). +- **Gating + reuse:** a `trigger` (`min_request_tokens`, `min_messages`; legacy `start_from_message` + folds into `min_messages`) gates the first summary so it fires only on a large/deep transcript. + After that, the summary is **checkpointed per session** and **reused verbatim** (no model call, and + byte-identical so the prefix stays KV-cache stable) until the un-summarized tail grows past + `resummarize_tokens`, when the checkpoint rolls forward with a fresh summary. This is what stops it + re-summarizing every turn. +- **Shines:** long agentic sessions where the bulk is stale middle context. **Inert:** transcript + below `trigger`, span below `min_tokens`, or no model available (no-op). Run it **alone** (its own + preset) — it restructures the whole transcript. --- diff --git a/internal/extract/contain.go b/internal/extract/contain.go index 08d6f9f..0e684a0 100644 --- a/internal/extract/contain.go +++ b/internal/extract/contain.go @@ -16,9 +16,12 @@ import ( ) // IsContained reports whether result is a lossless projection of original (both -// parsed values): string → substring; list → order-preserving subsequence of -// contained items; dict → keys subset with contained values; numbers/bools/nil → -// equal where present; nil result → always allowed (dropping is fine). +// parsed values): string → contiguous substring OR an order-preserving +// subsequence of whole lines (so a log/code/traceback reduction that drops whole +// lines still proves lossless — every kept line appears verbatim, in order); +// list → order-preserving subsequence of contained items; dict → keys subset +// with contained values; numbers/bools/nil → equal where present; nil result → +// always allowed (dropping is fine). func IsContained(result, original any) bool { return checkContained(result, original) } @@ -33,7 +36,7 @@ func checkContained(out, in any) bool { switch o := out.(type) { case string: s, ok := in.(string) - return ok && strings.Contains(s, o) + return ok && (strings.Contains(s, o) || linesSubsequence(o, s)) case bool: b, ok := in.(bool) return ok && b == o @@ -82,6 +85,31 @@ func checkContained(out, in any) bool { } } +// linesSubsequence reports whether every line of out appears, in order and +// byte-identical, as a line of in — i.e. out is in with whole lines dropped. +// This is the text analogue of the list-subsequence rule: a lossless projection +// that keeps whole lines verbatim (logs, source, tracebacks, search results). +func linesSubsequence(out, in string) bool { + outL := strings.Split(out, "\n") + inL := strings.Split(in, "\n") + i := 0 + for _, ol := range outL { + matched := false + for i < len(inL) { + if inL[i] == ol { + matched = true + i++ + break + } + i++ + } + if !matched { + return false + } + } + return true +} + // numbersEqual compares two JSON-decoded numbers (json.Number or float64) by value. func numbersEqual(a, b any) bool { return numStr(a) == numStr(b) diff --git a/internal/extract/contain_test.go b/internal/extract/contain_test.go new file mode 100644 index 0000000..f2fe01a --- /dev/null +++ b/internal/extract/contain_test.go @@ -0,0 +1,31 @@ +package extract + +import "testing" + +// TestContainmentLineSubsequence: a text reduction that drops whole lines (logs, +// code, tracebacks) is a valid lossless projection — it must pass containment +// even though it is NOT a contiguous substring of the original. +func TestContainmentLineSubsequence(t *testing.T) { + orig := "line one\nERROR: boom in widget.c\nnoise\nmore noise\nline five" + keep := "ERROR: boom in widget.c\nline five" // in-order subset of whole lines + + if !IsContained(keep, orig) { + t.Fatal("in-order whole-line subset must be contained") + } + // Contiguous substring still works (single kept region). + if !IsContained("noise\nmore noise", orig) { + t.Fatal("contiguous substring must be contained") + } + // Reordered lines are NOT a subsequence → rejected. + if IsContained("line five\nERROR: boom in widget.c", orig) { + t.Fatal("reordered lines must be rejected") + } + // An edited line (not byte-identical) is rejected. + if IsContained("ERROR: boom in widget.cpp", orig) { + t.Fatal("an altered line must be rejected") + } + // Fabricated line absent from the original is rejected. + if IsContained("totally invented line", orig) { + t.Fatal("a fabricated line must be rejected") + } +} diff --git a/internal/extract/prompt.go b/internal/extract/prompt.go index 18b33e8..2174d56 100644 --- a/internal/extract/prompt.go +++ b/internal/extract/prompt.go @@ -37,8 +37,8 @@ func buildPrompt(bodyText, goal string, keepIDs []string) string { if g == "" { g = "(no explicit goal stated)" } - if len(g) > 2000 { - g = g[:2000] + if len(g) > 8000 { + g = g[:8000] } keep := keepIDs if len(keep) > 60 { @@ -68,32 +68,47 @@ func buildPrompt(bodyText, goal string, keepIDs []string) string { // codeRules is the Starlark code-writing contract: the model writes a program that // runs over the real INPUT (it is NOT shown the full body), so the prompt stays cheap. // Same "select, never summarize, recall-first" discipline as buildPrompt, retargeted -// from "return the value" to "write the filter". +// from "return the value" to "write the filter". It is domain-agnostic: it works for +// ANY tool output (JSON records, logs, source files, tracebacks, search results, +// tables, command output) — not one benchmark or one shape. const codeRules = `Write a Starlark program (a safe Python subset) that filters this ONE tool output -down to only what the agent needs next. Contract: +down to only the parts the agent needs next. Works for ANY output — JSON, logs, +source code, search results, tracebacks, tables, command output. +Contract: - The global string INPUT holds the FULL tool output. The module ` + "`json`" + ` is available. -- Start: data = json.decode(INPUT) -- Select a SMALLER value of the SAME shape (e.g. drop irrelevant list records / fields). -- End: OUTPUT = json.encode(result) # OUTPUT must be a string +- Assign a string global OUTPUT holding the KEPT content, same kind as INPUT. +- If INPUT is JSON (it starts with { or [): result = json.decode(INPUT), keep a + SMALLER value of the SAME shape (drop irrelevant records/fields), then + OUTPUT = json.encode(result). +- Otherwise treat INPUT as TEXT: lines = INPUT.split("\n"); keep only the relevant + lines (an in-order subset — never reorder or edit them); OUTPUT = "\n".join(kept). Rules, in priority order: -1. RECALL FIRST. When unsure whether a record/field is relevant, KEEP IT. -2. SELECT, NEVER SUMMARIZE. Keep whole records/values byte-for-byte. Never paraphrase, - truncate, round, reformat, or invent values — only DROP clearly-irrelevant ones. -3. PRESERVE EXACTLY: ids, numbers, names, paths, timestamps, errors, stack traces — - and anything matching the KEEP list. -4. NO imports (no load()), NO I/O, NO network. Use only json.decode/json.encode and - plain Starlark (list comprehensions, dict/list ops, string ops). -5. If you cannot identify clearly-irrelevant content, set OUTPUT = INPUT. +1. RECALL FIRST. When unsure whether a line/record is relevant, KEEP IT. +2. SELECT, NEVER SUMMARIZE. Keep whole lines/records byte-for-byte. Never paraphrase, + truncate, reword, renumber, reformat, or invent — only DROP clearly-irrelevant ones. +3. PRESERVE EXACTLY: ids, numbers, names, paths, signatures, timestamps, error + messages, stack traces — and anything matching the KEEP list. +4. KEEP anything mentioning the goal or a KEEP identifier. DROP unrelated boilerplate: + banners, progress/spinner lines, repeated separators, unrelated files/records/tests. +5. NO imports (no load()), NO I/O, NO network. Use only json.decode/json.encode and + plain Starlark (list comprehensions, dict/list/string ops, .split/.join/in). +6. If you cannot identify clearly-irrelevant content, set OUTPUT = INPUT. Output ONLY the Starlark program — no prose, no markdown fences.` -const codeExample = `EXAMPLE -Goal: "Fix failing test test_auth_expiry; find the relevant hit." -KEEP: ["test_auth_expiry","auth/session.py"] +const codeExample = `EXAMPLE A (JSON records) +Goal: "Fix failing test test_auth_expiry." KEEP: ["test_auth_expiry","auth/session.py"] INPUT schema: list of {"path": str, "snippet": str} PROGRAM: data = json.decode(INPUT) result = [r for r in data if "auth" in r["path"] or "test_auth_expiry" in r["snippet"]] -OUTPUT = json.encode(result)` +OUTPUT = json.encode(result) + +EXAMPLE B (raw text / log) +Goal: "Why does build fail?" KEEP: ["compile_module","widget.c"] +PROGRAM: +lines = INPUT.split("\n") +kept = [ln for ln in lines if "error" in ln.lower() or "widget.c" in ln or "compile_module" in ln] +OUTPUT = "\n".join(kept)` // buildCodePrompt builds the prompt for the Starlark code-writing strategy. Unlike // buildPrompt it does NOT inline the full body — the program runs over the real INPUT. @@ -103,8 +118,8 @@ func buildCodePrompt(bodyText, goal string, keepIDs []string) string { if g == "" { g = "(no explicit goal stated)" } - if len(g) > 2000 { - g = g[:2000] + if len(g) > 8000 { + g = g[:8000] } keep := keepIDs if len(keep) > 60 { From f83aad58e0ff018ee44390cf9ac2002b0cb09240 Mon Sep 17 00:00:00 2001 From: Osher-Elhadad Date: Mon, 13 Jul 2026 11:07:07 +0300 Subject: [PATCH 5/5] docs(results): commit tuned SWE-bench sweep snapshot Preserve the full 10-task sweep behind docs/RESULTS.md as a dated snapshot (the live sweep-results.csv is gitignored scratch). One row per (task, config) with reward + gateway token-savings stats; LLM configs use the final tuned triggers. Adds a README documenting columns, config, and headline (resolved 6/6/6, extract-code ~193k / summarize ~15.6k tokens saved). Assisted-By: Claude Signed-off-by: Osher-Elhadad --- deploy/eval-containers/results/README.md | 27 ++++ .../results/swe-bench-sweep-2026-07-13.csv | 150 ++++++++++++++++++ 2 files changed, 177 insertions(+) create mode 100644 deploy/eval-containers/results/README.md create mode 100644 deploy/eval-containers/results/swe-bench-sweep-2026-07-13.csv diff --git a/deploy/eval-containers/results/README.md b/deploy/eval-containers/results/README.md new file mode 100644 index 0000000..f3922ef --- /dev/null +++ b/deploy/eval-containers/results/README.md @@ -0,0 +1,27 @@ +# SWE-bench sweep results (committed snapshots) + +The live `../sweep-results.csv` is gitignored (scratch, overwritten each run). This +directory holds dated snapshots of interesting runs so the numbers behind +[`docs/RESULTS.md`](../../../docs/RESULTS.md) are reproducible and preserved. + +## `swe-bench-sweep-2026-07-13.csv` + +Full sweep of every config over the 10-task set, via the eval-containers compose +stack against the IBM litellm upstream (`EVAL_MODEL=anthropic/claude-sonnet-4-6`, +`claude-code` agent). One row per (task, config). + +Columns: `task, config, agent, reward, passed, wall_s, gw_requests, gw_before, +gw_after, gw_saved, gw_pct, note`. `gw_*` are the gateway `/stats` rollups +(`tokens_before/after`, `saved_tokens`, `savings_pct`) for that cell. + +LLM-based configs use the final tuned triggers (`model.source: incoming`): + +- `cg-extract-code` — `extract {strategy: code, trigger: {min_output_tokens: 700, min_request_tokens: 4000}}` +- `cg-summarize` — `summarize {trigger: {min_request_tokens: 6000, min_messages: 10}, resummarize_tokens: 5000}` + +Headline (10 tasks): resolved 6/6/6 (baseline / extract-code / summarize) — zero +reward regression. extract-code reduced ~23% of its request tokens overall +(≈193k saved, fired on 5 tasks); summarize ≈4.8% (≈15.6k, fired on the 2 large +reward-0 tasks). Deterministic configs (`cg-mask`, `cg-dedup`, …) are included for +comparison. `pydata__xarray-4629` / `scikit-learn__scikit-learn-25931` runner images +are flaky under arm64 emulation; some cells show `up-failed`. diff --git a/deploy/eval-containers/results/swe-bench-sweep-2026-07-13.csv b/deploy/eval-containers/results/swe-bench-sweep-2026-07-13.csv new file mode 100644 index 0000000..0e7e6d1 --- /dev/null +++ b/deploy/eval-containers/results/swe-bench-sweep-2026-07-13.csv @@ -0,0 +1,150 @@ +task,config,agent,reward,passed,wall_s,gw_requests,gw_before,gw_after,gw_saved,gw_pct,note +sympy__sympy-13647,baseline,claude-code,1,True,67,,,,,,exited +sympy__sympy-13647,cg-balanced,claude-code,1,True,91,14,31738,27109,4629,14.585040015123827,exited +sympy__sympy-16766,baseline,claude-code,1,True,59,13,29156,29156,0,0,exited +sympy__sympy-16766,cg-balanced,claude-code,1,True,71,12,25660,22037,3623,14.119251753702262,exited +sympy__sympy-20438,baseline,claude-code,0,False,333,56,377181,377181,0,0,exited +sympy__sympy-20438,cg-balanced,claude-code,0,False,633,74,865706,843051,22655,2.616939238032311,exited +sphinx-doc__sphinx-7910,baseline,claude-code,1,True,56,10,10792,10792,0,0,exited +sphinx-doc__sphinx-7910,cg-balanced,claude-code,1,True,129,20,53155,42010,11145,20.9669833505785,exited +sphinx-doc__sphinx-9320,baseline,claude-code,1,True,53,7,29732,29732,0,0,exited +sphinx-doc__sphinx-9320,cg-balanced,claude-code,1,True,59,9,28698,28698,0,0,exited +scikit-learn__scikit-learn-12973,baseline,claude-code,1,True,28,4,5408,5408,0,0,exited +scikit-learn__scikit-learn-12973,cg-balanced,claude-code,1,True,36,4,5321,5321,0,0,exited +scikit-learn__scikit-learn-25931,baseline,claude-code,1,True,138,26,94707,94707,0,0,exited +scikit-learn__scikit-learn-25931,cg-balanced,claude-code,1,True,156,33,165339,126587,38752,23.437906362080334,exited +pydata__xarray-4629,cg-balanced,claude-code,0,False,29,3,4170,4170,0,0,exited +django__django-11820,baseline,claude-code,0,False,200,40,165659,165659,0,0,exited +django__django-11820,cg-balanced,claude-code,0,False,454,43,183022,157275,25747,14.067707707270163,exited +django__django-14089,baseline,claude-code,0,False,22,5,3042,3042,0,0,exited +django__django-14089,cg-balanced,claude-code,0,False,22,5,3495,3495,0,0,exited +sympy__sympy-13647,cg-format,claude-code,1,True,74,13,33021,33021,0,0,exited +sympy__sympy-13647,cg-dedup,claude-code,1,True,71,14,36309,35769,540,1.4872345699413367,exited +sympy__sympy-13647,cg-cmdfilter,claude-code,1,True,82,13,27092,27092,0,0,exited +sympy__sympy-13647,cg-cacheinject,claude-code,1,True,58,12,23364,23364,0,0,exited +sympy__sympy-13647,cg-failed_run,claude-code,1,True,86,13,33161,30176,2985,9.001537951207744,exited +sympy__sympy-13647,cg-skeleton,claude-code,1,True,80,13,38517,38517,0,0,exited +sympy__sympy-13647,cg-collapse,claude-code,1,True,76,15,32376,32376,0,0,exited +sympy__sympy-13647,cg-mask,claude-code,1,True,63,12,25583,20563,5020,19.62240550365477,exited +sympy__sympy-13647,cg-smartcrush,claude-code,1,True,65,11,26028,26028,0,0,exited +sympy__sympy-13647,cg-extract,claude-code,1,True,70,13,33813,28766,5047,14.926211812024961,exited +sympy__sympy-13647,cg-phi_evict,claude-code,1,True,79,14,28990,28990,0,0,exited +sympy__sympy-16766,cg-format,claude-code,1,True,60,14,20819,20819,0,0,exited +sympy__sympy-16766,cg-dedup,claude-code,1,True,59,10,19567,19567,0,0,exited +sympy__sympy-16766,cg-cmdfilter,claude-code,1,True,70,18,32571,32571,0,0,exited +sympy__sympy-16766,cg-cacheinject,claude-code,1,True,49,10,14569,14569,0,0,exited +sympy__sympy-16766,cg-failed_run,claude-code,1,True,95,15,26615,21368,5247,19.714446740559836,exited +sympy__sympy-16766,cg-skeleton,claude-code,1,True,82,21,58806,58806,0,0,exited +sympy__sympy-16766,cg-collapse,claude-code,1,True,82,20,50464,50464,0,0,exited +sympy__sympy-16766,cg-mask,claude-code,1,True,87,20,40868,23329,17539,42.916218067926,exited +sympy__sympy-16766,cg-smartcrush,claude-code,1,True,89,19,41281,41281,0,0,exited +sympy__sympy-16766,cg-extract,claude-code,1,True,96,23,62346,59201,3145,5.044429474224489,exited +sympy__sympy-16766,cg-phi_evict,claude-code,1,True,93,20,44091,44091,0,0,exited +sympy__sympy-20438,cg-format,claude-code,0,False,342,57,505437,505437,0,0,exited +sympy__sympy-20438,cg-dedup,claude-code,0,False,461,54,606521,606521,0,0,exited +sympy__sympy-20438,cg-cmdfilter,claude-code,0,False,251,35,219532,219532,0,0,exited +sympy__sympy-20438,cg-cacheinject,claude-code,0,False,1352,67,454999,454999,0,0,exited +sympy__sympy-20438,cg-failed_run,claude-code,0,False,375,54,435977,429882,6095,1.398009528025561,exited +sympy__sympy-20438,cg-skeleton,claude-code,0,False,253,28,128738,128738,0,0,exited +sympy__sympy-20438,cg-collapse,claude-code,0,False,1047,65,487908,457608,30300,6.210187166433016,exited +sympy__sympy-20438,cg-mask,claude-code,0,False,1833,226,27812982,1800277,26012705,93.52720610828426,exited +sympy__sympy-20438,cg-smartcrush,claude-code,0,False,705,64,492014,492014,0,0,exited +sympy__sympy-20438,cg-extract,claude-code,0,False,517,76,1320726,940115,380611,28.81831659254077,exited +sympy__sympy-20438,cg-phi_evict,claude-code,0,False,276,48,332596,332596,0,0,exited +sphinx-doc__sphinx-7910,cg-format,claude-code,1,True,60,8,7715,7715,0,0,exited +sphinx-doc__sphinx-7910,cg-dedup,claude-code,1,True,108,13,25166,25166,0,0,exited +sphinx-doc__sphinx-7910,cg-cmdfilter,claude-code,1,True,74,14,16898,16898,0,0,exited +sphinx-doc__sphinx-7910,cg-cacheinject,claude-code,1,True,119,19,34121,34121,0,0,exited +sphinx-doc__sphinx-7910,cg-failed_run,claude-code,1,True,126,18,28968,27415,1553,5.3610880972107156,exited +sphinx-doc__sphinx-7910,cg-skeleton,claude-code,1,True,60,9,9299,9299,0,0,exited +sphinx-doc__sphinx-7910,cg-collapse,claude-code,1,True,106,18,30295,30295,0,0,exited +sphinx-doc__sphinx-7910,cg-mask,claude-code,1,True,109,16,25701,20818,4883,18.999260729154507,exited +sphinx-doc__sphinx-7910,cg-smartcrush,claude-code,1,True,125,15,32625,32625,0,0,exited +sphinx-doc__sphinx-7910,cg-extract,claude-code,1,True,109,20,92688,74757,17931,19.345546349041946,exited +sphinx-doc__sphinx-7910,cg-phi_evict,claude-code,1,True,111,14,19957,19957,0,0,exited +sphinx-doc__sphinx-9320,cg-format,claude-code,1,True,52,10,27047,27047,0,0,exited +sphinx-doc__sphinx-9320,cg-dedup,claude-code,1,True,38,5,19215,19215,0,0,exited +sphinx-doc__sphinx-9320,cg-cmdfilter,claude-code,1,True,47,8,36290,36290,0,0,exited +sphinx-doc__sphinx-9320,cg-cacheinject,claude-code,1,True,37,7,19861,19861,0,0,exited +sphinx-doc__sphinx-9320,cg-failed_run,claude-code,1,True,53,9,27473,27473,0,0,exited +sphinx-doc__sphinx-9320,cg-skeleton,claude-code,1,True,45,9,41193,41193,0,0,exited +sphinx-doc__sphinx-9320,cg-collapse,claude-code,1,True,46,7,29431,29431,0,0,exited +sphinx-doc__sphinx-9320,cg-mask,claude-code,1,True,59,11,25116,15323,9793,38.99108138238573,exited +sphinx-doc__sphinx-9320,cg-smartcrush,claude-code,1,True,38,4,14237,14237,0,0,exited +sphinx-doc__sphinx-9320,cg-extract,claude-code,1,True,61,11,53881,51859,2022,3.7527143148790856,exited +sphinx-doc__sphinx-9320,cg-phi_evict,claude-code,1,True,37,7,19605,19605,0,0,exited +scikit-learn__scikit-learn-12973,cg-format,claude-code,1,True,31,4,5374,5374,0,0,exited +scikit-learn__scikit-learn-12973,cg-dedup,claude-code,1,True,35,4,5321,5321,0,0,exited +scikit-learn__scikit-learn-12973,cg-cmdfilter,claude-code,1,True,33,4,4209,4209,0,0,exited +scikit-learn__scikit-learn-12973,cg-cacheinject,claude-code,1,True,32,4,5382,5382,0,0,exited +scikit-learn__scikit-learn-12973,cg-failed_run,claude-code,1,True,32,4,5341,5341,0,0,exited +scikit-learn__scikit-learn-12973,cg-skeleton,claude-code,1,True,36,4,5218,5218,0,0,exited +scikit-learn__scikit-learn-12973,cg-collapse,claude-code,1,True,30,4,4610,4610,0,0,exited +scikit-learn__scikit-learn-12973,cg-mask,claude-code,1,True,33,4,5397,5397,0,0,exited +scikit-learn__scikit-learn-12973,cg-smartcrush,claude-code,1,True,30,4,5325,5325,0,0,exited +scikit-learn__scikit-learn-12973,cg-extract,claude-code,1,True,42,8,15031,13343,1688,11.23012440955359,exited +scikit-learn__scikit-learn-12973,cg-phi_evict,claude-code,1,True,35,4,5539,5539,0,0,exited +scikit-learn__scikit-learn-25931,cg-format,claude-code,1,True,129,18,58110,58110,0,0,exited +scikit-learn__scikit-learn-25931,cg-dedup,claude-code,1,True,154,27,99187,99187,0,0,exited +scikit-learn__scikit-learn-25931,cg-cmdfilter,claude-code,1,True,96,18,48742,48742,0,0,exited +scikit-learn__scikit-learn-25931,cg-cacheinject,claude-code,1,True,130,15,43833,43833,0,0,exited +scikit-learn__scikit-learn-25931,cg-failed_run,claude-code,1,True,91,14,37584,37279,305,0.8115155385270328,exited +scikit-learn__scikit-learn-25931,cg-skeleton,claude-code,1,True,106,18,61961,61961,0,0,exited +scikit-learn__scikit-learn-25931,cg-collapse,claude-code,1,True,121,20,69774,69774,0,0,exited +scikit-learn__scikit-learn-25931,cg-mask,claude-code,1,True,155,26,94545,56192,38353,40.565868105135124,exited +scikit-learn__scikit-learn-25931,cg-smartcrush,claude-code,1,True,132,20,78496,78496,0,0,exited +scikit-learn__scikit-learn-25931,cg-extract,claude-code,1,True,123,24,100608,88388,12220,12.146151399491094,exited +scikit-learn__scikit-learn-25931,cg-phi_evict,claude-code,1,True,140,23,80245,80245,0,0,exited +pydata__xarray-4629,cg-format,claude-code,0,False,25,3,4194,4194,0,0,exited +pydata__xarray-4629,cg-dedup,claude-code,0,False,21,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-cmdfilter,claude-code,0,False,22,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-cacheinject,claude-code,0,False,22,3,4195,4195,0,0,exited +pydata__xarray-4629,cg-failed_run,claude-code,0,False,22,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-skeleton,claude-code,0,False,24,3,4188,4188,0,0,exited +pydata__xarray-4629,cg-collapse,claude-code,0,False,29,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-mask,claude-code,0,False,33,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-smartcrush,claude-code,0,False,24,3,4170,4170,0,0,exited +pydata__xarray-4629,cg-extract,claude-code,0,False,23,3,4195,4195,0,0,exited +pydata__xarray-4629,cg-phi_evict,claude-code,0,False,26,3,4170,4170,0,0,exited +django__django-11820,cg-format,claude-code,0,False,113,18,40347,40347,0,0,exited +django__django-11820,cg-dedup,claude-code,0,False,172,30,125676,123840,1836,1.46089945574334,exited +django__django-11820,cg-cmdfilter,claude-code,0,False,93,17,44145,44145,0,0,exited +django__django-11820,cg-cacheinject,claude-code,0,False,118,18,55743,55743,0,0,exited +django__django-11820,cg-failed_run,claude-code,0,False,219,36,175236,110833,64403,36.752151384418724,exited +django__django-11820,cg-skeleton,claude-code,0,False,233,29,152725,152725,0,0,exited +django__django-11820,cg-collapse,claude-code,0,False,168,26,134458,134458,0,0,exited +django__django-11820,cg-mask,claude-code,0,False,116,17,43981,21978,22003,50.02842136377071,exited +django__django-11820,cg-smartcrush,claude-code,0,False,240,52,240944,240944,0,0,exited +django__django-11820,cg-extract,claude-code,0,False,218,40,274684,222188,52496,19.111415299034526,exited +django__django-11820,cg-phi_evict,claude-code,0,False,191,22,102560,102560,0,0,exited +django__django-14089,cg-format,claude-code,0,False,23,5,3289,3289,0,0,exited +django__django-14089,cg-dedup,claude-code,0,False,34,6,4259,4259,0,0,exited +django__django-14089,cg-cmdfilter,claude-code,0,False,25,6,4259,4259,0,0,exited +django__django-14089,cg-cacheinject,claude-code,0,False,28,6,4259,4259,0,0,exited +django__django-14089,cg-failed_run,claude-code,0,False,23,4,2575,2575,0,0,exited +django__django-14089,cg-skeleton,claude-code,0,False,22,5,3289,3289,0,0,exited +django__django-14089,cg-collapse,claude-code,0,False,25,5,3337,3337,0,0,exited +django__django-14089,cg-mask,claude-code,0,False,27,6,4259,4080,179,4.202864522188307,exited +django__django-14089,cg-smartcrush,claude-code,0,False,23,5,3495,3495,0,0,exited +django__django-14089,cg-extract,claude-code,0,False,22,5,3337,3337,0,0,exited +django__django-14089,cg-phi_evict,claude-code,0,False,38,6,4259,4259,0,0,exited +sympy__sympy-13647,cg-extract-code,claude-code,1,True,86,13,38416,35854,2562,6.669096209912537,exited +sympy__sympy-16766,cg-extract-code,claude-code,1,True,88,15,36053,36053,0,0,exited +sympy__sympy-20438,cg-extract-code,claude-code,0,False,446,55,527010,349096,177914,33.75913170528074,exited +sphinx-doc__sphinx-7910,cg-extract-code,claude-code,1,True,138,20,35362,35362,0,0,exited +sphinx-doc__sphinx-9320,cg-extract-code,claude-code,1,True,64,9,41567,35890,5677,13.657468665046792,exited +scikit-learn__scikit-learn-12973,cg-extract-code,claude-code,1,True,39,4,5339,5339,0,0,exited +scikit-learn__scikit-learn-25931,cg-extract-code,claude-code,1,True,118,16,51236,48284,2952,5.761573893356234,exited +pydata__xarray-4629,cg-extract-code,claude-code,0,False,24,3,4194,4194,0,0,exited +django__django-11820,cg-extract-code,claude-code,0,False,192,31,85915,81995,3920,4.562649129954025,exited +django__django-14089,cg-extract-code,claude-code,0,False,21,4,2417,2417,0,0,exited +sympy__sympy-13647,cg-summarize,claude-code,1,True,74,11,25174,25174,0,0,exited +sympy__sympy-16766,cg-summarize,claude-code,1,True,79,17,26227,26227,0,0,exited +sympy__sympy-20438,cg-summarize,claude-code,0,False,136,12,49819,38937,10882,21.843071920351672,exited +sphinx-doc__sphinx-7910,cg-summarize,claude-code,1,True,119,17,23283,23283,0,0,exited +sphinx-doc__sphinx-9320,cg-summarize,claude-code,1,True,57,8,36406,36406,0,0,exited +scikit-learn__scikit-learn-12973,cg-summarize,claude-code,1,True,39,4,4425,4425,0,0,exited +scikit-learn__scikit-learn-25931,cg-summarize,claude-code,1,True,123,22,83659,83659,0,0,exited +pydata__xarray-4629,cg-summarize,claude-code,0,False,25,3,4170,4170,0,0,exited +django__django-11820,cg-summarize,claude-code,0,False,152,21,67331,62636,4695,6.973013916323834,exited +django__django-14089,cg-summarize,claude-code,0,False,27,5,3337,3337,0,0,exited