diff --git a/deploy/pi-agent/extensions/tavily-search.mjs b/deploy/pi-agent/extensions/tavily-search.mjs new file mode 100644 index 00000000..7b0834ae --- /dev/null +++ b/deploy/pi-agent/extensions/tavily-search.mjs @@ -0,0 +1,300 @@ +/** + * Tavily Web Search Extension for Pi Coding Agent + * + * Features: + * - Real-time web search integration via Tavily Search API + * - API Key Pool support (via TAVILY_API_KEYS or TAVILY_API_KEY) + * - Automatic key rotation and failover on rate-limit (429), quota exhaustion, or auth errors (401/403) + * - Returns structured Markdown containing direct AI answers and top search results + */ + +export class TavilyKeyPool { + /** + * @param {string[]} keys + */ + constructor(keys = []) { + this.keys = keys.map(k => (typeof k === "string" ? k.trim() : "")).filter(k => k.length > 0); + this.currentIndex = 0; + this.stats = new Map(); // key -> { failures: number, lastError: string, lastTime: number } + } + + /** + * Initialize pool from environment variables + * Checks TAVILY_API_KEYS (comma/semicolon/newline separated) and TAVILY_API_KEY + * @param {Record} [env] + * @returns {TavilyKeyPool} + */ + static fromEnv(env = process.env) { + const raw = (env.TAVILY_API_KEYS || env.TAVILY_API_KEY || "").trim(); + if (!raw) { + return new TavilyKeyPool([]); + } + const splitKeys = raw + .split(/[\n,;\s]+/) + .map(k => k.trim()) + .filter(k => k.length > 0); + return new TavilyKeyPool(splitKeys); + } + + size() { + return this.keys.length; + } + + /** + * @param {string} key + * @returns {string} + */ + maskKey(key) { + if (!key || key.length < 8) return "***"; + return key.slice(0, 4) + "..." + key.slice(-4); + } + + /** + * @returns {string|null} + */ + getCurrentKey() { + if (this.keys.length === 0) return null; + return this.keys[this.currentIndex % this.keys.length]; + } + + /** + * Mark current key as failed and advance to next key + * @param {string} key + * @param {string} reason + * @param {number} [status] + */ + markFailure(key, reason, status) { + const current = this.stats.get(key) || { failures: 0, lastError: "", lastTime: 0 }; + current.failures += 1; + current.lastError = `[Status ${status || "ERR"}] ${reason}`; + current.lastTime = Date.now(); + this.stats.set(key, current); + + const oldIndex = this.currentIndex; + this.currentIndex = (this.currentIndex + 1) % this.keys.length; + console.warn( + `[TavilyKeyPool] Key ${this.maskKey(key)} failed (attempt #${current.failures}): ${reason}. ` + + `Rotated from #${oldIndex + 1} to #${this.currentIndex + 1}/${this.keys.length}` + ); + } + + /** + * Execute an async action with automatic retry and key rotation across the pool + * @template T + * @param {(apiKey: string) => Promise} action + * @returns {Promise} + */ + async executeWithPool(action) { + if (this.keys.length === 0) { + throw new Error( + "Tavily API key is not configured. Please set TAVILY_API_KEYS in .env (supports comma-separated keys)." + ); + } + + const total = this.keys.length; + const errors = []; + let attempts = 0; + + while (attempts < total) { + const activeKey = this.getCurrentKey(); + try { + const result = await action(activeKey); + return result; + } catch (err) { + attempts++; + const status = err.status || 0; + const msg = err.message || String(err); + this.markFailure(activeKey, msg, status); + errors.push(`[Key ${this.maskKey(activeKey)}] ${msg}`); + } + } + + throw new Error( + `All ${total} Tavily API key(s) in pool failed:\n- ` + errors.join("\n- ") + ); + } +} + +/** + * Execute HTTP POST query to Tavily Search API + * @param {string} apiKey + * @param {string} query + * @param {object} [options] + * @returns {Promise} + */ +export async function searchTavily(apiKey, query, options = {}) { + const endpoint = options.endpoint || "https://api.tavily.com/search"; + const searchDepth = options.search_depth === "advanced" ? "advanced" : "basic"; + const maxResults = Math.min(Math.max(Number(options.max_results) || 5, 1), 10); + + const payload = { + api_key: apiKey, + query: query, + search_depth: searchDepth, + include_answer: true, + max_results: maxResults, + }; + + const response = await fetch(endpoint, { + method: "POST", + headers: { + "Content-Type": "application/json", + "User-Agent": "pi-agent-tavily-tool/1.0", + }, + body: JSON.stringify(payload), + signal: AbortSignal.timeout(options.timeout || 25000), + }); + + if (!response.ok) { + let errDetail = ""; + try { + const data = await response.json(); + errDetail = data.detail || JSON.stringify(data); + } catch { + errDetail = await response.text(); + } + const err = new Error(`HTTP ${response.status} ${response.statusText}: ${errDetail}`); + err.status = response.status; + throw err; + } + + return await response.json(); +} + +/** + * Format Tavily search API response into clean Markdown for the Agent + * @param {any} data + * @param {string} query + * @returns {string} + */ +export function formatTavilyResults(data, query) { + let md = `## 🌐 Tavily 实时网络搜索结果: "${query}"\n\n`; + + if (data.answer) { + md += `> **智能总结**: ${data.answer.trim()}\n\n`; + } + + const results = Array.isArray(data.results) ? data.results : []; + if (results.length > 0) { + md += `### 检索到的参考条目 (${results.length} 条):\n\n`; + results.forEach((item, idx) => { + const title = item.title ? item.title.trim() : "无标题网页"; + const url = item.url ? item.url.trim() : "#"; + const snippet = item.content ? item.content.trim().replace(/\s+/g, " ") : "(无摘要)"; + md += `${idx + 1}. **[${title}](${url})**\n`; + md += ` - **摘要**: ${snippet}\n`; + if (item.score) { + md += ` - **相关度评分**: ${item.score.toFixed ? item.score.toFixed(3) : item.score}\n`; + } + md += "\n"; + }); + } else { + md += `*未检索到与 "${query}" 相关的网页结果。*\n\n`; + } + + return md; +} + +/** + * Pi Extension Entrypoint + * @param {object} pi - Pi Extension API + */ +export default function (pi) { + const pool = TavilyKeyPool.fromEnv(); + + // Register tavily_search tool + pi.registerTool({ + name: "tavily_search", + label: "Tavily Web Search", + description: + "Search the live internet for real-time information, competitor intelligence, regulatory news, product specs, or market data via Tavily (with automatic API key pool rotation).", + promptSnippet: "Search the web for up-to-date market, regulatory, or technical information with tavily_search", + promptGuidelines: [ + "Use tavily_search when you need external facts, recent regulations, competitor solutions, or live web data not found in local workspace or knowledge base.", + "Provide specific, high-signal query keywords in Chinese or English." + ], + parameters: { + type: "object", + properties: { + query: { + type: "string", + description: "The search query keywords in Chinese or English", + }, + search_depth: { + type: "string", + enum: ["basic", "advanced"], + description: "Search depth: 'basic' (fast summary) or 'advanced' (deep analysis, default 'basic')", + }, + max_results: { + type: "number", + description: "Maximum number of search results to retrieve (1 to 10, default 5)", + }, + }, + required: ["query"], + }, + async execute(toolCallId, params, signal, onUpdate, ctx) { + const query = (params.query || "").trim(); + if (!query) { + return { + content: [{ type: "text", text: "Error: query cannot be empty." }], + details: { error: "empty_query" }, + }; + } + + try { + const data = await pool.executeWithPool(async (apiKey) => { + return await searchTavily(apiKey, query, { + search_depth: params.search_depth, + max_results: params.max_results, + }); + }); + + const formatted = formatTavilyResults(data, query); + return { + content: [{ type: "text", text: formatted }], + details: { + query, + results_count: Array.isArray(data.results) ? data.results.length : 0, + has_answer: Boolean(data.answer), + }, + }; + } catch (err) { + return { + content: [ + { + type: "text", + text: `[Tavily Search Failed]: ${err.message}\n(Agent should rely on workspace materials or local knowledge base as fallback.)`, + }, + ], + details: { error: err.message }, + }; + } + }, + }); + + // Also register alias 'web_search' pointing to the same tool for convenience + pi.registerTool({ + name: "web_search", + label: "Web Search", + description: "Alias for tavily_search. Searches the live internet for information.", + parameters: { + type: "object", + properties: { + query: { + type: "string", + description: "The search query keywords", + }, + }, + required: ["query"], + }, + async execute(toolCallId, params, signal, onUpdate, ctx) { + const tool = pi.getAllTools ? pi.getAllTools().find(t => t.name === "tavily_search") : null; + if (tool && tool.execute) { + return await tool.execute(toolCallId, params, signal, onUpdate, ctx); + } + return { + content: [{ type: "text", text: "tavily_search is active. Please use tavily_search." }], + }; + }, + }); +} diff --git a/tools/pi-extensions/tavily-search.mjs b/tools/pi-extensions/tavily-search.mjs new file mode 100644 index 00000000..05f788c1 --- /dev/null +++ b/tools/pi-extensions/tavily-search.mjs @@ -0,0 +1,2 @@ +export * from "../../deploy/pi-agent/extensions/tavily-search.mjs"; +export { default } from "../../deploy/pi-agent/extensions/tavily-search.mjs";