diff options
Diffstat (limited to 'config/pi/extensions/parallel-search.ts')
| -rw-r--r-- | config/pi/extensions/parallel-search.ts | 401 |
1 files changed, 401 insertions, 0 deletions
diff --git a/config/pi/extensions/parallel-search.ts b/config/pi/extensions/parallel-search.ts new file mode 100644 index 0000000..3bca820 --- /dev/null +++ b/config/pi/extensions/parallel-search.ts @@ -0,0 +1,401 @@ +/** + * Parallel Search Extension - Web search via Parallel.ai + * + * Provides the search tool for searching the web using natural language objectives or keyword queries. + * + * Requires PARALLEL_API_KEY environment variable. + * + * https://docs.parallel.ai + */ + +import { + DEFAULT_MAX_BYTES, + DEFAULT_MAX_LINES, + type ExtensionAPI, + formatSize, + getMarkdownTheme, + keyHint, + type Theme, + truncateHead, +} from "@mariozechner/pi-coding-agent"; +import { Container, Markdown, Spacer, Text } from "@mariozechner/pi-tui"; +import { Type } from "typebox"; + +const API_BASE = "https://api.parallel.ai"; +const BETA_HEADER = "search-extract-2025-10-10"; + +function getApiKey(): string { + const key = process.env.PARALLEL_API_KEY; + if (!key) { + throw new Error( + "PARALLEL_API_KEY environment variable is not set. " + + "Get your API key at https://platform.parallel.ai", + ); + } + return key; +} + +// ── API types ────────────────────────────────────────────────────────────── + +interface SearchResult { + url: string; + title?: string | null; + publish_date?: string | null; + excerpts?: string[] | null; +} + +interface SearchResponse { + search_id: string; + results: SearchResult[]; +} + +// ── API helper ───────────────────────────────────────────────────────────── + +async function parallelSearch( + params: { + objective?: string; + search_queries?: string[]; + max_results?: number; + mode?: string; + include_domains?: string[]; + exclude_domains?: string[]; + after_date?: string; + max_chars_per_result?: number; + max_chars_total?: number; + }, + signal?: AbortSignal, +): Promise<SearchResponse> { + const body: Record<string, unknown> = { + mode: params.mode ?? "agentic", + }; + if (params.objective) body.objective = params.objective; + if (params.search_queries) body.search_queries = params.search_queries; + if (params.max_results) body.max_results = params.max_results; + + if (params.max_chars_per_result || params.max_chars_total) { + const excerpts: Record<string, unknown> = {}; + if (params.max_chars_per_result) excerpts.max_chars_per_result = params.max_chars_per_result; + if (params.max_chars_total) excerpts.max_chars_total = params.max_chars_total; + body.excerpts = excerpts; + } + + if (params.include_domains || params.exclude_domains || params.after_date) { + const sourcePolicy: Record<string, unknown> = {}; + if (params.include_domains) sourcePolicy.include_domains = params.include_domains; + if (params.exclude_domains) sourcePolicy.exclude_domains = params.exclude_domains; + if (params.after_date) sourcePolicy.after_date = params.after_date; + body.source_policy = sourcePolicy; + } + + const resp = await fetch(`${API_BASE}/v1beta/search`, { + method: "POST", + headers: { + "Content-Type": "application/json", + "x-api-key": getApiKey(), + "parallel-beta": BETA_HEADER, + }, + body: JSON.stringify(body), + signal, + }); + + if (!resp.ok) { + const text = await resp.text().catch(() => ""); + throw new Error(`Parallel search API error ${resp.status}: ${text}`); + } + + return (await resp.json()) as SearchResponse; +} + +// ── Formatting helper ────────────────────────────────────────────────────── + +function formatSearchResults(results: SearchResult[]): string { + if (results.length === 0) return "No results found."; + + const parts: string[] = []; + for (let i = 0; i < results.length; i++) { + const r = results[i]; + let entry = `[${i + 1}] ${r.title ?? "(no title)"}\n ${r.url}`; + if (r.publish_date) entry += `\n Published: ${r.publish_date}`; + if (r.excerpts?.length) { + entry += "\n" + r.excerpts.map((e) => ` ${e}`).join("\n"); + } + parts.push(entry); + } + return parts.join("\n\n"); +} + +// ── Tool details types ───────────────────────────────────────────────────── + +interface SearchResultDetail { + url: string; + title?: string | null; + publish_date?: string | null; + excerpts?: string[] | null; +} + +interface SearchDetails { + objective?: string; + search_queries?: string[]; + resultCount: number; + searchId: string; + results?: SearchResultDetail[]; + spinnerFrame?: number; +} + +// ── Spinner ───────────────────────────────────────────────────────────────── + +const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; + +type SpinnerUpdate = (partial: { + content: { type: "text"; text: string }[]; + details?: SearchDetails; +}) => void; + +function startSpinner( + onUpdate: SpinnerUpdate | undefined, + baseDetails: SearchDetails, + signal?: AbortSignal, +): () => void { + if (!onUpdate) return () => {}; + let frame = 0; + const interval = setInterval(() => { + if (signal?.aborted) { + clearInterval(interval); + return; + } + onUpdate({ + content: [{ type: "text", text: "Running..." }], + details: { ...baseDetails, spinnerFrame: frame }, + }); + frame = (frame + 1) % SPINNER_FRAMES.length; + }, 80); + return () => clearInterval(interval); +} + +function renderSpinner(details: SearchDetails | undefined, theme: Theme, label: string): Text { + const frame = details?.spinnerFrame ?? 0; + const spinner = SPINNER_FRAMES[frame % SPINNER_FRAMES.length]; + return new Text(`\n${theme.fg("dim", `${spinner} ${label} (escape/ctrl+c to cancel)`)}`, 0, 0); +} + +// ── Extension entry point ────────────────────────────────────────────────── + +export default function (pi: ExtensionAPI) { + if (!process.env.PARALLEL_API_KEY) return; + + const searchCountCache = new Map<string, number>(); + function searchCacheKey(objective?: string, search_queries?: string[]): string { + return JSON.stringify({ objective, search_queries }); + } + + pi.registerTool({ + name: "search", + label: "Search", + description: + `Search the web. Provide a natural language objective describing what you're looking for, ` + + `and/or keyword search queries. Returns relevant results with titles, URLs, and excerpts. ` + + `Output truncated to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)}.`, + promptSnippet: "Search the web via Parallel.ai (objective and/or keyword queries)", + parameters: Type.Object({ + objective: Type.Optional( + Type.String({ + description: + "Natural language description of what to find. At least one of objective or search_queries is required.", + }), + ), + search_queries: Type.Optional( + Type.Array(Type.String(), { + description: + "Keyword search queries. May contain search operators. At least one of objective or search_queries is required.", + }), + ), + max_results: Type.Optional( + Type.Number({ + description: "Max results to return (default: 5)", + }), + ), + mode: Type.Optional( + Type.String({ + description: + 'Search mode: "agentic" (default, concise), "one-shot" (comprehensive), or "fast" (lower latency)', + }), + ), + include_domains: Type.Optional( + Type.Array(Type.String(), { + description: 'Restrict results to these domains (e.g. ["wikipedia.org", ".gov"])', + }), + ), + exclude_domains: Type.Optional( + Type.Array(Type.String(), { + description: 'Exclude results from these domains (e.g. ["reddit.com"])', + }), + ), + after_date: Type.Optional( + Type.String({ + description: "Only return results published on or after this date (YYYY-MM-DD)", + }), + ), + max_chars_per_result: Type.Optional( + Type.Number({ + description: "Max characters of excerpts per result (min 1000)", + }), + ), + max_chars_total: Type.Optional( + Type.Number({ + description: "Max total characters of excerpts across all results (min 1000)", + }), + ), + }), + + async execute(_toolCallId, params, signal, onUpdate) { + const { + objective, + search_queries, + max_results, + mode, + include_domains, + exclude_domains, + after_date, + max_chars_per_result, + max_chars_total, + } = params; + + if (!objective && (!search_queries || search_queries.length === 0)) { + throw new Error("at least one of objective or search_queries must be provided."); + } + + const baseDetails: SearchDetails = { + objective, + search_queries, + resultCount: 0, + searchId: "", + }; + const stopSpinner = startSpinner(onUpdate as SpinnerUpdate | undefined, baseDetails, signal); + + let response: SearchResponse; + try { + response = await parallelSearch( + { + objective, + search_queries, + max_results: max_results ?? 5, + mode, + include_domains, + exclude_domains, + after_date, + max_chars_per_result: max_chars_per_result ?? 500, + max_chars_total: max_chars_total ?? 2500, + }, + signal, + ); + } finally { + stopSpinner(); + } + + let resultText = formatSearchResults(response.results); + + const truncation = truncateHead(resultText, { + maxLines: DEFAULT_MAX_LINES, + maxBytes: DEFAULT_MAX_BYTES, + }); + + resultText = truncation.content; + if (truncation.truncated) { + resultText += `\n\n[Output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines `; + resultText += `(${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)})]`; + } + + searchCountCache.set(searchCacheKey(objective, search_queries), response.results.length); + + const details: SearchDetails = { + objective, + search_queries, + resultCount: response.results.length, + searchId: response.search_id, + results: response.results.map((r) => ({ + url: r.url, + title: r.title, + publish_date: r.publish_date, + excerpts: r.excerpts, + })), + }; + + return { + content: [{ type: "text", text: resultText }], + details, + }; + }, + + renderCall(args, theme) { + let text = theme.fg("toolTitle", theme.bold("search ")); + if (args.objective) { + text += theme.fg("accent", `${args.objective}`); + } else if (args.search_queries?.length) { + text += theme.fg("accent", args.search_queries.join(", ")); + } + if (args.include_domains?.length) { + text += theme.fg("muted", ` site:${args.include_domains.join(",")}`); + } + const count = searchCountCache.get(searchCacheKey(args.objective, args.search_queries)); + if (count !== undefined) { + text += theme.fg("dim", ` [${count} results]`); + } + return new Text(text, 0, 0); + }, + + renderResult(result, { expanded, isPartial }, theme) { + const details = result.details as SearchDetails | undefined; + + if (isPartial) { + return renderSpinner(details, theme, "Searching..."); + } + + if (!details || !details.resultCount) { + const content = result.content[0]; + const msg = content?.type === "text" ? content.text : "No results found"; + return new Text(theme.fg("warning", msg), 0, 0); + } + + if (!expanded) { + let text = "\n"; + if (details.results?.length) { + for (let i = 0; i < details.results.length; i++) { + const r = details.results[i]; + if (i > 0) text += "\n"; + const url = r.url.replace(/^https?:\/\//, ""); + text += `${theme.fg("dim", `${i + 1}.`)} ${theme.fg("accent", url)}`; + } + } + text += `\n\n(${keyHint("app.tools.expand", "to expand")})`; + return new Text(text, 0, 0); + } + + const container = new Container(); + container.addChild(new Spacer(1)); + + if (details.results?.length) { + const mdTheme = getMarkdownTheme(); + for (let i = 0; i < details.results.length; i++) { + if (i > 0) { + container.addChild(new Spacer(1)); + } + const r = details.results[i]; + const url = r.url.replace(/^https?:\/\//, ""); + const header = `${theme.fg("dim", `${i + 1}.`)} ${theme.fg("accent", url)}`; + container.addChild(new Text(header, 0, 0)); + + if (r.excerpts?.length) { + const excerptMd = r.excerpts.join("\n\n---\n\n"); + container.addChild( + new Markdown(excerptMd, 3, 0, mdTheme, { color: (t) => theme.fg("dim", t) }), + ); + } + } + } + + container.addChild(new Spacer(1)); + container.addChild(new Text(`(${keyHint("app.tools.expand", "to collapse")})`, 0, 0)); + return container; + }, + }); +} |