Skip to content

Feat: add optimization module with SimplePromptOptimizer (Part 2/2) - #546

Open
AmaadMartin wants to merge 1 commit into
google:mainfrom
AmaadMartin:upstream-port/agent-optimization-part2
Open

Feat: add optimization module with SimplePromptOptimizer (Part 2/2)#546
AmaadMartin wants to merge 1 commit into
google:mainfrom
AmaadMartin:upstream-port/agent-optimization-part2

Conversation

@AmaadMartin

Copy link
Copy Markdown
Collaborator

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

  1. Link to an existing issue (if applicable):
    Closes: Feat: Port the agent optimization module so agents can automatically improve their own prompts (parity with adk-python optimization) #534
  2. Or, if no issue exists, describe the change:
    Problem: adk-python can automatically rewrite an agent's system prompt to make it score better (src/google/adk/optimization/), but adk-js has no equivalent. Tuning an agent's instruction in adk-js is guesswork rather than a measured search.

Solution: Port the core of that module plus a concrete SimplePromptOptimizer to @google/adk (new core/src/optimization/), with no new npm dependency and no eval-framework coupling.

  • data_types: SamplingResult, UnstructuredSamplingResult, AgentWithScores, OptimizerResult.
  • Sampler (abstract): the interface developers implement to plug in their own scoring/evaluation service. Keeping scoring behind this interface means the framework never decides what "good" means, and a local eval-backed sampler can slot in later without touching any optimizer.
  • AgentOptimizer (abstract): base class for optimizers.
  • SimplePromptOptimizer (@experimental): a hill-climbing prompt optimizer. Each round it (1) asks an LLM to rewrite the agent's instruction, telling it the current score; (2) builds a candidate via clone({instruction}); (3) scores the candidate on a random batch of training examples via the developer's Sampler; (4) keeps the candidate only if it beats the current best. After numIterations rounds it scores the winner on the validation set and returns it.

Notes:

  • optimizerModel accepts either a model-name string (resolved via LLMRegistry) or a BaseLlm instance. The instance overload keeps the string default for parity with the issue's API and lets unit tests inject a stub model with no network and no registry mocking.
  • Thought parts are never included in the generated prompt (part.text && !part.thought).
  • This optimizer is marked @experimental. It is not free to run: with the defaults (10 iterations, batch size 5) a single optimize() performs 50+ candidate scoring runs plus ~10 rewriter model calls, so expect real cost and latency.
  • This is Part 2 of a two-part stack; it builds on BaseAgent.clone() from Part 1: Feat: add clone() to BaseAgent (Part 1/2) #545. Out of scope (tracked separately): the GEPA optimizers, a local-eval-backed Sampler, and the eval framework itself.

Stacked PR — please review #545 first. This PR targets the agent-clone-part1 branch (not main), so the diff shown here contains only the optimization module, without Part 1's clone() changes. Once #545 lands in main, this PR will be retargeted to main and the base branch deleted.

Testing Plan

Please describe the tests that you ran to verify your changes. This is required for all PRs that are not small documentation or typo fixes.
Unit Tests:
[x] I have added or updated unit tests for my change.
[x] All unit tests pass locally.

Added core/test/optimization/simple_prompt_optimizer_test.ts using a hand-written MockLlm extends BaseLlm (no network) and a StubSampler extends Sampler. It covers: the happy path (adopts an improved prompt; exact sampleAndScore/generateContentAsync/getTrainExampleIds call counts; original agent left unmodified; batches are the requested size, unique, and in-set), adopt-if-better discard, thought-part and text-less-part skipping, the batchSize clamp + warning, empty-scores-return-0, validation-score averaging, the non-string-instruction fail-fast error (no LLM/sampler calls made), passing through the configured model configuration, and default model resolution via LLMRegistry. All new executable code has 100% line and branch coverage (data_types.ts is interfaces-only, so it has no executable code to cover).

Ran locally (from repo root):

  • npx vitest run --project unit:core simple_prompt_optimizer_test — 10/10 pass.
  • npm run build (core) — passes.
  • npm run lint, prettier --check, and npm run docs:check (TypeDoc, warnings-as-errors) — all clean.

Manual End-to-End (E2E) Tests:
Added dev/samples/simple_prompt_optimizer.ts, a runnable sample that doubles as a manual integration test. It builds an LlmAgent, implements a small deterministic Sampler (heuristic keyword scorer, no credentials), and runs the optimizer with a self-contained rewriter model, so it exercises the real optimize()clone() → scoring loop end-to-end with no mocks and no network. Run it with:

npx tsx dev/samples/simple_prompt_optimizer.ts

Observed output: the instruction is rewritten from "You are a customer support agent." to the stronger prompt and the validation score rises to 1. To try it against a real model, set optimizerModel: 'gemini-2.5-flash' and provide GOOGLE_GENAI_API_KEY (or GEMINI_API_KEY); the @experimental warning is logged once per process.

Checklist

[x] I have read the CONTRIBUTING.md document.
[x] I have performed a self-review of my own code.
[x] I have commented my code, particularly in hard-to-understand areas.
[x] I have added tests that prove my fix is effective or that my feature works.
[x] New and existing unit tests pass locally with my changes.

@AmaadMartin

Copy link
Copy Markdown
Collaborator Author

Note on CI for this PR: the validation workflow is configured with pull_request: branches: [main], so it does not trigger for a PR targeting the agent-clone-part1 stack base — only auto-assign and cla/google run here. Full validation will run automatically once this is retargeted to main after #545 merges.

In the meantime it was validated locally on the exact commit in this PR:

  • npx vitest run --project unit:core simple_prompt_optimizer_test — 10/10 pass
  • npx vitest run --project unit:core base_agent_test — 16/16 pass (Part 1's clone() still green with this on top)
  • npm run build — clean
  • npm run lint — clean

@AmaadMartin
AmaadMartin changed the base branch from agent-clone-part1 to main July 28, 2026 20:48
Port adk-python's agent optimization core to @google/adk under
core/src/optimization/:

- data_types: SamplingResult, UnstructuredSamplingResult, AgentWithScores,
  OptimizerResult.
- Sampler: abstract base developers implement to plug in their scoring/eval
  service, keeping scoring decisions out of the optimizer.
- AgentOptimizer: abstract base for optimizers.
- SimplePromptOptimizer (@experimental): hill-climbs an agent's string
  instruction. Each round it asks an LLM to rewrite the instruction (told the
  current score), builds a candidate via clone({instruction}), scores it on a
  random batch of training examples, and keeps it only if it beats the current
  best; then scores the winner on the validation set. optimizerModel accepts a
  string (resolved via LLMRegistry) or a BaseLlm instance (enables no-network
  tests). Thought parts are never included in the generated prompt.

No new npm dependency. New public symbols are exported from common.ts. Added
unit tests (stub LLM + sampler, no network; 100% coverage of new executable
code) and a runnable, offline sample that doubles as a manual integration test.

Related: google#534
@AmaadMartin
AmaadMartin force-pushed the upstream-port/agent-optimization-part2 branch from 6569d0c to f256026 Compare July 28, 2026 20:53
AmaadMartin pushed a commit to AmaadMartin/adk-js that referenced this pull request Jul 28, 2026
Ports adk-python's GEPARootAgentPromptOptimizer to adk-js so users can
optimize a root agent's system instruction with GEPA (reflective prompt
evolution).

- optimization/gepa/adapter.ts: engine-side protocol (EvaluationBatch,
  GepaAdapter) mirroring gepa.core.adapter.
- optimization/gepa/engine.ts: a native, minimal, first-party GEPA engine
  (optimize + default reflective proposer) with no new runtime dependencies,
  browser-safe, and deterministic when seeded.
- optimization/gepa_root_agent_prompt_optimizer.ts: the optimizer, its config
  and result, and AgentGepaAdapter (clones the agent per candidate and scores
  via the developer's Sampler).
- Exports wired through common.ts (`@google/adk`).

The external Python `gepa` package has no official JS build, so the engine is
reimplemented behind a mockable GepaAdapter/optimize seam. Adapter/wiring unit
tests mock the engine (parity with the adk-python test), a native engine test
suite validates the loop, and an integration test exercises the full pipeline
with a fake registered BaseLlm and Sampler (no mocks).

Builds on the optimization base module (Sampler, AgentOptimizer, data types)
from google#546 and the BaseAgent.clone() API landed in
google#545; this change adds only the GEPA-specific contribution.
AmaadMartin pushed a commit to AmaadMartin/adk-js that referenced this pull request Jul 28, 2026
Ports adk-python's GEPARootAgentPromptOptimizer to adk-js so users can
optimize a root agent's system instruction with GEPA (reflective prompt
evolution).

- optimization/gepa/adapter.ts: engine-side protocol (EvaluationBatch,
  GepaAdapter) mirroring gepa.core.adapter.
- optimization/gepa/engine.ts: a native, minimal, first-party GEPA engine
  (optimize + default reflective proposer) with no new runtime dependencies,
  browser-safe, and deterministic when seeded.
- optimization/gepa_root_agent_prompt_optimizer.ts: the optimizer, its config
  and result, and AgentGepaAdapter (clones the agent per candidate and scores
  via the developer's Sampler).
- Exports wired through common.ts (`@google/adk`).

The external Python `gepa` package has no official JS build, so the engine is
reimplemented behind a mockable GepaAdapter/optimize seam. Adapter/wiring unit
tests mock the engine (parity with the adk-python test), a native engine test
suite validates the loop, and an integration test exercises the full pipeline
with a fake registered BaseLlm and Sampler (no mocks).

Builds on the optimization base module (Sampler, AgentOptimizer, data types)
from google#546 and the BaseAgent.clone() API landed in
google#545; this change adds only the GEPA-specific contribution.
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.

Feat: Port the agent optimization module so agents can automatically improve their own prompts (parity with adk-python optimization)

1 participant