Skip to content

feat(optimization): GEPA prompt optimization for templates and crews - #67

Open
nehmetohme wants to merge 10 commits into
databrickslabs:mainfrom
nehmetohme:prompt-optimization
Open

feat(optimization): GEPA prompt optimization for templates and crews#67
nehmetohme wants to merge 10 commits into
databrickslabs:mainfrom
nehmetohme:prompt-optimization

Conversation

@nehmetohme

Copy link
Copy Markdown
Contributor

Summary

Adds a prompt-optimization stack built on MLflow GenAI + GEPA (reflective prompt evolution) that runs entirely inside Kasal — from mining real usage data, through crew-in-the-loop evaluation, to human grading that steers the next optimization round.

Template optimization (Configuration → Prompts)

  • The six generation templates (detect_intent, generate_agent, generate_task, generate_crew, generate_crew_plan, generate_job_name) each get an Optimize (✨) action.
  • Training data is mined from the LLM interaction logs (with junk filtering); scoring combines output-format validity and an LLM judge on correctness.
  • The Prompts configuration surface is consolidated into a single list; optimization opens in a dialog scoped to the chosen template.

Crew optimization (crew catalog → Optimize Prompts)

  • Full crew-in-the-loop GEPA: every metric call executes the crew for real and judges score the actual deliverable. Agent role/goal/backstory and task description/expected_output evolve together via a labeled crew-document serialization; a strict parser scores malformed candidates 0 without spending an execution.
  • Hard execution cap — the budget the user selects is a promise. GEPA evaluates candidates in parallel and can overshoot its metric-call budget, so the predict_fn refuses executions beyond the cap. A Stop button cancels between iterations. Runs show N/M executions live.
  • Apply writes the winning candidate back to the crew's agents and tasks through the repository layer.

Human judgment loop (all inside Kasal, persisted as MLflow assessments)

  • Grade past evaluation answers 0–10 with per-judge attribution, comments, and ground-truth Expectations (mlflow.log_feedback + mlflow.log_expectation on the eval traces).
  • Create custom LLM judges (mlflow.genai.make_judge) in-app; judges live in a shared library and are assigned per crew via a crew-scoped registered copy. A built-in graded 0–10 quality judge is always on.
  • Harvested crew 👍/👎 feedback and human assessments feed both the judge rubric and the GEPA reflection objective — so a non-expert's plain-language feedback autonomously steers prompt evolution.

Infrastructure notes

  • Respects the launch MLFLOW_TRACKING_URI (preserved as KASAL_LAUNCH_MLFLOW_TRACKING_URI before main.py forces databricks); a local OSS registry is only used when explicitly configured, otherwise databricks-uc.
  • Cross-loop safety: optimization worker threads reach the app DB only via run_coroutine_threadsafe onto the main loop with UserContext re-established; crew execution polling uses fresh sessions per tick.
  • MLflow autolog import-hook deadlocks avoided during optimization spans; a 1-token reflection-model preflight fails fast on a dead provider instead of burning budget.
  • New dependency: gepa (required by MLflow's GepaPromptOptimizer). The uv.lock churn beyond the gepa addition is uv re-canonicalization only — no other package versions changed.
  • Known Phase-1 tradeoff (documented in code): run state is kept in an in-memory registry and is lost on restart; a DB-backed run table is planned as a follow-up.

Test plan

  • 25 backend unit tests (tests/unit/services/test_prompt_optimization_service_unit.py): crew-doc serialization round-trips, judge value grading (categorical + numeric), budget accounting, registry gating, apply/cancel flows
  • python -c "from src.api import api_router" imports cleanly
  • Frontend tsc clean; ESLint 0 errors (remaining warnings pre-exist on main)
  • Manually exercised end-to-end locally: template runs with real improvements (e.g. 0.65 → 0.77), crew runs with hard-cap enforcement, stop/cancel, judge creation/assignment/deletion, grading with expectations feeding the next run's reflection objective

🤖 Generated with Claude Code

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

@nehmetohme
nehmetohme force-pushed the prompt-optimization branch from 7f199af to b3d3d3a Compare July 24, 2026 11:13
nehmetohme and others added 9 commits July 24, 2026 13:29
Adds a prompt-optimization stack built on MLflow GenAI + GEPA
(reflective prompt evolution) that works end-to-end inside Kasal:

Template optimization (Configuration -> Prompts):
- Six generation templates (detect_intent, generate_agent, generate_task,
  generate_crew, generate_crew_plan, generate_job_name) get a per-template
  Optimize action that runs GEPA against real usage data mined from the
  LLM interaction logs, with junk filtering and format+correctness scoring.
- Single consolidated "Prompts" surface; optimization opens in a scoped
  dialog next to each template.

Crew optimization (crew catalog -> Optimize Prompts):
- Full crew-in-the-loop GEPA: every metric call executes the crew for real
  and judges score the actual deliverable. Agent role/goal/backstory and
  task description/expected_output are evolved together via a labeled
  crew-document serialization with a strict parser (malformed candidates
  score 0 without spending an execution).
- Hard execution cap: the budget the user picks is a promise; the
  predict_fn refuses executions beyond the cap even when GEPA evaluates
  candidates in parallel. Stop button cancels between iterations.

Human judgment loop (all inside Kasal, persisted as MLflow assessments):
- Grade past evaluation answers 0-10 with per-judge attribution, comments,
  and ground-truth Expectations (log_feedback + log_expectation).
- Create custom LLM judges (make_judge) in-app; judges live in a library
  and are assigned per crew via a crew-scoped registered copy; built-in
  graded 0-10 quality judge always on.
- Harvested crew thumbs and human assessments feed both the judge rubric
  and the GEPA reflection objective, so non-expert feedback autonomously
  steers prompt evolution.

Infrastructure:
- MLflow registry/tracking handling that respects the launch
  MLFLOW_TRACKING_URI (preserved as KASAL_LAUNCH_MLFLOW_TRACKING_URI
  before main.py forces "databricks"), local OSS registry gated behind
  MCP_SERVER_ENABLED, databricks-uc otherwise.
- Cross-loop safety: worker threads reach the app DB only via
  run_coroutine_threadsafe onto the main loop with UserContext
  re-established; crew execution polling uses fresh sessions.
- Autolog import-hook deadlock avoidance during optimization spans;
  reflection-model preflight so a dead provider fails fast instead of
  burning budget.
- New dependency: gepa (required by MLflow's GepaPromptOptimizer).
  uv.lock churn beyond the gepa addition is uv re-canonicalization only;
  no other package versions changed.

25 unit tests cover doc serialization round-trips, judge value grading,
budget accounting, registry gating, and apply/cancel flows.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…make human feedback reach the mutator

Live diagnosis of flat runs (score X -> X, "baseline won") on a real crew
with 12 human grade-0 notes exposed four compounding causes:

1. Budget burned re-measuring the baseline. GEPA re-evaluates the SAME
   candidate doc repeatedly (upfront smoke test, baseline valset pass, and
   a fresh reflective-minibatch pass every iteration). Each re-run cost a
   real crew execution, so a 4-execution run bought exactly ONE distinct
   candidate (observed: total_metric_calls=7, candidates=1). Fix: cache
   deliverables by candidate-doc hash — each DISTINCT candidate costs one
   execution; re-evaluations are free and post-cap re-evals of the
   baseline stay truthful.

2. Judge noise made acceptance a coin flip. The same prompts drew grades
   of 0.0 and then 4/10 two minutes apart; the one candidate that DID
   incorporate the human requirements lost to a lucky baseline draw.
   Fix: cache the judge verdict by deliverable hash — within a run,
   identical outputs always score identically, so candidate-vs-baseline
   comparisons are stable.

3. Reflection was blind to WHY candidates failed. Scorers returned bare
   floats, and MLflow only forwards textual rationales to GEPA's
   reflective dataset from Feedback objects — so the mutator saw "0.0"
   but never "wrong region, rentals instead of sales". The judge was even
   instructed to reply with only the number. Fix: the judge now critiques
   first (quoting the violated requirement) and grades on the last line;
   output_correct returns Feedback(value, rationale) including registered
   judges' rationales; aggregation unwraps Feedback values. Harvested
   human Expectations additionally ride the train row's expectations
   channel, which the reflective dataset surfaces explicitly.

4. Invisible feedback usage + misleading timestamps. Users could not see
   that their grades were harvested ("we are not using the grading"), and
   naive-UTC created_at rendered a 01:20 local run as "11:20 PM". Fix:
   runs now report human_feedback_count ("guided by N human notes" chip)
   and candidates_tried ("N variants tried" chip); timestamps are
   timezone-aware so browsers render local time. Harvest also iterates
   traces oldest-first so the keep-last-12 slice keeps the NEWEST notes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…single-example datasets

A live run WITH the previous fixes still returned "baseline won" while
producing a candidate that fully incorporated the human requirements.
proposals.json showed why: minibatch score 0.9 vs 0.9 — a TIE, rejected
by GEPA's strict-improvement acceptance. Three compounding causes, each
verified against gepa 0.1.4 source and live A/B judge experiments:

1. No gradient: the judge graded EVERY real deliverable 0/10. Feeding
   the raw harvest ("human_grade: 0.0 ..." x13) anchored it — the same
   compliant answer scored 0/10 under the litany rubric and 6/10 under a
   requirements checklist. And raw complaint sentences as checklist
   items made a 30B judge fail every mark by quoting the requirement
   itself as evidence. Now: harvested notes are LLM-distilled once per
   run into <=5 testable imperatives; the judge marks each R<n>
   PASS/FAIL and must quote the violating passage VERBATIM from the
   answer (cannot quote -> PASS), plus a Q quality mark. The grade is
   COMPUTED from the marks (0.8 x pass-fraction + 0.2 x Q), never taken
   from the model's arithmetic — a judge writing "40" previously
   clamped to a perfect 10. Validated live: compliant answer 0.96,
   Geneva-containing answer 0.60 with correct verbatim quotes.

2. gepa defaults hostile to 1-example datasets: reflection_minibatch_size
   defaults to 3, so each proposal step evaluated the SAME single example
   three times — three concurrent predict calls raced past the result
   cache (two crew executions finished the same second) and burned 3
   executions per candidate. Now reflection_minibatch_size=1,
   acceptance_criterion="improvement_or_equal" (lateral moves survive
   flat regions), cache_evaluation=True (gepa-side result cache skips
   repeat metric calls entirely), and a lock serializes
   check-then-execute so concurrent identical calls cannot double-run
   the crew.

3. Metric budget conflated with execution budget: cached re-evaluations
   still consumed GEPA's metric-call budget, so a 10-execution run
   stopped after 4 executions with budget left. The optimizer now gets
   metric headroom (2x executions + 3) while the user's number remains
   a hard cap on real crew executions.

Also: judge context no longer includes the crew objective line (the
task text said "cities like Zurich, Geneva" while the human demanded
German-side only — the contradiction primed hallucinated FAILs), and
eval-trace logging failures are logged at warning (a baseline eval
vanished silently, leaving nothing to grade).

11 new unit tests: requirement distillation/parsing and checklist-mark
grading (including the model-arithmetic and duplicate-mark hazards).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ic reflection dead-loops the search

With everything else fixed, a live run still returned "baseline won"
after spending only 2 of 10 executions: proposals.json showed ELEVEN
byte-identical proposals (same 2268 chars, same 0.571 -> 0.356 scores),
each a free cache-hit rejection until the metric budget drained.

The reflection prompt is identical every iteration in this regime —
same parent (the baseline never gets displaced), same single training
example, same cached judge rationale — so a reflection endpoint with a
deterministic default (temperature ~0) regenerates the exact same
candidate forever. The result caches introduced earlier correctly made
those repeats free, which turned determinism into a zero-exploration
loop instead of a budget fire.

gepa's LM wrapper forwards reflection_lm_kwargs to litellm.completion;
an explicit temperature=1.0 makes each iteration propose a genuinely
different variant, so the execution budget is spent on NEW candidates
and a good draw can beat the baseline. (Judge and distillation calls
already run at temperature 0 — grading stays deterministic.)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…tion and pin the mutation output format

Two more live-run killers found after the temperature change shipped:

1. Reflection was served from Kasal's own cache. llm_manager enables a
   process-global litellm disk cache at import time, and gepa's LM
   wrapper rides the same litellm in the same process — so the identical
   per-iteration reflection prompt got the identical cached response
   forever (observed live: duration=0.00s, byte-identical proposals at
   temperature 1.0). Reflection calls now send cache={"no-cache": True}.

2. gepa's default proposal template says "write a new instruction ...
   within ``` blocks", which invited the reflection model to restructure:
   it returned {"instruction": "..."} JSON blobs that lost the
   [AGENT]/[TASK] document structure, so every proposal free-rejected
   without ever executing (observed live: 11/11 malformed, run ended
   after 1 execution). A custom reflection_prompt_template now pins the
   output contract to the crew-doc format. Validated offline against the
   live reflection endpoint: with the fenced-output instruction 2 of 3
   samples degenerated to a bare "```"; with the no-fence contract at
   temperature 0.8, 4 of 4 samples parsed, were distinct, and carried
   the human requirements into the mutated fields.

Also: _parse_crew_doc strips surviving markdown fences before parsing
(a recoverable wrapper should not cost a candidate), with tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ing panel

Crew deliverables are mostly GFM tables; the read-only TextField showed
them as raw pipe soup. The grading panel now renders them with
react-markdown + remark-gfm (the Documentation viewer's stack), with
compact table styling, scroll containment, and external links opening
in a new tab.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…elete from library, fix cross-crew visibility

Judges were create-only from Kasal. Now:

- Edit: assigned judge chips and library rows open an edit dialog for
  instructions and/or model. Updating registers a new version under the
  same registry name (MLflow scorers are versioned; latest wins —
  verified live against the local registry). Editing an assigned copy
  changes what that crew's runs use; editing a library judge leaves
  already-assigned snapshots untouched, and the dialog says so.
  list_judges now returns full instructions (4000 chars, was a 500-char
  preview) so the edit round-trip cannot truncate-corrupt a judge.

- Delete: library judges get a delete action with a confirm dialog
  (assigned copies survive; the chip's unassign already covered those).

- Cross-crew visibility fix: creating a judge from a crew's dialog
  registered ONLY the crew-scoped copy — no library original existed,
  so the judge never appeared in any other crew's Assign menu (observed
  live). Creation now registers the shared library original AND the
  crew-scoped copy.

- Experiment pinning: scorers are per-experiment, and none of the judge
  CRUD paths pinned one — a fresh worker's active experiment is
  Default/0, so judges could register into or list from the wrong
  experiment and silently vanish. All five judge bodies now pin the
  launch experiment ('kasal' fallback), matching the optimization runs.

Verified end-to-end against the live API: create-from-crew registers
both names, update re-versions the target only, delete cleans up.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ce, backend and frontend

Backend — 93 tests:
- service (67): crew-doc serialize/parse round-trip incl. multiline
  continuation and malformed-candidate rejection (the GEPA mutation
  contract), judge verdict normalization (numeric/bool/categorical),
  requirement distillation + checklist-mark grading (incl. the
  model-arithmetic and duplicate-mark hazards), judge lifecycle against
  a faked mlflow registry (dual library+scoped registration, {{ outputs }}
  autofix, update keeps omitted fields, assign snapshots, delete all
  versions, per-operation experiment pinning, local-mode gating), run
  registry behaviors (public-field allowlist never leaks task handles,
  timezone-aware sort, cancel transitions/guards, prune spares active
  runs), plus the existing mining/scorer/registry/apply coverage
- router (26): every endpoint's handler — ValueError → 400 for client
  input, missing runs → 404, response-schema wrapping incl. the
  progress-chip fields, judge CRUD passthrough, eval-feedback numeric
  coercion, user-context publication, and the request-schema bounds
  (budget 4..40 default 10; template name is a closed set)

Frontend — 46 tests:
- PromptOptimizationService (19): every method's endpoint, payload
  undefined-stripping, judge-name URL encoding, zero-grade preservation,
  list unwrapping and boolean coercion
- CrewOptimizeDialog (10): judge scoping by crew registry prefix
  (assigned chips vs library menu; other crews' judges invisible;
  same-name library duplicate excluded), edit-by-full-name, delete
  confirm, create auto-assign, honest-progress chips, budget-as-cap
  start request, ungraded-first eval ordering with counts, markdown
  answer rendering, grade submission
- PromptConfiguration (5): Optimize action gated to optimizable
  templates and to an onOptimize handler; edit dialog
- PromptOptimization (4): fixedTemplate hides the picker and filters
  runs; start payload; backend rejection surfaced
- Prompts (3): dialog opens scoped to the chosen template and closes
- optimizableTemplates (2): mirrors the backend TEMPLATE_TASKS set
- CrewFlowDialog optimize entry (3): source-level guards in the
  established CrewFeedback.test.tsx style (heavy dialog): click
  isolation and dialog wiring

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…any explicit temperature

A crew run on kimi-k2.7-code-highspeed "completed" after a single
execution: gepa made 10 reflection attempts and produced ZERO proposals
(no proposals.json artifact, no subsample scores). Direct reproduction
against Moonshot with gepa's exact request showed why — every call
400'd with "invalid temperature: only 1 is allowed for this model";
the same call without a temperature returns a clean, parseable crew
doc (thinking rides the separate reasoning_content field).

_resolve_reflection_model now returns (uri, env, provider) and the
crew body builds reflection_lm_kwargs per provider: Kimi gets only the
litellm cache bypass (its forced default sampling is already diverse);
everyone else keeps temperature 0.8. The template-mode optimizer gains
the same cache bypass (it never sent a temperature, so it was
Kimi-safe but cache-exposed).

New test pins the kimi resolver path (key lookup, URI, provider tag);
resolver tests updated for the 3-tuple.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@nehmetohme
nehmetohme force-pushed the prompt-optimization branch from b3d3d3a to 36a97c9 Compare July 24, 2026 12:18
…provider plumbing outside the manager

The optimization service had grown its own provider layer for the calls
MLflow/gepa make directly (reflection, registered judges, preflight):
a URI resolver, a provider→env-key table, per-tenant API keys written
into the shared process env, and request quirks re-implemented per
provider. Each of those independently produced live failures (retired
DeepSeek names, Kimi temperature 400s, key-lookup drift) that LLMManager
already handles centrally for the rest of Kasal.

Now every LLM call goes through LLMManager.completion:

- GEPA reflection: gepa accepts a LanguageModel callable, but MLflow's
  GepaPromptOptimizer pins reflection_lm to a litellm string. A one-time
  idempotent patch on gepa.optimize swaps in a per-worker-thread
  callable (each run owns one thread) backed by _sync_llm_completion.
  The callable adds a unique system-line cache-buster (the process-
  global litellm cache replayed one proposal forever for the identical
  per-iteration prompt) and passes temperature=0.8 — the MANAGER drops
  it for providers that reject it (Kimi), so the provider-aware kwargs
  from the previous commit are deleted again.

- Registered judges: rendered and invoked HERE via LLMManager instead
  of mlflow's own model client. Judges store a plain Kasal model key
  (wrapped as 'openai:/<key>' only to satisfy make_judge's URI shape);
  legacy URI-storing judges are stripped back to their best key.

- Preflight: a manager-routed ping instead of a raw litellm call with
  env juggling.

Deleted: _resolve_reflection_model, _REFLECTION_KEY_ENV, reflection
env application/restore in both sync bodies, reflection_lm_kwargs, and
the reflection_provider threading. Security side effect: per-tenant
API keys are no longer written into os.environ.

Tests: resolver suite replaced with coverage of the new pure helpers
(key stripping, grade parsing, bridge install/override/idempotency,
preflight error wrapping, reflection message shaping incl. distinct
cache-busters); judge lifecycle updated for key storage. 95 backend
tests pass.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants