diff --git a/sample.config.toml b/sample.config.toml index 980e99d..e2fa352 100644 --- a/sample.config.toml +++ b/sample.config.toml @@ -26,4 +26,8 @@ API_URL = "" # Ollama API URL - http://host.docker.internal:11434 API_KEY = "" [API_ENDPOINTS] -SEARXNG = "" # SearxNG API URL - http://localhost:32768 \ No newline at end of file +SEARXNG = "" # SearxNG API URL - http://localhost:32768 +TAVILY = "" # Tavily API key + +[SEARCH] +ENGINE = "searxng" # "searxng" or "tavily" \ No newline at end of file diff --git a/src/app/api/config/route.ts b/src/app/api/config/route.ts index 39c1f84..f203ea0 100644 --- a/src/app/api/config/route.ts +++ b/src/app/api/config/route.ts @@ -8,6 +8,8 @@ import { getOllamaApiEndpoint, getOpenaiApiKey, getDeepseekApiKey, + getSearchEngine, + getTavilyApiKey, updateConfig, } from '@/lib/config'; import { @@ -58,6 +60,8 @@ export const GET = async (req: Request) => { config['customOpenaiApiUrl'] = getCustomOpenaiApiUrl(); config['customOpenaiApiKey'] = getCustomOpenaiApiKey(); config['customOpenaiModelName'] = getCustomOpenaiModelName(); + config['searchEngine'] = getSearchEngine(); + config['tavilyApiKey'] = getTavilyApiKey(); return Response.json({ ...config }, { status: 200 }); } catch (err) { @@ -99,6 +103,12 @@ export const POST = async (req: Request) => { MODEL_NAME: config.customOpenaiModelName, }, }, + SEARCH: { + ENGINE: config.searchEngine, + }, + API_ENDPOINTS: { + TAVILY: config.tavilyApiKey || '', + }, }; updateConfig(updatedConfig); diff --git a/src/app/api/discover/route.ts b/src/app/api/discover/route.ts index 8c1f470..502295c 100644 --- a/src/app/api/discover/route.ts +++ b/src/app/api/discover/route.ts @@ -1,4 +1,4 @@ -import { searchSearxng } from '@/lib/searxng'; +import { searchSearxng } from '../../../lib/searchEngines/searxng'; const articleWebsites = [ 'yahoo.com', diff --git a/src/app/settings/page.tsx b/src/app/settings/page.tsx index 8eee9a4..6988a3d 100644 --- a/src/app/settings/page.tsx +++ b/src/app/settings/page.tsx @@ -24,6 +24,8 @@ interface SettingsType { customOpenaiApiKey: string; customOpenaiApiUrl: string; customOpenaiModelName: string; + searchEngine: string; + tavilyApiKey?: string; } interface InputProps extends React.InputHTMLAttributes { @@ -145,6 +147,7 @@ const Page = () => { const [automaticImageSearch, setAutomaticImageSearch] = useState(false); const [automaticVideoSearch, setAutomaticVideoSearch] = useState(false); const [systemInstructions, setSystemInstructions] = useState(''); + const [searchEngine, setSearchEngine] = useState('searxng'); const [savingStates, setSavingStates] = useState>({}); useEffect(() => { @@ -207,6 +210,7 @@ const Page = () => { ); setSystemInstructions(localStorage.getItem('systemInstructions')!); + setSearchEngine(localStorage.getItem('searchEngine') || 'searxng'); setIsLoading(false); }; @@ -366,6 +370,10 @@ const Page = () => { localStorage.setItem('embeddingModel', value); } else if (key === 'systemInstructions') { localStorage.setItem('systemInstructions', value); + } else if (key === 'searchEngine') { + localStorage.setItem('searchEngine', value); + } else if (key === 'tavilyApiKey') { + localStorage.setItem('tavilyApiKey', value); } } catch (err) { console.error('Failed to save:', err); @@ -508,6 +516,32 @@ const Page = () => { /> + +
+

+ Search Engine +

+ { + setConfig((prev) => ({ + ...prev!, + tavilyApiKey: e.target.value, + })); + }} + onSave={(value) => saveConfig('tavilyApiKey', value)} + /> +
diff --git a/src/lib/chains/imageSearchAgent.ts b/src/lib/chains/imageSearchAgent.ts index 4fd684f..1381c3c 100644 --- a/src/lib/chains/imageSearchAgent.ts +++ b/src/lib/chains/imageSearchAgent.ts @@ -7,7 +7,7 @@ import { PromptTemplate } from '@langchain/core/prompts'; import formatChatHistoryAsString from '../utils/formatHistory'; import { BaseMessage } from '@langchain/core/messages'; import { StringOutputParser } from '@langchain/core/output_parsers'; -import { searchSearxng } from '../searxng'; +import { searchSearxng } from '../searchEngines/searxng'; import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; const imageSearchChainPrompt = ` diff --git a/src/lib/chains/videoSearchAgent.ts b/src/lib/chains/videoSearchAgent.ts index f7cb156..0b05d57 100644 --- a/src/lib/chains/videoSearchAgent.ts +++ b/src/lib/chains/videoSearchAgent.ts @@ -7,7 +7,7 @@ import { PromptTemplate } from '@langchain/core/prompts'; import formatChatHistoryAsString from '../utils/formatHistory'; import { BaseMessage } from '@langchain/core/messages'; import { StringOutputParser } from '@langchain/core/output_parsers'; -import { searchSearxng } from '../searxng'; +import { searchSearxng } from '../searchEngines/searxng'; import type { BaseChatModel } from '@langchain/core/language_models/chat_models'; const VideoSearchChainPrompt = ` diff --git a/src/lib/config.ts b/src/lib/config.ts index 2831214..54d5c75 100644 --- a/src/lib/config.ts +++ b/src/lib/config.ts @@ -36,6 +36,10 @@ interface Config { }; API_ENDPOINTS: { SEARXNG: string; + TAVILY: string; + }; + SEARCH: { + ENGINE: string; }; } @@ -64,6 +68,12 @@ export const getGeminiApiKey = () => loadConfig().MODELS.GEMINI.API_KEY; export const getSearxngApiEndpoint = () => process.env.SEARXNG_API_URL || loadConfig().API_ENDPOINTS.SEARXNG; +export const getTavilyApiKey = () => + process.env.TAVILY_API_KEY || loadConfig().API_ENDPOINTS.TAVILY; + +export const getSearchEngine = () => + process.env.SEARCH_ENGINE || loadConfig().SEARCH?.ENGINE || 'searxng'; + export const getOllamaApiEndpoint = () => loadConfig().MODELS.OLLAMA.API_URL; export const getDeepseekApiKey = () => loadConfig().MODELS.DEEPSEEK.API_KEY; diff --git a/src/lib/search/metaSearchAgent.ts b/src/lib/search/metaSearchAgent.ts index 67b7c58..94e1692 100644 --- a/src/lib/search/metaSearchAgent.ts +++ b/src/lib/search/metaSearchAgent.ts @@ -17,7 +17,9 @@ import LineListOutputParser from '../outputParsers/listLineOutputParser'; import LineOutputParser from '../outputParsers/lineOutputParser'; import { getDocumentsFromLinks } from '../utils/documents'; import { Document } from 'langchain/document'; -import { searchSearxng } from '../searxng'; +import { searchTavily } from '../searchEngines/tavily'; +import { searchSearxng } from '../searchEngines/searxng'; +import { getSearchEngine } from '../config'; import path from 'node:path'; import fs from 'node:fs'; import computeSimilarity from '../utils/computeSimilarity'; @@ -205,25 +207,42 @@ class MetaSearchAgent implements MetaSearchAgentType { } else { question = question.replace(/.*?<\/think>/g, ''); - const res = await searchSearxng(question, { - language: 'en', - engines: this.config.activeEngines, - }); + const searchEngine = getSearchEngine(); - const documents = res.results.map( - (result) => - new Document({ - pageContent: - result.content || - (this.config.activeEngines.includes('youtube') - ? result.title - : '') /* Todo: Implement transcript grabbing using Youtubei (source: https://www.npmjs.com/package/youtubei) */, - metadata: { - title: result.title, - url: result.url, - ...(result.img_src && { img_src: result.img_src }), - }, - }), + let res; + + if (searchEngine === 'tavily') { + res = await searchTavily(question, { + search_depth: 'basic', + max_results: 15, + include_images: true, + }); + } else { + // Default to SearxNG + res = await searchSearxng(question, { + language: 'en', + engines: this.config.activeEngines, + }); + } + + let documents: Document[] = []; + + documents = documents.concat( + res.results.map( + (result) => + new Document({ + pageContent: + result.content || + (this.config.activeEngines.includes('youtube') + ? result.title + : ''), + metadata: { + title: result.title, + url: result.url, + ...(result.img_src ? { img_src: result.img_src } : {}), + }, + }), + ) ); return { query: question, docs: documents }; diff --git a/src/lib/searxng.ts b/src/lib/searchEngines/searxng.ts similarity index 95% rename from src/lib/searxng.ts rename to src/lib/searchEngines/searxng.ts index ae19db2..1209fea 100644 --- a/src/lib/searxng.ts +++ b/src/lib/searchEngines/searxng.ts @@ -1,5 +1,5 @@ import axios from 'axios'; -import { getSearxngApiEndpoint } from './config'; +import { getSearxngApiEndpoint } from '../config'; interface SearxngSearchOptions { categories?: string[]; diff --git a/src/lib/searchEngines/tavily.ts b/src/lib/searchEngines/tavily.ts new file mode 100644 index 0000000..b6e651d --- /dev/null +++ b/src/lib/searchEngines/tavily.ts @@ -0,0 +1,79 @@ +import axios from 'axios'; +import { getTavilyApiKey } from '../config'; + +interface TavilySearchOptions { + topic?: 'general' | 'news'; + search_depth?: 'basic' | 'advanced'; + chunks_per_source?: number; + max_results?: number; + time_range?: 'day' | 'week' | 'month' | 'year' | 'd' | 'w' | 'm' | 'y'; + days?: number; + include_answer?: boolean | 'basic' | 'advanced'; + include_raw_content?: boolean; + include_images?: boolean; + include_image_descriptions?: boolean; + include_domains?: string[]; + exclude_domains?: string[]; +} + +interface TavilySearchResult { + title: string; + url: string; + content: string; + score: number; + raw_content?: string; +} + +interface TavilySearchResponse { + query: string; + answer?: string; + images?: Array<{ + url: string; + description?: string; + }>; + results: TavilySearchResult[]; + response_time: string; +} + +export const searchTavily = async ( + query: string, + opts?: TavilySearchOptions, +) => { + const tavilyApiKey = getTavilyApiKey(); + + if (!tavilyApiKey) { + throw new Error('Tavily API key is not configured'); + } + + const url = 'https://api.tavily.com/search'; + + const response = await axios.post( + url, + { + query, + ...opts, + }, + { + headers: { + 'Content-Type': 'application/json', + 'Authorization': `Bearer ${tavilyApiKey}`, + }, + } + ); + + const results = response.data.results; + + // Convert Tavily results to match the format expected by the rest of the application + const formattedResults = results.map(result => ({ + title: result.title, + url: result.url, + content: result.content, + img_src: undefined, // Tavily doesn't provide image URLs in the standard response + })); + + return { + results: formattedResults, + suggestions: [], // Tavily doesn't provide suggestions, so return empty array + answer: response.data.answer, // Include the AI-generated answer if available + }; +}; \ No newline at end of file