# Agent Learning & Optimization

URL: https://cogitator.app/docs/advanced/learning

> Improve agents over time with execution traces, DSPy-style compilation of instructions and demos, A/B tests, degradation monitoring and automated optimization.

## Overview

The learning system captures execution traces, scores them, and optimizes instructions and few-shot examples through an iterative compilation process. It brings ideas from DSPy — bootstrapped demonstrations, instruction optimization, and automated evaluation — into the Cogitator runtime.

For deploying, rolling back and A/B testing instructions on live traffic, use [Prompt Versions](/docs/advanced/prompt-versions) (`cog.prompts`): it is built on the same A/B testing and version classes described below and wires them into every run.

## AgentOptimizer

The `AgentOptimizer` coordinates trace capture, demo selection and instruction optimization:

```typescript
import { AgentOptimizer, Cogitator } from '@cogitator-ai/core';

const cogitator = new Cogitator({
  llm: { providers: { openai: { apiKey: process.env.OPENAI_API_KEY! } } },
});

const { backend, model } = cogitator.route('openai/gpt-6.1-sol');

const optimizer = new AgentOptimizer({
  llm: backend,
  model, // 'gpt-6.1-sol': sent to the backend as is
  cogitator, // runs the trainset of compile(); optional
  config: {
    maxDemosPerAgent: 5,
    minScoreForDemo: 0.8,
  },
});
```

The optimizer's own model calls (instruction analysis, candidate generation and evaluation) go to `llm` with `model`. Traces live in an `InMemoryTraceStore` unless you pass `traceStore` (any `TraceStore`); `insightStore` adds [reflection](/docs/advanced/reflection) insights to instruction candidates.

## Capturing Traces

Every agent run can be captured as an `ExecutionTrace` with step-by-step details and computed metrics:

```typescript
const result = await cogitator.run(agent, { input: 'Summarize this article...' });

const trace = await optimizer.captureTrace(result, 'Summarize this article...', {
  expected: 'A concise 3-sentence summary covering key points',
  labels: ['summarization', 'production'],
});

console.log(`Score: ${trace.score}`);
console.log(`Tool accuracy: ${trace.metrics.toolAccuracy}`);
console.log(`Efficiency: ${trace.metrics.efficiency}`);
```

The trace records the run's model calls and tool executions with token usage, latency and cost. By default its `score` is a weighted average of `success` (0.4, no errors), `tool_accuracy` (0.3, the share of tool calls without errors) and `efficiency` (0.3, tokens and duration). `expected` and `labels` are stored with the trace; for a run with tool calls, a string `expected` also moves `tool_accuracy` up or down by 0.2 depending on whether the output contains it. When the run's instructions came from [`cog.prompts`](/docs/advanced/prompt-versions), the trace's `prompt` holds the version or A/B variant it used (`result.prompt`).

### Metrics and Retention

The optimizer's `config` decides how traces are scored and kept:

```typescript
const optimizer = new AgentOptimizer({
  llm: backend,
  model,
  config: {
    defaultMetrics: ['success', 'completeness', 'coherence'],
    customMetrics: [
      { name: 'success', type: 'boolean', description: 'No errors', weight: 0.6 },
      { name: 'cites_sources', type: 'numeric', description: 'Output cites a source', weight: 0.4 },
    ],
    captureTraces: true,
    traceRetention: 1000,
  },
});

optimizer.getMetricEvaluator().registerMetric('cites_sources', (trace) => {
  const value = trace.output.includes('http') ? 1 : 0;
  return { name: 'cites_sources', value, passed: value === 1 };
});
```

* `defaultMetrics` — the built-in metrics traces are scored with: `success`, `tool_accuracy`, `efficiency`, `completeness` and `coherence` (default the first three). `completeness` and `coherence` are judged by `llm`, so each costs a model call per trace; their default weight is 0.3.
* `customMetrics` — more `MetricDefinition`s, or new weights for built-in metrics of the same name. A metric that is not built in needs its function, registered with `getMetricEvaluator().registerMetric(name, fn)`; without one it scores a neutral 0.5.
* `registerMetric(name, fn)` alone also works: a registered metric that is not in the config is scored with weight 1, and one with a built-in name replaces the built-in function.
* `captureTraces` — `false` scores traces without storing them (default `true`).
* `traceRetention` — the number of traces kept per agent; older and lower-scoring ones are pruned after each capture.

`autoOptimize`, `optimizeAfterRuns` and `traceStore` in the learning config are deprecated and not read: use the [`AutoOptimizer`](#autooptimizer) and the optimizer's `traceStore` option.

## Compiling (DSPy-Style Optimization)

`compile()` runs multi-round optimization of an agent's instructions and in-context examples:

```typescript
const compiled = await optimizer.compile(agent, trainset, {
  maxRounds: 3,
  maxBootstrappedDemos: 5,
  optimizeInstructions: true,
});

console.log(`Score: ${compiled.scoreBefore} → ${compiled.scoreAfter}`);
console.log(`Improvement: ${(compiled.improvement * 100).toFixed(1)}%`);
console.log(`Demos added: ${compiled.demosAdded.length}`);
console.log(`New instructions: ${compiled.instructionsAfter}`);
```

`trainset` is an array of `{ input, expected? }`. With a `cogitator`, `compile()` first runs the trainset with the original agent: those runs give `scoreBefore` and are stored as traces. Each round, high-scoring traces become bootstrapped demos, and the `InstructionOptimizer` analyzes failing traces and generates improved instruction candidates, keeping the best one by LLM-based evaluation. Finally the trainset runs again with the optimized instructions, which gives `scoreAfter`.

Without a `cogitator`, a trainset cannot run: `compile()` reports that in `errors`, works on the traces already captured, and `scoreAfter` is the instruction optimizer's estimate. Running a trainset also needs an agent created with `new Agent()`.

`compile()` does not change the agent. Apply the result yourself, for example as a new [prompt version](/docs/advanced/prompt-versions):

```typescript
if (compiled.improvement > 0 && compiled.instructionsAfter) {
  await cogitator.prompts.deploy(agent, compiled.instructionsAfter, 'optimization');
}
```

## Bootstrapping Demos

Demos are high-quality execution examples you can put into the agent's prompt as few-shot examples:

```typescript
const demos = await optimizer.bootstrapDemos(agent.id);

const relevantDemos = await optimizer.getDemosForPrompt(
  agent.id,
  'Find papers on transformer architectures',
  3
);

const fewShot = optimizer.formatDemosForPrompt(relevantDemos);
const withDemos = agent.clone({ instructions: `${agent.instructions}\n\n${fewShot}` });
```

Only traces scoring at least `minScoreForDemo` (default 0.8) become demos, up to `maxDemosPerAgent` (default 5) per agent. `getDemosForPrompt` ranks them by word overlap with the input, their score and how often they were used, and keeps the selection diverse. Demos are not added to runs automatically.

## Learning Stats

```typescript
const stats = await optimizer.getStats(agent.id);

console.log(`Traces: ${stats.traces.totalTraces}`);
console.log(`Demos: ${stats.demos.totalDemos}`);
console.log(`Optimization runs: ${stats.optimization.runsOptimized}`);
console.log(`Avg improvement: ${(stats.optimization.averageImprovement * 100).toFixed(1)}%`);
```

`getTraceStore()`, `getMetricEvaluator()`, `getDemoSelector()` and `getInstructionOptimizer()` give access to the parts.

## A/B Testing

`ABTestingFramework` compares two sets of instructions with Welch's t-test. `cog.prompts.startABTest()` runs it on live traffic for you; use the framework directly when you route traffic yourself:

```typescript
import { ABTestingFramework, InMemoryABTestStore } from '@cogitator-ai/core';

const abTesting = new ABTestingFramework({
  store: new InMemoryABTestStore(),
  defaultConfidenceLevel: 0.95,
  defaultMinSampleSize: 50,
});

const test = await abTesting.createTest({
  agentId: 'support-agent',
  name: 'Concise vs detailed instructions',
  controlInstructions: 'You are a helpful support agent...',
  treatmentInstructions: 'You are a support agent. Be concise, use bullet points...',
  treatmentAllocation: 0.5,
});

await abTesting.startTest(test.id);
```

For each request, pick a variant, run with its instructions and record the outcome:

```typescript
const variant = abTesting.selectVariant(test);
const instructions = abTesting.getInstructionsForVariant(test, variant);

const started = Date.now();
const run = await cogitator.run(agent.clone({ instructions }), { input });
await abTesting.recordResult(test.id, variant, score, Date.now() - started, run.usage.cost);

const outcome = await abTesting.checkAndCompleteIfReady(test.id);
if (outcome?.isSignificant) {
  console.log(`Winner: ${outcome.winner} (p=${outcome.pValue.toFixed(4)})`);
  console.log(outcome.recommendation);
}
```

`selectVariant` is random per call. A test completes once the two variants together have twice `minSampleSize` results and the difference is significant at `confidenceLevel`, or once `maxDuration` (default 7 days) has passed. `analyzeResults(test)` gives the current outcome without completing; `pauseTest`, `resumeTest`, `completeTest` and `cancelTest` manage the lifecycle.

## Instruction Versions

`RollbackManager` keeps the deployed instructions of each agent as versions with run metrics:

```typescript
import { RollbackManager, InMemoryInstructionVersionStore } from '@cogitator-ai/core';

const rollbackManager = new RollbackManager({
  store: new InMemoryInstructionVersionStore(),
  maxVersionsToKeep: 20,
});

await rollbackManager.deployVersion(
  'support-agent',
  'You are a helpful support agent...',
  'manual'
);
await rollbackManager.recordMetrics('support-agent', 0.9, 1200, 0.002, true);

const history = await rollbackManager.getVersionHistory('support-agent');
await rollbackManager.rollbackToPrevious('support-agent');
```

`cog.prompts` uses the same class; reach for it directly in custom pipelines.

## Prompt Monitoring

`PromptMonitor` compares recent executions of an agent with a baseline and raises alerts on degradation:

```typescript
import { PromptMonitor } from '@cogitator-ai/core';

const monitor = new PromptMonitor({
  windowSize: 60 * 60 * 1000, // 1 hour window
  scoreDropThreshold: 0.15, // 15% score drop
  latencySpikeThreshold: 2.0, // 2x latency
  errorRateThreshold: 0.1, // 10% error rate
  enableAutoRollback: true,
  onAlert: (alert) => {
    console.log(
      `[${alert.severity}] ${alert.type}: ${alert.currentValue} vs baseline ${alert.baselineValue}`
    );
  },
});

monitor.setBaseline(
  'support-agent',
  monitor.computeBaselineFromHistory('support-agent', pastTraces)
);

const alerts = monitor.recordExecution(trace);
for (const alert of alerts) {
  if (alert.severity === 'critical') {
    console.log('Critical degradation detected!');
  }
}
```

The monitor raises four alert types: `score_drop`, `latency_spike`, `error_rate_increase` and `cost_spike`, as `warning` or — beyond twice the threshold — `critical`. The monitor does not change anything itself: with `enableAutoRollback`, critical alerts carry `autoAction: 'rollback'`, which the `AutoOptimizer` acts on. `getCurrentMetrics`, `getAlerts`, `getActiveAlerts` and `resolveAlert` inspect its state.

## AutoOptimizer

The `AutoOptimizer` ties the pieces into an automated loop:

```typescript
import { AutoOptimizer } from '@cogitator-ai/core';

const autoOptimizer = new AutoOptimizer({
  enabled: true,
  triggerAfterRuns: 100,
  minRunsForOptimization: 20,
  requireABTest: false,
  maxOptimizationsPerDay: 3,
  agentOptimizer: optimizer,
  abTesting,
  monitor,
  rollbackManager,
  onOptimizationComplete: (run) => {
    console.log(`Optimization ${run.id}: ${run.status}`);
  },
  onRollback: (agentId, reason) => console.log(`Rolled back ${agentId}: ${reason}`),
});

const trace = await optimizer.captureTrace(result, input);
await autoOptimizer.recordExecution(trace);
```

`enabled` defaults to `false`. Each recorded trace goes to the monitor (a critical alert with `autoAction: 'rollback'` rolls the agent back to its previous version). Every `triggerAfterRuns` traces, once at least `minRunsForOptimization` were recorded and at most `maxOptimizationsPerDay` times a day, it compiles the agent's current version (from the `RollbackManager`) against the traces in the optimizer's store. Deploy a version first: without one, the optimization run ends with `status: 'failed'` and an `error`. Improved instructions are deployed as a new version, or with `requireABTest` (the default) put into an A/B test against the current ones; the optimization run completes when the test does, deploying the treatment if it won. `triggerOptimization(agentId)` and `forceRollback(agentId, reason)` run the steps by hand.

### A/B Tests on Live Traffic

To serve an A/B test to real traffic, let the Cogitator run it: give its [`prompts`](/docs/advanced/prompt-versions) config the same stores as the `AutoOptimizer`, and give the agent an explicit `id` (versions and tests are keyed by it, and traces carry it):

```typescript
import {
  ABTestingFramework,
  Agent,
  AutoOptimizer,
  Cogitator,
  InMemoryABTestStore,
  InMemoryInstructionVersionStore,
  PromptMonitor,
  RollbackManager,
} from '@cogitator-ai/core';

const abTests = new InMemoryABTestStore();
const versions = new InMemoryInstructionVersionStore();

const cogitator = new Cogitator({
  llm: { providers: { openai: { apiKey: process.env.OPENAI_API_KEY! } } },
  prompts: { abTests, versions },
});

const agent = new Agent({
  id: 'support-agent',
  name: 'support',
  model: 'openai/gpt-5.5',
  instructions: 'You are a helpful support agent.',
});
await cogitator.prompts.deploy(agent, agent.instructions);

const autoOptimizer = new AutoOptimizer({
  enabled: true,
  agentOptimizer: optimizer,
  abTesting: new ABTestingFramework({ store: abTests }),
  monitor: new PromptMonitor(),
  rollbackManager: new RollbackManager({ store: versions }),
});

const result = await cogitator.run(agent, { input, threadId });
await autoOptimizer.recordExecution(await optimizer.captureTrace(result, input));
```

While the test runs, the Cogitator gives each thread one variant and records every run's result against it; the trace's `prompt` tells the `AutoOptimizer` that the run was already recorded, so it is not counted twice, and the optimizer settles its run once the test has completed. Leave `prompts.autoDeployWinner` off: the `AutoOptimizer` deploys the winner itself.

Traces without a `prompt` — runs of a Cogitator without `prompts` — are recorded by the `AutoOptimizer`: against the A/B test as the variant whose instructions are currently deployed, and to that version's metrics.

## Storage

Everything above keeps its data in process memory by default. `PostgresTraceStore` persists all of it in one schema (default `cogitator`):

```typescript
import { PostgresTraceStore } from '@cogitator-ai/core';

const store = new PostgresTraceStore({ connectionString: process.env.DATABASE_URL! });
await store.connect();

const optimizer = new AgentOptimizer({ llm: backend, model, traceStore: store.traces() });
const abTesting = new ABTestingFramework({ store: store.abTests() });
const rollbackManager = new RollbackManager({ store: store.instructionVersions() });

await store.storeTrace(trace);
const recent = await store.queryTraces({ agentId: 'support-agent', limit: 50 });
```

It stores traces (`storeTrace`, `queryTraces`, `getAllTraces`), captured prompts (`capture`, `getByRun`, `query`), A/B tests and instruction versions. `traces()` adapts it to a `TraceStore` for `AgentOptimizer`, `DemoSelector` or [time travel](/docs/advanced/time-travel); `abTests()` and `instructionVersions()` adapt it for A/B tests and versions, including the [`prompts`](/docs/advanced/prompt-versions#stores) config. Traces keep their `prompt` in a `prompt` column, which `connect()` adds to tables created by older versions.

## Prompt Logging

`wrapWithPromptLogger` wraps a backend so that every request and response is captured into a prompt store (anything with `capture()`, such as `PostgresTraceStore`), for debugging and auditing:

```typescript
import { OpenAIBackend, wrapWithPromptLogger } from '@cogitator-ai/core';

const logged = wrapWithPromptLogger(
  new OpenAIBackend({ apiKey: process.env.OPENAI_API_KEY! }),
  store,
  {
    captureContent: true,
    captureTools: true,
  }
);

logged.setContext({ runId: 'run_1', agentId: 'support-agent', threadId: 'thread_1' });
await logged.chat({ model: 'gpt-6-luna', messages: [{ role: 'user', content: 'Hello' }] });
logged.clearContext();
```

Calls are captured only while a context is set; without one the wrapper passes requests through.