From 19b1600584429dcb3485001e559155ebad975529 Mon Sep 17 00:00:00 2001 From: Marcel Samyn Date: Wed, 28 May 2025 12:28:26 +0200 Subject: [PATCH 1/3] feat(search): iterative deep research --- src/lib/formatting.ts | 29 ++++++- src/lib/jobs/deep-research.ts | 148 ++++++++++++++++++++++++++++++++-- 2 files changed, 170 insertions(+), 7 deletions(-) diff --git a/src/lib/formatting.ts b/src/lib/formatting.ts index f57dca3..c32f4d7 100644 --- a/src/lib/formatting.ts +++ b/src/lib/formatting.ts @@ -91,7 +91,6 @@ ${xmlItems} `; } - // Group definitions for reranked search results type SearchGroups = { similarNodes: NodeSearchResult; @@ -149,3 +148,31 @@ export function formatSearchResultsAsXml(results: SearchResults): string { : ""; return body; } + +export type SearchResultWithId = SearchResults[number] & { tempId: string }; + +/** + * Format search results with temporary IDs so the LLM can reference them. + */ +export function formatSearchResultsWithIds( + results: SearchResultWithId[], +): string { + const body = results.length + ? results + .map((r) => { + const inner = (() => { + switch (r.group) { + case "similarNodes": + return formatSearchNode(r.item); + case "similarEdges": + return formatSearchEdge(r.item); + case "connections": + return formatSearchConnection(r.item); + } + })(); + return `${inner}`; + }) + .join("\n") + : ""; + return body; +} diff --git a/src/lib/jobs/deep-research.ts b/src/lib/jobs/deep-research.ts index 78a82e3..066b16d 100644 --- a/src/lib/jobs/deep-research.ts +++ b/src/lib/jobs/deep-research.ts @@ -1,6 +1,11 @@ import { performStructuredAnalysis } from "../ai"; import { storeDeepResearchResult } from "../cache/deep-research-cache"; import { generateEmbeddings } from "../embeddings"; +import { + formatSearchResultsWithIds, + type SearchResultWithId, + type SearchResults, +} from "../formatting"; import { findOneHopNodes, findSimilarEdges, @@ -14,6 +19,7 @@ import { DeepResearchJobInput, DeepResearchResult, } from "../schemas/deep-research"; +import { TemporaryIdMapper } from "../temporary-id-mapper"; import { z } from "zod"; import { DrizzleDB } from "~/db"; import { useDatabase } from "~/utils/db"; @@ -28,6 +34,8 @@ type SearchGroups = { // Default TTL for deep research results (24 hours) const DEFAULT_TTL_SECONDS = 24 * 60 * 60; +// Maximum number of refinement loops +const MAX_SEARCH_LOOPS = 4; /** * Main job handler for deep research @@ -42,8 +50,7 @@ export async function performDeepResearch( console.log(`Starting deep research for conversation ${conversationId}`); try { - // Get search queries based on recent conversation turns - // Filter to only include user and assistant messages + // Prepare initial queries based on recent conversation turns const recentMessages = messages .slice(-lastNMessages) .filter((m) => m.role === "user" || m.role === "assistant"); @@ -54,11 +61,16 @@ export async function performDeepResearch( return; } - // Execute search queries and aggregate results - const searchResults = await executeDeepSearchQueries(db, userId, queries); + // Run iterative search/refine loop + const searchResults = await runIterativeSearch( + db, + userId, + recentMessages, + queries, + ); - // Process results and cache them - await cacheDeepResearchResults(userId, conversationId, searchResults); + // Cache the combined results + await cacheDeepResearchResults(userId, conversationId, [searchResults]); console.log(`Deep research completed for conversation ${conversationId}`); } catch (error) { @@ -110,6 +122,130 @@ Come up with 1-5 search queries that explore adjacent or less obvious connection } } +/** + * Run iterative search with LLM refinement. + */ +async function runIterativeSearch( + db: DrizzleDB, + userId: string, + messages: DeepResearchJobInput["messages"], + initialQueries: string[], +): Promise> { + const queue = [...initialQueries]; + const history: string[] = []; + let results: SearchResultWithId[] = []; + const mapper = new TemporaryIdMapper( + (_item, idx) => `r${idx + 1}`, + ); + const seen = new Set(); + let loops = 0; + + while (loops < MAX_SEARCH_LOOPS && queue.length > 0) { + const query = queue.shift()!; + history.push(query); + + const embResp = await generateEmbeddings({ + model: "jina-embeddings-v3", + task: "retrieval.query", + input: [query], + truncate: true, + }); + const embedding = embResp.data[0]?.embedding; + if (embedding) { + const res = await executeSearchWithEmbedding( + db, + userId, + query, + embedding, + 20, + ); + if (res) { + const dedup = res.filter((r) => { + const key = `${r.group}:${(r.item as any).id}`; + if (seen.has(key)) return false; + seen.add(key); + return true; + }); + results.push(...mapper.mapItems(dedup)); + } + } + + loops++; + if (loops >= MAX_SEARCH_LOOPS) break; + + const refinement = await refineSearchResults( + userId, + messages, + history, + results, + ); + if (refinement.dropIds.length) { + const drop = new Set(refinement.dropIds); + results = results.filter((r) => !drop.has(r.tempId)); + } + if (refinement.done) break; + if (refinement.nextQuery) queue.push(refinement.nextQuery); + } + + return results.map(({ tempId, ...rest }) => rest); +} + +interface RefinementResult { + dropIds: string[]; + done: boolean; + nextQuery?: string; +} + +/** + * Ask the LLM to refine search results. + */ +async function refineSearchResults( + userId: string, + messages: DeepResearchJobInput["messages"], + queries: string[], + results: SearchResultWithId[], +): Promise { + const schema = z + .object({ + dropIds: z.array(z.string()).default([]), + done: z.boolean(), + nextQuery: z.string().optional(), + }) + .describe("DeepResearchRefinement"); + + const messageContext = messages + .map((m) => `${m.content}`) + .join("\n"); + const queriesXml = queries.map((q) => `${q}`).join("\n"); + const resultsXml = formatSearchResultsWithIds(results); + + try { + return await performStructuredAnalysis({ + userId, + systemPrompt: "You refine background search results.", + prompt: ` +${messageContext} + + + +${queriesXml} + + + +${resultsXml} + + + +Remove irrelevant results by listing their ids in dropIds. If more searching is needed, set done=false and provide nextQuery. If satisfied, set done=true. +`, + schema, + }); + } catch (error) { + console.error("Failed to refine deep search results:", error); + return { dropIds: [], done: true }; + } +} + /** * Execute multiple search queries in parallel with higher limits * and return combined results From 6edecd6f168bf9c08d593cb1654dac05b38a3e20 Mon Sep 17 00:00:00 2001 From: Marcel Samyn Date: Tue, 7 Oct 2025 13:05:07 +0200 Subject: [PATCH 2/3] Harden deep research search loop --- src/lib/formatting.ts | 34 ++++++++++++++++++++++++++-------- src/lib/jobs/deep-research.ts | 14 +++++++++----- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/src/lib/formatting.ts b/src/lib/formatting.ts index c32f4d7..f562ece 100644 --- a/src/lib/formatting.ts +++ b/src/lib/formatting.ts @@ -29,7 +29,7 @@ export function formatConversationAsXml(messages: Message[]): string { } /** Escape special characters for XML */ -function escapeXml(str: string): string { +export function escapeXml(str: string): string { return str .replace(/&/g, "&") .replace(/ ` - - ${node.description || ""} + + ${escapeXml(node.description || "")} `, ) .join("\n"); @@ -82,8 +82,8 @@ export function formatLabelDescList( const xmlItems = items .map( - (item) => `${item.description ?? ""}`, + (item) => + `${escapeXml(item.description ?? "")}`, ) .join("\n"); return ` @@ -106,8 +106,8 @@ export type SearchResults = RerankResult; // Helpers for formatting individual result items function formatSearchNode(node: NodeSearchResult): string { return ` - - ${node.description ?? ""} + + ${escapeXml(node.description ?? "")} `; } @@ -123,10 +123,14 @@ function formatSearchConnection(conn: OneHopNode): string { return ` - ${conn.description ?? ""} + ${escapeXml(conn.description ?? "")} `; } +function assertNever(value: never, message: string): never { + throw new Error(message); +} + /** * Formats reranked search results as an XML-like structure for LLM prompts. * Items are ordered by descending relevance and tagged by their group. @@ -142,6 +146,13 @@ export function formatSearchResultsAsXml(results: SearchResults): string { return formatSearchEdge(r.item); case "connections": return formatSearchConnection(r.item); + default: + return assertNever( + r.group, + `[formatSearchResultsAsXml] Unhandled search result group: ${String( + r.group, + )}`, + ); } }) .join("\n") @@ -168,6 +179,13 @@ export function formatSearchResultsWithIds( return formatSearchEdge(r.item); case "connections": return formatSearchConnection(r.item); + default: + return assertNever( + r.group, + `[formatSearchResultsWithIds] Unhandled search result group: ${String( + r.group, + )}`, + ); } })(); return `${inner}`; diff --git a/src/lib/jobs/deep-research.ts b/src/lib/jobs/deep-research.ts index 066b16d..106d181 100644 --- a/src/lib/jobs/deep-research.ts +++ b/src/lib/jobs/deep-research.ts @@ -2,6 +2,7 @@ import { performStructuredAnalysis } from "../ai"; import { storeDeepResearchResult } from "../cache/deep-research-cache"; import { generateEmbeddings } from "../embeddings"; import { + escapeXml, formatSearchResultsWithIds, type SearchResultWithId, type SearchResults, @@ -94,7 +95,7 @@ async function generateSearchQueries( // Format messages for context const messageContext = messages - .map((m) => `${m.content}`) + .map((m) => `${escapeXml(m.content)}`) .join("\n"); // Use structured analysis to generate tangential search queries @@ -134,8 +135,9 @@ async function runIterativeSearch( const queue = [...initialQueries]; const history: string[] = []; let results: SearchResultWithId[] = []; + let tempIdCounter = 0; const mapper = new TemporaryIdMapper( - (_item, idx) => `r${idx + 1}`, + () => `r${++tempIdCounter}`, ); const seen = new Set(); let loops = 0; @@ -161,7 +163,7 @@ async function runIterativeSearch( ); if (res) { const dedup = res.filter((r) => { - const key = `${r.group}:${(r.item as any).id}`; + const key = `${r.group}:${r.item.id}`; if (seen.has(key)) return false; seen.add(key); return true; @@ -214,9 +216,11 @@ async function refineSearchResults( .describe("DeepResearchRefinement"); const messageContext = messages - .map((m) => `${m.content}`) + .map((m) => `${escapeXml(m.content)}`) + .join("\n"); + const queriesXml = queries + .map((q) => `${escapeXml(q)}`) .join("\n"); - const queriesXml = queries.map((q) => `${q}`).join("\n"); const resultsXml = formatSearchResultsWithIds(results); try { From 6877f90c12388cc59bea1ab1c5d00797d3910ae3 Mon Sep 17 00:00:00 2001 From: Marcel Samyn Date: Tue, 7 Oct 2025 13:08:51 +0200 Subject: [PATCH 3/3] fixup --- package.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/package.json b/package.json index 047a0ae..473e736 100644 --- a/package.json +++ b/package.json @@ -70,7 +70,9 @@ }, "pnpm": { "onlyBuiltDependencies": [ - "esbuild" + "@parcel/watcher", + "esbuild", + "msgpackr-extract" ] } }