301 lines
9.2 KiB
JavaScript
301 lines
9.2 KiB
JavaScript
/**
|
|
* 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." }],
|
|
};
|
|
},
|
|
});
|
|
}
|