From 748ee4d3c2b74bab077beac12f4213916c981e55 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Sat, 13 Dec 2025 22:22:17 +0530 Subject: [PATCH] feat(actions): add uploads search action --- .../search/researcher/actions/fileSearch.ts | 10 -- .../agents/search/researcher/actions/index.ts | 2 + .../researcher/actions/uploadsSearch.ts | 102 ++++++++++++++++++ 3 files changed, 104 insertions(+), 10 deletions(-) delete mode 100644 src/lib/agents/search/researcher/actions/fileSearch.ts create mode 100644 src/lib/agents/search/researcher/actions/uploadsSearch.ts diff --git a/src/lib/agents/search/researcher/actions/fileSearch.ts b/src/lib/agents/search/researcher/actions/fileSearch.ts deleted file mode 100644 index 0905e5d..0000000 --- a/src/lib/agents/search/researcher/actions/fileSearch.ts +++ /dev/null @@ -1,10 +0,0 @@ -import z from "zod"; -import { ResearchAction } from "../../types"; - -const schema = z.object({ - queries: z.array(z.string()).describe("A list of queries to search in files."), -}) - -const fileSearhAction: ResearchAction = { - name: "file_search", -} \ No newline at end of file diff --git a/src/lib/agents/search/researcher/actions/index.ts b/src/lib/agents/search/researcher/actions/index.ts index 15c5734..6d7e58f 100644 --- a/src/lib/agents/search/researcher/actions/index.ts +++ b/src/lib/agents/search/researcher/actions/index.ts @@ -2,11 +2,13 @@ import doneAction from './done'; import planAction from './plan'; import ActionRegistry from './registry'; import scrapeURLAction from './scrapeURL'; +import uploadsSearchAction from './uploadsSearch'; import webSearchAction from './webSearch'; ActionRegistry.register(webSearchAction); ActionRegistry.register(doneAction); ActionRegistry.register(planAction); ActionRegistry.register(scrapeURLAction); +ActionRegistry.register(uploadsSearchAction); export { ActionRegistry }; diff --git a/src/lib/agents/search/researcher/actions/uploadsSearch.ts b/src/lib/agents/search/researcher/actions/uploadsSearch.ts new file mode 100644 index 0000000..8195063 --- /dev/null +++ b/src/lib/agents/search/researcher/actions/uploadsSearch.ts @@ -0,0 +1,102 @@ +import z from 'zod'; +import { ResearchAction } from '../../types'; +import UploadStore from '@/lib/uploads/store'; + +const schema = z.object({ + queries: z + .array(z.string()) + .describe( + 'A list of queries to search in user uploaded files. Can be a maximum of 3 queries.', + ), +}); + +const uploadsSearchAction: ResearchAction = { + name: 'uploads_search', + enabled: (config) => + (config.classification.classification.personalSearch && + config.fileIds.length > 0) || + config.fileIds.length > 0, + schema, + getToolDescription: () => + `Use this tool to perform searches over the user's uploaded files. This is useful when you need to gather information from the user's documents to answer their questions. You can provide up to 3 queries at a time. You will have to use this every single time if this is present and relevant.`, + getDescription: () => ` + Use this tool to perform searches over the user's uploaded files. This is useful when you need to gather information from the user's documents to answer their questions. You can provide up to 3 queries at a time. You will have to use this every single time if this is present and relevant. + Always ensure that the queries you use are directly relevant to the user's request and pertain to the content of their uploaded files. + + For example, if the user says "Please find information about X in my uploaded documents", you can call this tool with a query related to X to retrieve the relevant information from their files. + Never use this tool to search the web or for information that is not contained within the user's uploaded files. + `, + execute: async (input, additionalConfig) => { + input.queries = input.queries.slice(0, 3); + + const researchBlock = additionalConfig.session.getBlock( + additionalConfig.researchBlockId, + ); + + if (researchBlock && researchBlock.type === 'research') { + researchBlock.data.subSteps.push({ + id: crypto.randomUUID(), + type: 'upload_searching', + queries: input.queries, + }); + + additionalConfig.session.updateBlock(additionalConfig.researchBlockId, [ + { + op: 'replace', + path: '/data/subSteps', + value: researchBlock.data.subSteps, + }, + ]); + } + + const uploadStore = new UploadStore({ + embeddingModel: additionalConfig.embedding, + fileIds: additionalConfig.fileIds, + }); + + const results = await uploadStore.query(input.queries, 10); + + const seenIds = new Map(); + + const filteredSearchResults = results + .map((result, index) => { + if (result.metadata.url && !seenIds.has(result.metadata.url)) { + seenIds.set(result.metadata.url, index); + return result; + } else if (result.metadata.url && seenIds.has(result.metadata.url)) { + const existingIndex = seenIds.get(result.metadata.url)!; + const existingResult = results[existingIndex]; + + existingResult.content += `\n\n${result.content}`; + + return undefined; + } + + return result; + }) + .filter((r) => r !== undefined); + + if (researchBlock && researchBlock.type === 'research') { + researchBlock.data.subSteps.push({ + id: crypto.randomUUID(), + type: 'upload_search_results', + results: filteredSearchResults, + }); + + additionalConfig.session.updateBlock(additionalConfig.researchBlockId, [ + { + op: 'replace', + path: '/data/subSteps', + value: researchBlock.data.subSteps, + }, + ]); + } + + return { + type: 'search_results', + results: filteredSearchResults, + }; + }, +}; + +export default uploadsSearchAction;