Realizando pesquisa aprofundada com subagentes

const spawnSubagentTool = {
  type: "function" as const,
  name: "spawn_subagent",
  description: "Delegue uma tarefa de pesquisa focada a uma resposta filha e retorne um memorando conciso.",
  strict: true,
  parameters: {
    type: "object",
    properties: {
      task: { type: "string" },
      instructions: { type: "string" },
      model: { type: "string" },
    },
    required: ["task"],
    additionalProperties: false,
  },
};

type SpawnSubagentArgs = {
  task: string;
  instructions?: string;
  model?: string;
};

const DEEP_RESEARCH_PARENT_INSTRUCTIONS = `
Você é uma orquestradora de pesquisa aprofundada.

Você não pesquisa diretamente.
Use a ferramenta spawn_subagent para delegar tarefas de pesquisa focadas a respostas filhas.

Regras de delegação:
- Para consultas factuais simples, inicie pelo menos 3 subagentes antes de responder.
- Para a maioria das tarefas de pesquisa, inicie pelo menos 4 subagentes antes de responder.
- Para pesquisas amplas, ambíguas ou de alto impacto, inicie 5 a 8 subagentes em rodadas.
- Inicie vários subagentes no mesmo turno sempre que seus trabalhos puderem prosseguir de forma independente.
- Se a primeira rodada estiver incompleta, repetitiva, fraca ou conflitante, inicie outra rodada em vez de adivinhar.

Exemplos de delegação:
- Um subagente para o caminho de resposta mais direto.
- Um subagente focado em fontes primárias ou oficiais.
- Um subagente focado em confirmação ou contradições.
- Subagentes adicionais para cronologia, geografia, perspectiva das partes interessadas, detalhes técnicos ou desenvolvimentos recentes.

Após o retorno de todos os resultados das ferramentas, sintetize as evidências em uma resposta para o usuário.
`;

const DEEP_RESEARCH_CHILD_INSTRUCTIONS = `
Você é um subagente especialista em pesquisa aprofundada.

Você tem a ferramenta web_search e deve usá-la agressivamente.
Concentre-se apenas na tarefa delegada.

Disciplina de busca:
- Para tarefas não triviais, faça pelo menos 3 buscas web distintas antes de retornar.
- Se a primeira busca for fraca, vazia, genérica ou composta principalmente por resultados de diretórios, continue buscando.
- Não retorne "próximos passos". Execute essas buscas você mesmo antes de retornar.
- Prefira fontes primárias, oficiais e outras fontes confiáveis.
- Trate diretórios e sites de perfis como pistas a confirmar, não como única prova quando uma fonte mais forte deveria existir.

Retorne um memorando conciso para outro modelo:
1. Principais descobertas
2. Evidências e fontes
3. Questões em aberto ou ressalvas
`;

import { SDK } from "@meetkai/mka1";
import type { ResponseObject } from "@meetkai/mka1/models/components";

const sdk = new SDK({
  bearerAuth: `Bearer ${process.env.MKA1_API_KEY}`,
});

async function runResearchChild(args: SpawnSubagentArgs) {
  const child = await sdk.llm.responses.create({
    responsesCreateRequest: {
      model: args.model ?? "meetkai:functionary-pt",
      instructions: args.instructions ?? DEEP_RESEARCH_CHILD_INSTRUCTIONS,
      input: args.task,
      tools: [
        {
          type: "web_search",
          userLocation: { country: "pakistan" },
          searchContextSize: "medium",
        },
      ],
      toolChoice: "auto",
      parallelToolCalls: true,
      maxToolCalls: 60,
      store: true,
    },
  }) as ResponseObject;

  return {
    response_id: child.id,
    output_text: child.outputText,
  };
}

const MIN_CHILD_SEARCHES = 3;
const MAX_CHILD_PASSES = 3;
const WEAK_MEMO_PATTERNS = [
  /next steps/i,
  /let me try/i,
  /need to search/i,
  /should check/i,
  /didn't find/i,
  /no direct results/i,
];

function isWeakMemo(outputText: string) {
  const trimmed = outputText.trim();
  if (!trimmed) return true;
  if (trimmed.length < 250) return true;
  return WEAK_MEMO_PATTERNS.some((pattern) => pattern.test(trimmed));
}

function buildChildContinuationPrompt(args: {
  task: string;
  searchCount: number;
  previousOutput: string;
}) {
  return [
    `Continue a tarefa de pesquisa delegada: ${args.task}`,
    `Você realizou apenas ${args.searchCount} buscas web até agora.`,
    `Faça pelo menos ${MIN_CHILD_SEARCHES} buscas web distintas antes de retornar, a menos que já tenha uma resposta direta de evidências confiáveis.`,
    `Não retorne "próximos passos" nem diga que é necessário pesquisar mais. Execute você mesmo as buscas adicionais.`,
    `Priorize fontes primárias, oficiais e outras fontes confiáveis. Use diretórios apenas como pistas para confirmação.`,
    ``,
    `Memorando anterior:`,
    args.previousOutput,
  ].join("\n");
}

async function runResearchChildWithContinuation(args: SpawnSubagentArgs) {
  let previousResponseId: string | undefined;
  let nextInput = args.task;
  let totalSearchCount = 0;

  for (let pass = 0; pass < MAX_CHILD_PASSES; pass += 1) {
    const child = await sdk.llm.responses.create({
      responsesCreateRequest: {
        model: args.model ?? "meetkai:functionary-pt",
        instructions: args.instructions ?? DEEP_RESEARCH_CHILD_INSTRUCTIONS,
        ...(previousResponseId
          ? { previousResponseId, input: nextInput }
          : { input: nextInput }),
        tools: [
          {
            type: "web_search",
            userLocation: { country: "pakistan" },
            searchContextSize: "medium",
          },
        ],
        toolChoice: "auto",
        parallelToolCalls: true,
        maxToolCalls: 60,
        store: true,
      },
    }) as ResponseObject;

    const passSearchCount = countCompletedWebSearches(child.output);
    totalSearchCount += passSearchCount;

    if (totalSearchCount >= MIN_CHILD_SEARCHES && !isWeakMemo(child.outputText ?? "")) {
      return {
        response_id: child.id,
        output_text: child.outputText,
        search_count: totalSearchCount,
      };
    }

    previousResponseId = child.id;
    nextInput = buildChildContinuationPrompt({
      task: args.task,
      searchCount: totalSearchCount,
      previousOutput: child.outputText ?? "",
    });
  }

  throw new Error("A resposta filha permaneceu muito fraca após várias etapas de pesquisa");
}

export async function runDeepResearch(input: string) {
  let response = await sdk.llm.responses.create({
    responsesCreateRequest: {
      model: "meetkai:functionary-urdu-large",
      instructions: DEEP_RESEARCH_PARENT_INSTRUCTIONS,
      input,
      tools: [spawnSubagentTool],
      toolChoice: "auto",
      parallelToolCalls: true,
      maxToolCalls: 18,
      store: true,
    },
  }) as ResponseObject;

  while (true) {
    const toolCalls = response.output.filter(
      (item): item is {
        type: "function_call";
        name: string;
        callId: string;
        arguments: string;
      } => item.type === "function_call" && item.name === "spawn_subagent",
    );

    if (toolCalls.length === 0) {
      return response.outputText;
    }

    const toolOutputs = await Promise.all(
      toolCalls.map(async (toolCall) => {
        const args = JSON.parse(toolCall.arguments) as SpawnSubagentArgs;
        const childResult = await runResearchChildWithContinuation(args);

        return {
          type: "function_call_output" as const,
          callId: toolCall.callId,
          output: JSON.stringify(childResult),
        };
      }),
    );

    response = await sdk.llm.responses.create({
      responsesCreateRequest: {
        model: response.model,
        previousResponseId: response.id,
        input: toolOutputs,
        tools: [spawnSubagentTool],
        parallelToolCalls: true,
        maxToolCalls: 18,
        store: true,
      },
    }) as ResponseObject;
  }
}