feat(agent): add Tavily search extension with API key pool and auto-rotation

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-13 02:50:18 +00:00
parent 1208c74030
commit b78c9f5071
2 changed files with 302 additions and 0 deletions
@@ -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<string, string|undefined>} [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<T>} action
+   * @returns {Promise<T>}
+   */
+  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<any>}
+ */
+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." }],
+      };
+    },
+  });
+}
@@ -0,0 +1,2 @@
+export * from "../../deploy/pi-agent/extensions/tavily-search.mjs";
+export { default } from "../../deploy/pi-agent/extensions/tavily-search.mjs";