From 33b895b75eb349856f66528ef05817ff4cb472a5 Mon Sep 17 00:00:00 2001 From: ItzCrazyKns <95534749+ItzCrazyKns@users.noreply.github.com> Date: Thu, 20 Mar 2025 12:29:52 +0530 Subject: [PATCH] feat(app): add search API --- docs/API/SEARCH.md | 12 +-- src/app/api/chat/route.ts | 58 +------------ src/app/api/search/route.ts | 164 ++++++++++++++++++++++++++++++++++++ src/lib/search/index.ts | 59 +++++++++++++ 4 files changed, 230 insertions(+), 63 deletions(-) create mode 100644 src/app/api/search/route.ts create mode 100644 src/lib/search/index.ts diff --git a/docs/API/SEARCH.md b/docs/API/SEARCH.md index 9405bc5..0169976 100644 --- a/docs/API/SEARCH.md +++ b/docs/API/SEARCH.md @@ -20,11 +20,11 @@ The API accepts a JSON object in the request body, where you define the focus mo { "chatModel": { "provider": "openai", - "model": "gpt-4o-mini" + "name": "gpt-4o-mini" }, "embeddingModel": { "provider": "openai", - "model": "text-embedding-3-large" + "name": "text-embedding-3-large" }, "optimizationMode": "speed", "focusMode": "webSearch", @@ -38,18 +38,18 @@ The API accepts a JSON object in the request body, where you define the focus mo ### Request Parameters -- **`chatModel`** (object, optional): Defines the chat model to be used for the query. For model details you can send a GET request at `http://localhost:3001/api/models`. Make sure to use the key value (For example "gpt-4o-mini" instead of the display name "GPT 4 omni mini"). +- **`chatModel`** (object, optional): Defines the chat model to be used for the query. For model details you can send a GET request at `http://localhost:3000/api/models`. Make sure to use the key value (For example "gpt-4o-mini" instead of the display name "GPT 4 omni mini"). - `provider`: Specifies the provider for the chat model (e.g., `openai`, `ollama`). - - `model`: The specific model from the chosen provider (e.g., `gpt-4o-mini`). + - `name`: The specific model from the chosen provider (e.g., `gpt-4o-mini`). - Optional fields for custom OpenAI configuration: - `customOpenAIBaseURL`: If you’re using a custom OpenAI instance, provide the base URL. - `customOpenAIKey`: The API key for a custom OpenAI instance. -- **`embeddingModel`** (object, optional): Defines the embedding model for similarity-based searching. For model details you can send a GET request at `http://localhost:3001/api/models`. Make sure to use the key value (For example "text-embedding-3-large" instead of the display name "Text Embedding 3 Large"). +- **`embeddingModel`** (object, optional): Defines the embedding model for similarity-based searching. For model details you can send a GET request at `http://localhost:3000/api/models`. Make sure to use the key value (For example "text-embedding-3-large" instead of the display name "Text Embedding 3 Large"). - `provider`: The provider for the embedding model (e.g., `openai`). - - `model`: The specific embedding model (e.g., `text-embedding-3-large`). + - `name`: The specific embedding model (e.g., `text-embedding-3-large`). - **`focusMode`** (string, required): Specifies which focus mode to use. Available modes: diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index 918fd48..d9f9c6b 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -20,67 +20,11 @@ import { getCustomOpenaiApiUrl, getCustomOpenaiModelName, } from '@/lib/config'; +import { searchHandlers } from '@/lib/search'; export const runtime = 'nodejs'; export const dynamic = 'force-dynamic'; -const searchHandlers: Record = { - webSearch: new MetaSearchAgent({ - activeEngines: [], - queryGeneratorPrompt: prompts.webSearchRetrieverPrompt, - responsePrompt: prompts.webSearchResponsePrompt, - rerank: true, - rerankThreshold: 0.3, - searchWeb: true, - summarizer: true, - }), - academicSearch: new MetaSearchAgent({ - activeEngines: ['arxiv', 'google scholar', 'pubmed'], - queryGeneratorPrompt: prompts.academicSearchRetrieverPrompt, - responsePrompt: prompts.academicSearchResponsePrompt, - rerank: true, - rerankThreshold: 0, - searchWeb: true, - summarizer: false, - }), - writingAssistant: new MetaSearchAgent({ - activeEngines: [], - queryGeneratorPrompt: '', - responsePrompt: prompts.writingAssistantPrompt, - rerank: true, - rerankThreshold: 0, - searchWeb: false, - summarizer: false, - }), - wolframAlphaSearch: new MetaSearchAgent({ - activeEngines: ['wolframalpha'], - queryGeneratorPrompt: prompts.wolframAlphaSearchRetrieverPrompt, - responsePrompt: prompts.wolframAlphaSearchResponsePrompt, - rerank: false, - rerankThreshold: 0, - searchWeb: true, - summarizer: false, - }), - youtubeSearch: new MetaSearchAgent({ - activeEngines: ['youtube'], - queryGeneratorPrompt: prompts.youtubeSearchRetrieverPrompt, - responsePrompt: prompts.youtubeSearchResponsePrompt, - rerank: true, - rerankThreshold: 0.3, - searchWeb: true, - summarizer: false, - }), - redditSearch: new MetaSearchAgent({ - activeEngines: ['reddit'], - queryGeneratorPrompt: prompts.redditSearchRetrieverPrompt, - responsePrompt: prompts.redditSearchResponsePrompt, - rerank: true, - rerankThreshold: 0.3, - searchWeb: true, - summarizer: false, - }), -}; - type Message = { messageId: string; chatId: string; diff --git a/src/app/api/search/route.ts b/src/app/api/search/route.ts new file mode 100644 index 0000000..a2c1bd0 --- /dev/null +++ b/src/app/api/search/route.ts @@ -0,0 +1,164 @@ +import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; +import type { Embeddings } from '@langchain/core/embeddings'; +import { ChatOpenAI } from '@langchain/openai'; +import { + getAvailableChatModelProviders, + getAvailableEmbeddingModelProviders, +} from '@/lib/providers'; +import { AIMessage, BaseMessage, HumanMessage } from '@langchain/core/messages'; +import { MetaSearchAgentType } from '@/lib/search/metaSearchAgent'; +import { + getCustomOpenaiApiKey, + getCustomOpenaiApiUrl, + getCustomOpenaiModelName, +} from '@/lib/config'; +import { searchHandlers } from '@/lib/search'; + +interface chatModel { + provider: string; + name: string; + customOpenAIKey?: string; + customOpenAIBaseURL?: string; +} + +interface embeddingModel { + provider: string; + name: string; +} + +interface ChatRequestBody { + optimizationMode: 'speed' | 'balanced'; + focusMode: string; + chatModel?: chatModel; + embeddingModel?: embeddingModel; + query: string; + history: Array<[string, string]>; +} + +export const POST = async (req: Request) => { + try { + const body: ChatRequestBody = await req.json(); + + if (!body.focusMode || !body.query) { + return Response.json( + { message: 'Missing focus mode or query' }, + { status: 400 }, + ); + } + + body.history = body.history || []; + body.optimizationMode = body.optimizationMode || 'balanced'; + + const history: BaseMessage[] = body.history.map((msg) => { + return msg[0] === 'human' + ? new HumanMessage({ content: msg[1] }) + : new AIMessage({ content: msg[1] }); + }); + + const [chatModelProviders, embeddingModelProviders] = await Promise.all([ + getAvailableChatModelProviders(), + getAvailableEmbeddingModelProviders(), + ]); + + const chatModelProvider = + body.chatModel?.provider || Object.keys(chatModelProviders)[0]; + const chatModel = + body.chatModel?.name || + Object.keys(chatModelProviders[chatModelProvider])[0]; + + const embeddingModelProvider = + body.embeddingModel?.provider || Object.keys(embeddingModelProviders)[0]; + const embeddingModel = + body.embeddingModel?.name || + Object.keys(embeddingModelProviders[embeddingModelProvider])[0]; + + let llm: BaseChatModel | undefined; + let embeddings: Embeddings | undefined; + + if (body.chatModel?.provider === 'custom_openai') { + llm = new ChatOpenAI({ + modelName: body.chatModel?.name || getCustomOpenaiModelName(), + openAIApiKey: + body.chatModel?.customOpenAIKey || getCustomOpenaiApiKey(), + temperature: 0.7, + configuration: { + baseURL: + body.chatModel?.customOpenAIBaseURL || getCustomOpenaiApiUrl(), + }, + }) as unknown as BaseChatModel; + } else if ( + chatModelProviders[chatModelProvider] && + chatModelProviders[chatModelProvider][chatModel] + ) { + llm = chatModelProviders[chatModelProvider][chatModel] + .model as unknown as BaseChatModel | undefined; + } + + if ( + embeddingModelProviders[embeddingModelProvider] && + embeddingModelProviders[embeddingModelProvider][embeddingModel] + ) { + embeddings = embeddingModelProviders[embeddingModelProvider][ + embeddingModel + ].model as Embeddings | undefined; + } + + if (!llm || !embeddings) { + return Response.json( + { message: 'Invalid model selected' }, + { status: 400 }, + ); + } + + const searchHandler: MetaSearchAgentType = searchHandlers[body.focusMode]; + + if (!searchHandler) { + return Response.json({ message: 'Invalid focus mode' }, { status: 400 }); + } + + const emitter = await searchHandler.searchAndAnswer( + body.query, + history, + llm, + embeddings, + body.optimizationMode, + [], + ); + + return new Promise((resolve, reject) => { + let message = ''; + let sources: any[] = []; + + emitter.on('data', (data) => { + try { + const parsedData = JSON.parse(data); + if (parsedData.type === 'response') { + message += parsedData.data; + } else if (parsedData.type === 'sources') { + sources = parsedData.data; + } + } catch (error) { + reject( + Response.json({ message: 'Error parsing data' }, { status: 500 }), + ); + } + }); + + emitter.on('end', () => { + resolve(Response.json({ message, sources }, { status: 200 })); + }); + + emitter.on('error', (error) => { + reject( + Response.json({ message: 'Search error', error }, { status: 500 }), + ); + }); + }); + } catch (err: any) { + console.error(`Error in getting search results: ${err.message}`); + return Response.json( + { message: 'An error has occurred.' }, + { status: 500 }, + ); + } +}; diff --git a/src/lib/search/index.ts b/src/lib/search/index.ts new file mode 100644 index 0000000..01fb87f --- /dev/null +++ b/src/lib/search/index.ts @@ -0,0 +1,59 @@ +import MetaSearchAgent from '@/lib/search/metaSearchAgent'; +import prompts from '../prompts'; + +export const searchHandlers: Record = { + webSearch: new MetaSearchAgent({ + activeEngines: [], + queryGeneratorPrompt: prompts.webSearchRetrieverPrompt, + responsePrompt: prompts.webSearchResponsePrompt, + rerank: true, + rerankThreshold: 0.3, + searchWeb: true, + summarizer: true, + }), + academicSearch: new MetaSearchAgent({ + activeEngines: ['arxiv', 'google scholar', 'pubmed'], + queryGeneratorPrompt: prompts.academicSearchRetrieverPrompt, + responsePrompt: prompts.academicSearchResponsePrompt, + rerank: true, + rerankThreshold: 0, + searchWeb: true, + summarizer: false, + }), + writingAssistant: new MetaSearchAgent({ + activeEngines: [], + queryGeneratorPrompt: '', + responsePrompt: prompts.writingAssistantPrompt, + rerank: true, + rerankThreshold: 0, + searchWeb: false, + summarizer: false, + }), + wolframAlphaSearch: new MetaSearchAgent({ + activeEngines: ['wolframalpha'], + queryGeneratorPrompt: prompts.wolframAlphaSearchRetrieverPrompt, + responsePrompt: prompts.wolframAlphaSearchResponsePrompt, + rerank: false, + rerankThreshold: 0, + searchWeb: true, + summarizer: false, + }), + youtubeSearch: new MetaSearchAgent({ + activeEngines: ['youtube'], + queryGeneratorPrompt: prompts.youtubeSearchRetrieverPrompt, + responsePrompt: prompts.youtubeSearchResponsePrompt, + rerank: true, + rerankThreshold: 0.3, + searchWeb: true, + summarizer: false, + }), + redditSearch: new MetaSearchAgent({ + activeEngines: ['reddit'], + queryGeneratorPrompt: prompts.redditSearchRetrieverPrompt, + responsePrompt: prompts.redditSearchResponsePrompt, + rerank: true, + rerankThreshold: 0.3, + searchWeb: true, + summarizer: false, + }), +};