Multi-Agent

/**
 * Multi-Agent Web Reading & Summarization using AG-Kit Agents
 */

import { Agent, OpenAIProvider } from "@ag-kit/agents";
import { tool } from "@ag-kit/tools";
import { z } from "zod";

// Web Reading Session State
class WebReadingState {
  urls: string[] = [];
  targetLanguage: string = "en";
  crawlScope: "single_page" | "same_domain" = "single_page";
  maxPages: number = 5;
  maxDepth: number = 1;
  outputFormat: "summary_only" | "summary_and_schema" | "full" =
    "summary_and_schema";
  currentPhase:
    | "planning"
    | "retrieving"
    | "translating"
    | "extracting"
    | "evaluating"
    | "complete" = "planning";
  documents: any[] = [];
  translatedDocuments: any[] = [];
  extractedSchema: any = null;
  mindMap: any = null;
  evaluationScore: any = null;
  errors: string[] = [];
  logs: string[] = [];

  [key: string]: unknown;
}

/**
 * Create Multi-Agent Web Reading agent
 */
export function createMultiAgentOrchestrator(): Agent<WebReadingState> {
  if (!process.env.OPENAI_API_KEY) {
    throw new Error("OPENAI_API_KEY is required");
  }

  const provider = new OpenAIProvider({
    apiKey: process.env.OPENAI_API_KEY!,
    defaultModel: process.env.OPENAI_MODEL || "gpt-4o",
    baseURL: process.env.OPENAI_BASE_URL,
  });

  // Create specialized agents
  const pageProcessorAgent = createPageProcessorAgent();
  const reportGeneratorAgent = createReportGeneratorAgent();

  return new Agent<WebReadingState>({
    name: "multi-agent-orchestrator",
    description:
      "Orchestrates multi-page web reading and summarization workflow",
    model: provider,
    stateType: WebReadingState,
    instructions: `You are the orchestrator for multi-page web reading. When given URLs, you should:

1. **Task Planning**: Use todo_write to plan tasks for each page (no session_id needed)
2. **Task Tracking**: Use todo_read with the session_id returned from todo_write
3. **Page Processing**: Coordinate page processing for each URL using process_page
4. **Report Generation**: Coordinate report generation from all results using generate_report
5. **Final Output**: Generate a comprehensive report in the user's language

**Multi-Agent Workflow**:
- Use todo_write to break down complex multi-page tasks into manageable steps
- Each page can be processed independently by PageProcessorAgent
- ReportGeneratorAgent integrates all results into a final report
- This demonstrates true multi-agent coordination with task management`,
    modelSettings: {
      temperature: 0.7,
      maxTokens: 8192 * 2,
    },
    tools: [
      todoWriteTool,
      todoReadTool,
      crawlTool,
      createProcessPageTool(pageProcessorAgent),
      createGenerateReportTool(reportGeneratorAgent),
    ],
  });
}

// ==================== Specialized Agents ====================

function createPageProcessorAgent(): Agent {
  return new Agent({
    name: "PageProcessorAgent",
    description: "Process single page: crawl, translate, summarize",
    model: new OpenAIProvider({
      apiKey: process.env.OPENAI_API_KEY!,
      defaultModel: process.env.OPENAI_MODEL || "gpt-4o-mini",
      baseURL: process.env.OPENAI_BASE_URL,
    }),
    instructions: `You are PageProcessorAgent. Process a single page:
1. Crawl the page content
2. Translate if needed to target language
3. Extract key points and entities
4. Return structured data for report generation`,
  });
}

function createReportGeneratorAgent(): Agent {
  return new Agent({
    name: "ReportGeneratorAgent",
    description: "Generate comprehensive report from multiple page results",
    model: new OpenAIProvider({
      apiKey: process.env.OPENAI_API_KEY!,
      defaultModel: process.env.OPENAI_MODEL || "gpt-4o",
      baseURL: process.env.OPENAI_BASE_URL,
    }),
    instructions: `You are ReportGeneratorAgent. Generate comprehensive reports:
1. Integrate results from multiple pages
2. Create structured summary with key insights
3. Generate mind map (mermaid format)
4. Provide quality assessment
5. Output in user's target language`,
  });
}

// ==================== Tools ====================

// In-memory todo storage
const todoStorage = new Map<
  string,
  Array<{
    id: string;
    content: string;
    status: "pending" | "completed";
    priority: "high" | "medium" | "low";
  }>
>();

const todoWriteTool = tool(
  async ({ todos }) => {
    const sessionId = `session_${Date.now()}`;

    // Store in memory
    todoStorage.set(sessionId, todos);

    return {
      message: "Tasks planned successfully",
      session_id: sessionId,
      todos,
      total_tasks: todos.length,
    };
  },
  {
    name: "todo_write",
    description:
      "Create and manage structured task lists for multi-page processing sessions",
    schema: z.object({
      todos: z
        .array(
          z.object({
            id: z.string().describe("Unique identifier for the task"),
            content: z
              .string()
              .min(1)
              .describe("Description of the task to be completed"),
            status: z
              .enum(["pending", "completed"])
              .describe("Current status of the task"),
            priority: z
              .enum(["high", "medium", "low"])
              .describe("Priority level of the task"),
          })
        )
        .min(1)
        .describe("The complete todo list to store for this session"),
    }),
  }
);

const todoReadTool = tool(
  async ({ session_id }) => {
    const todos = todoStorage.get(session_id) || [];

    return {
      message: "Tasks retrieved successfully",
      session_id,
      todos,
      total_tasks: todos.length,
    };
  },
  {
    name: "todo_read",
    description: "Read current task list for the session",
    schema: z.object({
      session_id: z
        .string()
        .describe("Session identifier returned from todo_write"),
    }),
  }
);

const crawlTool = tool(
  async ({ url }) => {
    const base = process.env.CRAWL4AI_BASE_URL;
    if (!base) {
      throw new Error("CRAWL4AI_BASE_URL is not set");
    }

    try {
      const resp = await fetch(`${base}/md`, {
        method: "POST",
        headers: { "content-type": "application/json" },
        body: JSON.stringify({ url, f: "fit", c: "0", q: null }),
      });

      if (!resp.ok) {
        const txt = await resp.text();
        throw new Error(`crawl4ai error ${resp.status}: ${txt}`);
      }

      const content = await resp.text();

      // Extract title
      const titleMatch = content.match(/^#\s+(.+)$/m);
      const title = titleMatch ? titleMatch[1] : "Untitled";

      return {
        url,
        title,
        content,
        fetchedAt: new Date().toISOString(),
      };
    } catch (error) {
      throw new Error(
        `Error fetching ${url}: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  },
  {
    name: "crawl",
    description: "Get a page content from the web using crawl4ai",
    schema: z.object({
      url: z.string().url().describe("URL to crawl"),
    }),
  }
);

function createProcessPageTool(agent: Agent) {
  return tool(
    async ({ url, targetLanguage }) => {
      const prompt = `Process this page: ${url}\n\nTarget language: ${targetLanguage || "en"}`;

      const result = await agent.run(prompt);
      return result.data;
    },
    {
      name: "process_page",
      description: getToolDescription("process_page"),
      schema: TOOL_SCHEMA.process_page,
    }
  );
}

function createGenerateReportTool(agent: Agent) {
  return tool(
    async ({ targetLanguage, pageResults }) => {
      const prompt = `Generate comprehensive report from these page results:\n\n${JSON.stringify(pageResults)}\n\nTarget language: ${targetLanguage || "en"}`;

      const result = await agent.run(prompt);
      return result.data;
    },
    {
      name: "generate_report",
      description: getToolDescription("generate_report"),
      schema: TOOL_SCHEMA.generate_report,
    }
  );
}

function createDefaultTool(agent: Agent) {
  return tool(
    async ({ content }) => {
      const prompt = content;
      const result = await agent.run(prompt);
      return result.data;
    },
    {
      name: "default",
      description: getToolDescription("default"),
      schema: TOOL_SCHEMA.default,
    }
  );
}

function getToolDescription(toolName: string): string {
  switch (toolName) {
    case "process_page":
      return "Process single page: crawl, translate, summarize";
    case "generate_report":
      return "Generate comprehensive report from multiple page results";
    default:
      return `Tool for ${toolName}`;
  }
}

const TOOL_SCHEMA = {
  process_page: z.object({
    url: z.string().describe("URL to process"),
    targetLanguage: z.string().optional().describe("Target language code"),
  }),
  generate_report: z.object({
    pageResults: z.array(z.any()).describe("Results from processed pages"),
    targetLanguage: z.string().optional().describe("Target language code"),
  }),
  default: z.object({
    content: z.string().describe("Content to process"),
  }),
};

// No pre-created agent instance to avoid environment variable issues