Map-Reduce

import { Agent, Cogitator } from '@cogitator-ai/core';
import {
  WorkflowBuilder,
  WorkflowExecutor,
  agentNode,
  collect,
  executeMapReduce,
  functionNode,
} from '@cogitator-ai/workflows';

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

interface Doc {
  id: string;
  title: string;
  content: string;
}

interface DocAnalysis {
  title: string;
  findings: string;
  sentiment: string;
}

interface AnalysisState {
  [key: string]: unknown;
  documents: Doc[];
  analyses: DocAnalysis[];
  summary: string;
}

const documents: Doc[] = [
  { id: 'doc-1', title: 'Q4 Revenue Report', content: 'Revenue grew 23% to $4.2B. Cloud led at 35%; hardware fell 5%.' },
  { id: 'doc-2', title: 'Engineering Retro', content: 'Deploys went from weekly to daily. Incident response time fell from 45 to 12 minutes.' },
  { id: 'doc-3', title: 'Security Audit', content: '3 critical vulnerabilities found and patched within 48 hours. Compliance score is 94%.' },
];

const cog = new Cogitator({ llm: { providers: { google: { apiKey } } } });
const model = 'google/gemini-3.5-flash-lite';

const analyzer = new Agent({
  name: 'document-analyzer',
  model,
  instructions: 'Reply exactly as:\nFINDINGS: <one line>\nSENTIMENT: <positive|negative|neutral|mixed>',
  temperature: 0.2,
});

const summarizer = new Agent({
  name: 'executive-summarizer',
  model,
  instructions: 'Write an executive summary of the analyses in under 120 words, as bullet points.',
  temperature: 0.3,
});

const analyzeAll = functionNode<AnalysisState, DocAnalysis[]>(
  'analyze-all',
  async (state) => {
    const mapped = await executeMapReduce<AnalysisState, DocAnalysis, DocAnalysis[]>(state, {
      name: 'doc-analysis',
      map: {
        items: (s) => s.documents,
        mapper: async (item) => {
          const doc = item as Doc;
          const { output } = await cog.run(analyzer, { input: `${doc.title}\n\n${doc.content}` });
          return {
            title: doc.title,
            findings: output.match(/FINDINGS:\s*(.+)/i)?.[1]?.trim() ?? output,
            sentiment: output.match(/SENTIMENT:\s*(\w+)/i)?.[1]?.toLowerCase() ?? 'unknown',
          };
        },
        concurrency: 3,
        continueOnError: true,
        onProgress: (progress) => console.log(`mapped ${progress.completed}/${progress.total}`),
      },
      reduce: { ...collect<DocAnalysis>() },
    });
    return mapped.reduced;
  },
  { stateMapper: (output) => ({ analyses: output as DocAnalysis[] }) }
);

const summarize = agentNode<AnalysisState>(summarizer, {
  inputMapper: (state) =>
    state.analyses.map((a) => `${a.title}: ${a.findings} (${a.sentiment})`).join('\n'),
  stateMapper: (result) => ({ summary: result.output }),
});

const workflow = new WorkflowBuilder<AnalysisState>('document-analysis')
  .initialState({ documents: [], analyses: [], summary: '' })
  .addNode(analyzeAll.name, analyzeAll.fn)
  .addNode(summarize.name, summarize.fn, { after: [analyzeAll.name] })
  .entryPoint(analyzeAll.name)
  .build();

const result = await new WorkflowExecutor(cog).execute(workflow, { documents });
for (const analysis of result.state.analyses) {
  console.log(`[${analysis.sentiment}] ${analysis.title}: ${analysis.findings}`);
}
console.log('\n' + result.state.summary);

await cog.close();

pnpm add @cogitator-ai/core @cogitator-ai/workflows

GOOGLE_API_KEY=your-key npx tsx map-reduce.ts

npx tsx examples/workflows/03-map-reduce.ts