Merge branch 'main' of https://github.com/hatchet-dev/pickaxe

Author: abelanger5

pnpm-lock.yaml

@@ -0,0 +1,191 @@
+import { search } from "@/agents/deep-research/tools/search.tool";
+import { planSearch, PlanSearchOutput } from "@/agents/deep-research/tools/plan-search.tool";
+import { websiteToMd } from "@/agents/deep-research/tools/website-to-md.tool";
+import { summarize } from "@/agents/deep-research/tools/summarize.tool";
+import { judgeResults } from "@/agents/deep-research/tools/judge-results.tool";
+import { extractFacts } from "@/agents/deep-research/tools/extract-facts.tool";
+import { judgeFacts, JudgeFactsOutput } from "@/agents/deep-research/tools/judge-facts.tool";
+import { z } from "zod";
+import { pickaxe } from "@/pickaxe-client";
+
+const MessageSchema = z.object({
+  message: z.string(),
+});
+
+const SourceSchema = z.object({
+  url: z.string(),
+  title: z.string().optional(),
+  index: z.number(),
+});
+
+const ResponseSchema = z.object({
+  result: z.object({
+    isComplete: z.boolean(),
+    reason: z.string(),
+    sources: z.array(SourceSchema),
+    summary: z.string().optional(),
+    facts: z.array(z.object({
+      text: z.string(),
+      sourceIndex: z.number(),
+    })).optional(),
+    iterations: z.number().optional(),
+    factsJudgment: z.object({
+      reason: z.string(),
+      hasEnoughFacts: z.boolean(),
+      missingAspects: z.array(z.string()),
+    }).optional(),
+    searchPlans: z.string().optional(),
+  }),
+});
+
+
+type Source = z.infer<typeof SourceSchema>;
+type Fact = {
+  text: string;
+  sourceIndex: number;
+};
+
+export const deepResearchAgent = pickaxe.agent({
+  name: "deep-research-agent",
+  description: "A tool that performs deep research on a given query",
+  inputSchema: MessageSchema,
+  outputSchema: ResponseSchema,
+  executionTimeout: "15m",
+  fn: async (input, ctx) => {
+    ctx.logger.info(`Starting deep research agent with query: ${input.message}`);
+
+    let iteration = 0;
+    const maxIterations = 3;
+    const allFacts: Fact[] = [];
+    const allSources: Source[] = [];
+    let missingAspects: string[] = [];
+    let plan: PlanSearchOutput | undefined = undefined;
+    let factsJudgment: JudgeFactsOutput | undefined = undefined;
+
+    while (!ctx.cancelled && iteration < maxIterations) {
+      iteration++;
+      ctx.logger.info(`Starting iteration ${iteration}/${maxIterations}`);
+
+      // Plan the search based on the query, existing facts, and missing aspects
+      ctx.logger.info(
+        `Planning search with ${allFacts.length} existing facts and ${missingAspects.length} missing aspects`
+      );
+
+      plan = await planSearch.run({
+        query: input.message,
+        existingFacts: allFacts.map((f) => f.text),
+        missingAspects: missingAspects,
+      });
+
+      ctx.logger.info(
+        `Search plan for iteration ${iteration}: ${plan.reasoning}. Queries:`
+      );
+
+      for (const query of plan.queries) {
+        ctx.logger.info(`${query}`);
+      }
+
+      ctx.logger.info(`Executing ${plan.queries.length} search queries`);
+      const results = await search.run  (
+        plan.queries.map((query: string) => ({ query }))
+      );
+
+      // Flatten and deduplicate sources
+      const newSources = results.flatMap((result) => result.sources);
+      const uniqueSources = new Map(
+        newSources.map((source, index) => [source.url, { ...source, index }])
+      );
+
+      ctx.logger.info(
+        `Found ${newSources.length} new sources, ${uniqueSources.size} unique sources`
+      );
+
+      // Add new sources to all sources
+      allSources.push(...Array.from(uniqueSources.values()));
+
+      // Convert sources to markdown
+      ctx.logger.info(`Converting ${uniqueSources.size} sources to markdown`);
+      const mdResults = await websiteToMd.run(
+        Array.from(uniqueSources.values())
+          .sort((a, b) => a.index - b.index)
+          .map((source) => ({
+              url: source.url,
+              index: source.index,
+              title: source.title || "",
+            }))
+      );
+
+      // Extract facts from each source
+      ctx.logger.info("Extracting facts from markdown content");
+      const factsResults = await extractFacts.run(
+        mdResults.map((result) => ({
+            source: result.markdown,
+            query: input.message,
+            sourceInfo: {
+              url: result.url,
+              title: result.title,
+              index: result.index,
+            },
+        }))
+      );
+
+      // Add new facts to all facts
+      const newFacts = factsResults.flatMap((result) => result.facts);
+      allFacts.push(...newFacts);
+      ctx.logger.info(
+        `Extracted ${newFacts.length} new facts, total facts: ${allFacts.length}`
+      );
+
+      // Judge if we have enough facts
+      ctx.logger.info("Judging if we have enough facts");
+      factsJudgment = await judgeFacts.run({
+        query: input.message,
+        facts: allFacts.map((f) => f.text),
+      });
+
+      // Update missing aspects for next iteration
+      missingAspects = factsJudgment.missingAspects;
+      ctx.logger.info(`Missing aspects: ${missingAspects.join(", ")}`);
+
+      // If we have enough facts or reached max iterations, generate final summary
+      if (factsJudgment.hasEnoughFacts || iteration >= maxIterations) {
+        ctx.logger.info(
+          `Generating final summary (hasEnoughFacts: ${
+            factsJudgment.hasEnoughFacts
+          }, reachedMaxIterations: ${iteration >= maxIterations})`
+        );
+        break;
+      }
+    }
+
+    // Always summarize and judge results after the loop
+    const summarizeResult = await summarize.run({
+      text: input.message,
+      facts: allFacts,
+      sources: allSources,
+    });
+
+    ctx.logger.info("Judging final results");
+    const judgeResult = await judgeResults.run({
+      query: input.message,
+      result: summarizeResult.summary,
+    });
+
+    ctx.logger.info(
+      `Deep research complete (isComplete: ${judgeResult.isComplete}, totalFacts: ${allFacts.length}, totalSources: ${allSources.length}, iterations: ${iteration})`
+    );
+
+    return {
+      result: {
+        isComplete: judgeResult.isComplete,
+        reason: judgeResult.reason,
+        sources: allSources,
+        summary: summarizeResult.summary,
+        facts: allFacts,
+        iterations: iteration,
+        factsJudgment: factsJudgment,
+        searchPlans: plan?.reasoning,
+      },
+    };
+  },
+});

scaffolds/src/agents/deep-research/deep-research.agent.ts

@@ -0,0 +1,10 @@
+import { search } from '@/agents/deep-research/tools/search.tool';
+import { summarize } from '@/agents/deep-research/tools/summarize.tool';
+import { pickaxe } from '@/pickaxe-client';
+
+export const deepResearchTaskbox = pickaxe.toolbox({
+  tools: [
+    search,
+    summarize,
+  ],
+});

scaffolds/src/agents/deep-research/deep-research.toolbox.ts

@@ -0,0 +1,3 @@
+export * from './deep-research.agent';
+export * from './deep-research.toolbox';
+export * from './tools';

scaffolds/src/agents/deep-research/index.ts

@@ -0,0 +1,57 @@
+import { z } from "zod";
+import { generateObject } from "ai";
+import { openai } from "@ai-sdk/openai";
+import { pickaxe } from "@/pickaxe-client";
+
+const ExtractFactsInputSchema = z.object({
+  source: z.string(),
+  query: z.string(),
+  sourceInfo: z.object({
+    url: z.string(),
+    title: z.string().optional(),
+    index: z.number(),
+  }),
+});
+
+type ExtractFactsInput = z.infer<typeof ExtractFactsInputSchema>;
+
+const FactSchema = z.object({
+  text: z.string(),
+  sourceIndex: z.number(),
+});
+
+const ExtractFactsOutputSchema = z.object({
+  facts: z.array(FactSchema),
+});
+
+export const extractFacts = pickaxe.tool({
+  name: "extract-facts",
+  description: "Extract relevant facts from a source that are related to a query",
+  inputSchema: ExtractFactsInputSchema,
+  outputSchema: ExtractFactsOutputSchema,
+  fn: async (input, ctx) => {
+    const result = await generateObject({
+      abortSignal: ctx.abortController.signal,
+      prompt: `
+Extract relevant facts from the following source that are related to this query:
+"""${input.query}"""
+
+Source:
+"""${input.source}"""
+
+Extract only factual statements that are directly relevant to the query. Each fact should be a complete, standalone statement.
+`,
+      model: openai("gpt-4.1-mini"),
+      schema: z.object({
+        facts: z.array(z.string()),
+      }),
+    });
+
+    return {
+      facts: result.object.facts.map((fact) => ({
+        text: fact,
+        sourceIndex: input.sourceInfo.index,
+      })),
+    };
+  },
+});

scaffolds/src/agents/deep-research/tools/extract-facts.tool.ts

@@ -0,0 +1,6 @@
+export * from './plan-search.tool';
+export * from './search.tool';
+export * from './summarize.tool';
+export * from './extract-facts.tool';
+export * from './judge-results.tool';
+export * from './judge-facts.tool';

scaffolds/src/agents/deep-research/tools/index.ts

@@ -0,0 +1,54 @@
+import { z } from "zod";
+import { generateObject } from "ai";
+import { openai } from "@ai-sdk/openai";
+import { pickaxe } from "@/pickaxe-client";
+
+const JudgeFactsInputSchema = z.object({
+  query: z.string(),
+  facts: z.array(z.string()),
+});
+
+const JudgeFactsOutputSchema = z.object({
+  hasEnoughFacts: z.boolean(),
+  reason: z.string(),
+  missingAspects: z.array(z.string()),
+});
+
+export type JudgeFactsOutput = z.infer<typeof JudgeFactsOutputSchema>;
+
+export const judgeFacts = pickaxe.tool({
+  name: "judge-facts",
+  description: "Judge if we have enough facts to comprehensively answer a query",
+  inputSchema: JudgeFactsInputSchema,
+  outputSchema: JudgeFactsOutputSchema,
+  fn: async (input, ctx) => {
+    const result = await generateObject({
+      abortSignal: ctx.abortController.signal,
+      prompt: `
+Evaluate if we have enough facts to comprehensively answer this query:
+"""${input.query}"""
+
+Current facts:
+${input.facts.map((fact, i) => `${i + 1}. ${fact}`).join("\n")}
+
+Consider:
+1. Are there any key aspects of the query that aren't covered by the current facts?
+2. Are the facts diverse enough to provide a complete picture?
+3. Are there any gaps in the information that would prevent a comprehensive answer?
+4. Are there any technical jargon words that are not defined in the facts that require additional research?
+`,
+      model: openai("gpt-4.1-mini"),
+      schema: z.object({
+        hasEnoughFacts: z.boolean(),
+        reason: z.string(),
+        missingAspects: z.array(z.string()),
+      }),
+    });
+
+    return {
+      hasEnoughFacts: result.object.hasEnoughFacts,
+      reason: result.object.reason,
+      missingAspects: result.object.missingAspects,
+    };
+  },
+});