feat: LLM-based context components (summarize + extract code strategy)

.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

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)

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, <summary>, 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.<i> 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)

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, <summary>, 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, "<<cg:") {
+		t.Fatalf("summary message missing wrapper/marker: %q", s)
+	}
+}
+
 func TestNoMessagesForwardsUnchanged(t *testing.T) {
 	cfg := pipe(t, "pipeline: [dedup]\n")
 	p, _ := cfg.Build(nil)

cmd/context-guru-proxy/main.go

@@ -16,8 +16,10 @@ import (
 	"net/http"
 	"os"
 
+	"github.com/kagenti/context-guru/components"
 	_ "github.com/kagenti/context-guru/components/all"
 	"github.com/kagenti/context-guru/config"
+	"github.com/kagenti/context-guru/internal/cheapmodel"
 	"github.com/kagenti/context-guru/metrics"
 	"github.com/kagenti/context-guru/proxy"
 )
@@ -52,6 +54,7 @@ func main() {
 		OpenAIKey:    os.Getenv("OPENAI_API_KEY"),
 		AnthropicKey: os.Getenv("ANTHROPIC_API_KEY"),
 		ForceModel:   os.Getenv("FORCE_MODEL"), // eval-containers pins EVAL_MODEL's model here
+		CheapModel:   cheapModelFromEnv(),      // static "config"-source LLM for NeedsModel components
 	})
 
 	slog.Info("context-guru-proxy listening", "addr", addr, "pipeline", cfg.Pipeline)
@@ -73,3 +76,32 @@ func envOr(key, def string) string {
 	}
 	return def
 }
+
+// cheapModelFromEnv builds the static "config"-source LLM client for NeedsModel
+// components (extract code/rlm, summarize with model.source=config). Returns nil
+// when CHEAP_MODEL is unset, so those components fall back / no-op.
+//
+//	CHEAP_MODEL           model id (e.g. claude-haiku-4-5); unset => 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"),
+		}
+	}
+}