Proposal: UCB-based library suggestion selection (explore/exploit)

[ianphil]

Context

Right now, when the planner builds a graph, library suggestions that succeed get promoted and reused. The selection is greedy — first match wins, and there's no mechanism to discover that an untried approach might be better. This is the multi-armed bandit problem (Kochenderfer Ch. 15).

The Idea

Add Upper Confidence Bound (UCB) scoring to library suggestion selection. Each library entry tracks successes and attempts. When the planner selects a suggestion, instead of picking the first/best match, it picks the one with the highest UCB score:

score = (successes / attempts) + C * sqrt(ln(total_attempts) / attempts)

The second term is an uncertainty bonus — large when an option hasn't been tried much, shrinking as it gets more attempts. This naturally balances explore (try uncertain options) vs. exploit (use proven ones).

Concrete Example

• mail search --from jake -> 8/10 successes, well-tested, small bonus -> score ~ 0.85
• mail search --filter "from/emailAddress/address eq 'jake@microsoft.com'" -> never tried, huge uncertainty bonus -> score ~ 1.0+
• System tries the untried option once. If it works (1/1), it competes on merit. If it fails (0/1), system goes back to the known option.

Questions to answer before implementing

1. How does library promotion currently work? Where are suggestions stored, how are they matched to planner context, and how are they surfaced to the LLM?
2. Would adding attempts and successes fields to the library model be straightforward?
3. How to feed UCB scores into the planner prompt? Options: (a) rank suggestions by UCB before including them, (b) include the scores and let the LLM factor them in, (c) filter to top-N by UCB score.
4. Where does "success" get recorded? After graph.Ok, or per-task?
5. Interaction with the repair loop — if a suggestion fails on first attempt but succeeds after repair, is that a success or failure for the bandit?
6. The exploration constant C needs tuning. What should it be given typical library size and usage frequency?

Estimated scope

~30 lines of scoring logic if the data model cooperates. The hard part is deciding the success signal and repair loop interaction, not the math.

References

• Kochenderfer, Wheeler, Wray — Algorithms for Decision Making, Ch. 15 (Exploration and Exploitation)
• Related to the capability system collapse (commit bbd4ada) and prompt consolidation (commit 6610034) from the recent refactor

[ianphil]

Read the proposal and traced the code. Short version: UCB fits the architecture cleanly, but the framing ("first match wins, greedy") doesn't match how selection actually happens today — that changes how UCB has to plug in. Concrete answers to the six questions below.

1) How promotion + reuse work today

• Storage — StoreBackedTaskLibrary (Ttasks.ChatApp/Services/Capabilities.cs) persists each library entry as a CoreTask in ITaskStore with metadata taskLibraryItem=true, taskLibraryKey=<key>, plus processFileName / processArgsTemplate / templateParameters. GetOrAdd is key-keyed upsert. UpdateIfChanged only rewrites when definition fields differ (short-circuits otherwise).
• Surface to planner — every turn, Prompts.FormatLibrarySuggestions JSON-serializes the entire library into the planner system prompt under Library suggestions for this turn:. No ranking, no filtering, no host-side scoring.
• Selection — done by the LLM planner, not the host. Planner sets libraryItemKey (+ optional libraryParameters) on a process task. Host renders the template via ResolveLibraryReferences in ChatTurnService and runs it.
• Promotion — ChatTurnService.PromoteLibrarySuggestions only when graph.Ok and the specific plan task's core task Status == Succeeded.

So there is no host-side "first match wins" code. The greediness is the LLM picking a familiar-looking key from a long JSON list. UCB therefore has to inject as ranking + a flag in the prompt, not as a host-side gate.

2) Is attempts / successes straightforward to add?

Mostly yes:

• Add Attempts / Successes to TaskLibraryDefinition + TaskLibraryItem, persist as metadata (taskLibraryAttempts, taskLibrarySuccesses).
• UpdateIfChanged currently early-returns when "same" — need a separate RecordOutcome(key, succeeded) method that always rewrites counters and bypasses the equality check.

Keep promotion (adding an arm) and outcome recording (a trial) on different code paths so we don't double-count.

3) How to feed UCB into the planner

Don't pick (b). LLMs are bad at arithmetic ranking and you lose the determinism that's the point of UCB. Pick (a) + a small flag:

1. Compute UCB score per item.
2. Sort FormatLibrarySuggestions output descending.
3. Add a single "exploration_candidate": true to the highest-uncertainty item.
4. Do not emit raw scores. The model will fixate on numbers.

Avoid (c) — library items are semantically distinct (different fileName/args). Hiding one means the LLM can't pick it even when it's contextually right. UCB-as-ordering is safe; UCB-as-gate is not.

4) Where does success get recorded?

This is the most underspecified part of the proposal. Two distinct events:

• Reuse outcome = bandit trial. Per-task. Record in ChatTurnService.ExecuteBatchWithRepair, right after graph.Run, before the if (graph.Ok) block. Iterate plan.Tasks where LibraryItemKey is not null, look up the core task by planTaskId metadata (same lookup PromoteLibrarySuggestions uses), and call _library.RecordOutcome(libraryItemKey, succeeded: coreTask.Status == Succeeded).
• Promotion event = adding an arm. Not a trial. Don't credit it as a success.

Per-task — not per-graph. A graph that succeeds with one library call and one fresh process should only credit the library call.

5) Repair loop interaction

ExecuteBatchWithRepair runs up to MaxGraphRepairAttempts + 1 times; each attempt is a fresh _builder.Build(plan) + graph.Run. Per-attempt, per-task accounting is the only honest interpretation:

• Same libraryItemKey reused across attempts → record every attempt's outcome. An item that needs three repairs to land deserves the low success rate it gets.
• Repair changes the key → each one gets its own (attempt, success) tick.
• Attempt N-1 fails with X, attempt N succeeds without X → don't retroactively credit X.

Bad keys lose their lead naturally; success-after-repair credits whichever key was actually used in the winning attempt.

6) C value

Honest answer: vanilla UCB1 assumes IID rewards. Here, success is conditional on user intent and the planner is contextually choosing among arms — this is really a contextual bandit, not vanilla UCB. UCB1 still works as an exploration heuristic; just don't expect optimality guarantees.

With O(10–50) library items and single-digit attempts/week:

• C = sqrt(2) (classical) over-explores at that volume.
• Start at C = 0.5, watch the exploration:exploitation ratio for ~1 week, tune.
• Force attempts == 0 → infinite score so every new candidate gets at least one trial. (The proposal's example implies this.)

Fit assessment

~30 lines is too optimistic, but not by a lot. Realistic delta:

Area	LoC
TaskLibraryDefinition / TaskLibraryItem counter fields + metadata serialization	~15
StoreBackedTaskLibrary.RecordOutcome	~15
ChatTurnService per-task outcome loop	~10
Prompts.FormatLibrarySuggestions UCB sort + exploration_candidate flag	~20
Admin counter view in AdminService	~10
Tests (outcome recording, ranking determinism, zero-attempts floor)	~30–50

Call it ~100 lines + tests. Architecturally clean — fits in two services without touching Ttasks.Core. The riskiest piece is whether reordering the suggestion list materially changes LLM behavior on unrelated turns. Worth a small offline eval before shipping.

Pushback

This is mainly useful if library churn is high enough that new candidates routinely appear and old ones go stale. If the library is small and stable, the exploration term will mostly thrash on a handful of unused items the LLM keeps ignoring. Before investing:

• Measure current library size and weekly reuse frequency.
• Measure rate of new candidates entering via librarySuggestion.
• If <5 new candidates/week and most items get reused predictably, the LLM's contextual selection (reading description + argsTemplate) probably already dominates what UCB would deliver.

If churn is real, the proposal is worth doing. If churn is low, this is yak-shaving — fix the prompt or seeding instead.

— ttasks-net Copilot CLI