UNPKG

mcp-tavily-search

Version:

MCP server for integrating Tavily search API with LLMs, providing web search, RAG context generation, and QnA capabilities

362 lines (361 loc) 13.9 kB
#!/usr/bin/env node import { Server } from '@modelcontextprotocol/sdk/server/index.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { CallToolRequestSchema, ErrorCode, ListToolsRequestSchema, McpError, } from '@modelcontextprotocol/sdk/types.js'; import { readFileSync } from 'fs'; import { dirname, join } from 'path'; import { fileURLToPath } from 'url'; import { TAVILY_TOOLS } from './tools.js'; const __filename = fileURLToPath(import.meta.url); const __dirname = dirname(__filename); const pkg = JSON.parse(readFileSync(join(__dirname, '..', 'package.json'), 'utf8')); const { name, version } = pkg; const TAVILY_API_KEY = process.env.TAVILY_API_KEY; if (!TAVILY_API_KEY) { throw new Error('TAVILY_API_KEY environment variable is required'); } // Error Handling class TavilyError extends Error { constructor(message, code, details) { super(message); this.code = code; this.details = details; this.name = 'TavilyError'; } } const error_handler = { format_error(error) { if (error instanceof TavilyError) { return `Tavily API Error (${error.code}): ${error.message}`; } return `Error: ${error instanceof Error ? error.message : String(error)}`; }, get_error_code(error) { if (error instanceof TavilyError) { return error.code; } return 'UNKNOWN_ERROR'; }, get_error_details(error) { if (error instanceof TavilyError) { return error.details; } return null; }, }; // Response Formatters const response_formatters = { text(data) { let output = `Search Results for "${data.query}":\n\n`; if (data.answer) { output += `Summary: ${data.answer}\n\nDetailed Sources:\n`; } return (output + data.results .map((result, i) => { let source = `${i + 1}. ${result.title}\n`; source += ` URL: ${result.url}\n`; if (result.published_date) { source += ` Published: ${result.published_date}\n`; } source += ` Content: ${result.content}\n`; return source; }) .join('\n')); }, json(data) { return JSON.stringify(data, null, 2); }, markdown(data) { let output = `# Search Results: ${data.query}\n\n`; if (data.answer) { output += `## Summary\n${data.answer}\n\n## Sources\n`; } return (output + data.results .map((result) => { let source = `### ${result.title}\n`; source += `- URL: [${result.url}](${result.url})\n`; if (result.published_date) { source += `- Published: ${result.published_date}\n`; } source += `\n${result.content}\n`; return source; }) .join('\n---\n')); }, }; class TavilyCache { constructor() { this.cache = new Map(); } set(key, data, ttl) { this.cache.set(key, { data, timestamp: Date.now(), ttl, }); } get(key) { const entry = this.cache.get(key); if (!entry) return null; if (Date.now() - entry.timestamp > entry.ttl * 1000) { this.cache.delete(key); return null; } return entry.data; } clear() { this.cache.clear(); } } class TavilySearchServer { constructor() { // No default domains - let users specify their trusted/excluded sources this.default_include_domains = []; this.default_exclude_domains = []; this.server = new Server({ name, version }, { capabilities: { tools: { tavily_search: true, tavily_get_search_context: true, tavily_qna_search: true, }, caching: true, formatting: ['text', 'json', 'markdown'], }, }); this.cache = new TavilyCache(); this.setup_tool_handlers(); } setup_tool_handlers() { this.server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: TAVILY_TOOLS, })); this.server.setRequestHandler(CallToolRequestSchema, async (request) => { const tool = TAVILY_TOOLS.find((t) => t.name === request.params.name); if (!tool) { throw new McpError(ErrorCode.MethodNotFound, `Unknown tool: ${request.params.name}`); } try { switch (request.params.name) { case 'tavily_search': { const rawArgs = request.params.arguments; if (!rawArgs?.query || typeof rawArgs.query !== 'string') { throw new McpError(ErrorCode.InvalidParams, 'Query parameter is required and must be a string'); } const args = { query: rawArgs.query, search_depth: rawArgs.search_depth, topic: rawArgs.topic, days: rawArgs.days, time_range: rawArgs.time_range, max_results: rawArgs.max_results, include_images: rawArgs.include_images, include_image_descriptions: rawArgs.include_image_descriptions, include_answer: rawArgs.include_answer, include_raw_content: rawArgs.include_raw_content, include_domains: rawArgs.include_domains, exclude_domains: rawArgs.exclude_domains, }; return await this.handle_search(args); } case 'tavily_get_search_context': { const rawArgs = request.params.arguments; if (!rawArgs?.query || typeof rawArgs.query !== 'string') { throw new McpError(ErrorCode.InvalidParams, 'Query parameter is required and must be a string'); } const args = { query: rawArgs.query, max_tokens: rawArgs.max_tokens, search_depth: rawArgs.search_depth, topic: rawArgs.topic, days: rawArgs.days, time_range: rawArgs.time_range, max_results: rawArgs.max_results, include_domains: rawArgs.include_domains, exclude_domains: rawArgs.exclude_domains, }; return await this.handle_context(args); } case 'tavily_qna_search': { const rawArgs = request.params.arguments; if (!rawArgs?.query || typeof rawArgs.query !== 'string') { throw new McpError(ErrorCode.InvalidParams, 'Query parameter is required and must be a string'); } const args = { query: rawArgs.query, search_depth: rawArgs.search_depth, topic: rawArgs.topic, days: rawArgs.days, time_range: rawArgs.time_range, max_results: rawArgs.max_results, include_domains: rawArgs.include_domains, exclude_domains: rawArgs.exclude_domains, }; return await this.handle_qna(args); } default: throw new McpError(ErrorCode.MethodNotFound, `Unimplemented tool: ${request.params.name}`); } } catch (error) { return { content: [ { type: 'text', text: error_handler.format_error(error), }, ], isError: true, }; } }); } async handle_search(args) { const { query, search_depth = 'basic', topic = 'general', days, time_range, max_results = 5, include_images = false, include_image_descriptions = false, include_answer = false, include_raw_content = false, include_domains = this.default_include_domains, exclude_domains = this.default_exclude_domains, } = args; // Check cache if enabled const cache_key = JSON.stringify({ query, search_depth, topic, include_answer, include_images, include_raw_content, }); const cached_data = this.cache.get(cache_key); if (cached_data) { return { content: [ { type: 'text', text: JSON.stringify(cached_data, null, 2), }, ], }; } const start_time = Date.now(); const response = await fetch('https://api.tavily.com/search', { method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${TAVILY_API_KEY}`, }, body: JSON.stringify({ query, search_depth, topic, days, time_range, max_results, include_images, include_image_descriptions, include_answer, include_raw_content, include_domains, exclude_domains, }), }); if (!response.ok) { throw new TavilyError(`API request failed: ${response.statusText}`, 'API_ERROR', { status: response.status }); } const data = await response.json(); const response_time = (Date.now() - start_time) / 1000; const search_response = { ...data, response_time, }; // Cache the results this.cache.set(cache_key, search_response, 3600); return { content: [ { type: 'text', text: JSON.stringify(search_response, null, 2), }, ], }; } async handle_context(args) { const { query, max_tokens = 2000, search_depth = 'advanced', topic = 'general', days, time_range, max_results = 5, include_domains = this.default_include_domains, exclude_domains = this.default_exclude_domains, } = args; const response = await fetch('https://api.tavily.com/search', { method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${TAVILY_API_KEY}`, }, body: JSON.stringify({ query, search_depth, topic, days, time_range, max_results, include_domains, exclude_domains, max_tokens, include_answer: false, }), }); if (!response.ok) { throw new TavilyError(`API request failed: ${response.statusText}`, 'API_ERROR', { status: response.status }); } const data = await response.json(); const context = data.results .map((result) => result.content) .join('\n\n') .slice(0, max_tokens); return { content: [ { type: 'text', text: context, }, ], }; } async handle_qna(args) { const { query, search_depth = 'advanced', topic = 'general', days, time_range, max_results = 5, include_domains = this.default_include_domains, exclude_domains = this.default_exclude_domains, } = args; const response = await fetch('https://api.tavily.com/search', { method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${TAVILY_API_KEY}`, }, body: JSON.stringify({ query, search_depth, topic, days, time_range, max_results, include_domains, exclude_domains, include_answer: true, }), }); if (!response.ok) { throw new TavilyError(`API request failed: ${response.statusText}`, 'API_ERROR', { status: response.status }); } const data = await response.json(); return { content: [ { type: 'text', text: data.answer || 'No answer found.', }, ], }; } async run() { const transport = new StdioServerTransport(); await this.server.connect(transport); console.error('Tavily Search MCP server running on stdio'); } } const server = new TavilySearchServer(); server.run().catch(console.error);