diff --git a/src/lib/agents/search/api.ts b/src/lib/agents/search/api.ts new file mode 100644 index 0000000..924bc68 --- /dev/null +++ b/src/lib/agents/search/api.ts @@ -0,0 +1,99 @@ +import { ResearcherOutput, SearchAgentInput } from './types'; +import SessionManager from '@/lib/session'; +import { classify } from './classifier'; +import Researcher from './researcher'; +import { getWriterPrompt } from '@/lib/prompts/search/writer'; +import { WidgetExecutor } from './widgets'; + +class APISearchAgent { + async searchAsync(session: SessionManager, input: SearchAgentInput) { + const classification = await classify({ + chatHistory: input.chatHistory, + enabledSources: input.config.sources, + query: input.followUp, + llm: input.config.llm, + }); + + const widgetPromise = WidgetExecutor.executeAll({ + classification, + chatHistory: input.chatHistory, + followUp: input.followUp, + llm: input.config.llm, + }); + + let searchPromise: Promise | null = null; + + if (!classification.classification.skipSearch) { + const researcher = new Researcher(); + searchPromise = researcher.research(SessionManager.createSession(), { + chatHistory: input.chatHistory, + followUp: input.followUp, + classification: classification, + config: input.config, + }); + } + + const [widgetOutputs, searchResults] = await Promise.all([ + widgetPromise, + searchPromise, + ]); + + if (searchResults) { + session.emit('data', { + type: 'searchResults', + data: searchResults.searchFindings, + }); + } + + session.emit('data', { + type: 'researchComplete', + }); + + const finalContext = + searchResults?.searchFindings + .map( + (f, index) => + `${f.content}`, + ) + .join('\n') || ''; + + const widgetContext = widgetOutputs + .map((o) => { + return `${o.llmContext}`; + }) + .join('\n-------------\n'); + + const finalContextWithWidgets = `\n${finalContext}\n\n\n${widgetContext}\n`; + + const writerPrompt = getWriterPrompt( + finalContextWithWidgets, + input.config.systemInstructions, + input.config.mode, + ); + + const answerStream = input.config.llm.streamText({ + messages: [ + { + role: 'system', + content: writerPrompt, + }, + ...input.chatHistory, + { + role: 'user', + content: input.followUp, + }, + ], + }); + + for await (const chunk of answerStream) { + session.emit('data', { + type: 'response', + data: chunk.contentChunk, + }); + } + + session.emit('end', {}); + } +} + +export default APISearchAgent;