Optimize an Agent

import { Agent, AgentOptimizer, Cogitator, InMemoryTraceStore, createLLMBackend, tool } from '@cogitator-ai/core';
import { z } from 'zod';

const apiKey = process.env.GOOGLE_API_KEY;
if (!apiKey) throw new Error('Set GOOGLE_API_KEY');

const answers: Record<string, string> = {
  'What is the capital of France?': 'Paris',
  'Who wrote "1984"?': 'George Orwell',
  'What planet is closest to the Sun?': 'Mercury',
};

const checkAnswer = tool({
  name: 'check_answer',
  description: 'Check if an answer to a trivia question is correct',
  parameters: z.object({ question: z.string(), answer: z.string() }),
  execute: async ({ question, answer }) => {
    const correct = answers[question];
    if (!correct) return { found: false, question };
    return { correct: answer.toLowerCase().includes(correct.toLowerCase()), expected: correct };
  },
});

const cog = new Cogitator({ llm: { providers: { google: { apiKey } } } });

const optimizer = new AgentOptimizer({
  llm: createLLMBackend('google', { providers: { google: { apiKey } } }),
  model: 'gemini-3.8-flash',
  traceStore: new InMemoryTraceStore(),
  cogitator: cog,
});

const agent = new Agent({
  name: 'trivia-bot',
  model: 'google/gemini-3.5-flash-lite',
  instructions: 'You are a trivia bot. Answer, verify with check_answer, then state the answer clearly.',
  tools: [checkAnswer],
  temperature: 0.3,
  maxIterations: 5,
});

for (const [question, expected] of Object.entries(answers)) {
  const result = await cog.run(agent, { input: question });
  const trace = await optimizer.captureTrace(result, question, { expected, labels: ['trivia'] });
  console.log(`${question} -> score ${trace.score.toFixed(2)}`);
}

const optimization = await optimizer.compile(
  agent,
  Object.entries(answers).map(([input, expected]) => ({ input, expected })),
  { maxRounds: 2, maxBootstrappedDemos: 3 }
);
console.log(`score ${optimization.scoreBefore.toFixed(2)} -> ${optimization.scoreAfter.toFixed(2)}`);
console.log('demos added:', optimization.demosAdded.length);
console.log('new instructions:', optimization.instructionsAfter ?? '(unchanged)');

const demos = await optimizer.bootstrapDemos(agent.id);
console.log(optimizer.formatDemosForPrompt(demos.slice(0, 2)));

await cog.close();

pnpm add @cogitator-ai/core zod

GOOGLE_API_KEY=your-key npx tsx agent-optimizer.ts

npx tsx examples/core/07-agent-optimizer.ts