diff options
Diffstat (limited to 'config/pi/extensions/parallel-extract.ts')
| -rw-r--r-- | config/pi/extensions/parallel-extract.ts | 277 |
1 files changed, 277 insertions, 0 deletions
diff --git a/config/pi/extensions/parallel-extract.ts b/config/pi/extensions/parallel-extract.ts new file mode 100644 index 0000000..2ab7a87 --- /dev/null +++ b/config/pi/extensions/parallel-extract.ts @@ -0,0 +1,277 @@ +/** + * Parallel Extract Extension - Content extraction via Parallel.ai + * + * Provides the extract tool for extracting content from URLs. + * + * 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 ExtractResult { + url: string; + title?: string | null; + publish_date?: string | null; + excerpts?: string[] | null; + full_content?: string | null; +} + +interface ExtractError { + url: string; + error_type: string; + http_status_code?: number | null; + content?: string | null; +} + +interface ExtractResponse { + extract_id: string; + results: ExtractResult[]; + errors: ExtractError[]; +} + +// ── API helper ───────────────────────────────────────────────────────────── + +async function parallelExtract( + params: { + url: string; + objective?: string; + }, + signal?: AbortSignal, +): Promise<ExtractResponse> { + const body: Record<string, unknown> = { + urls: [params.url], + excerpts: false, + full_content: true, + }; + if (params.objective) body.objective = params.objective; + + const resp = await fetch(`${API_BASE}/v1beta/extract`, { + 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(`Extract API error ${resp.status}: ${text}`); + } + + return (await resp.json()) as ExtractResponse; +} + +// ── Formatting helper ────────────────────────────────────────────────────── + +function formatExtractResult(result?: ExtractResult, error?: ExtractError): string { + if (error) { + return ( + `## Error: ${error.url}\nType: ${error.error_type}` + + (error.http_status_code ? ` (HTTP ${error.http_status_code})` : "") + + (error.content ? `\n${error.content}` : "") + ); + } + + if (!result) return "No content extracted."; + + return "\n" + (result.full_content ?? ""); +} + +// ── Tool details type ────────────────────────────────────────────────────── + +interface ExtractDetails { + url: string; + error: boolean; + extractId: string; + truncated?: boolean; + spinnerFrame?: number; +} + +// ── Spinner ───────────────────────────────────────────────────────────────── + +const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; + +type SpinnerUpdate = (partial: { + content: { type: "text"; text: string }[]; + details?: ExtractDetails; +}) => void; + +function startSpinner( + onUpdate: SpinnerUpdate | undefined, + baseDetails: ExtractDetails, + 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: ExtractDetails | 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; + + pi.registerTool({ + name: "extract", + label: "Extract", + description: + `Extract content from a URL. Returns page title and full content. ` + + `Optionally provide an objective to focus extraction on specific information. ` + + `Output truncated to ${DEFAULT_MAX_LINES} lines or ${formatSize(DEFAULT_MAX_BYTES)}.`, + promptSnippet: "Extract full content from a URL via Parallel.ai", + parameters: Type.Object({ + url: Type.String({ + description: "URL to extract content from", + }), + objective: Type.Optional( + Type.String({ + description: "Focus extraction on this objective (e.g. 'find the pricing information')", + }), + ), + }), + + async execute(_toolCallId, params, signal, onUpdate) { + const { url, objective } = params; + + const baseDetails: ExtractDetails = { url, error: false, extractId: "" }; + const stopSpinner = startSpinner(onUpdate as SpinnerUpdate | undefined, baseDetails, signal); + + let response: ExtractResponse; + try { + response = await parallelExtract({ url, objective }, signal); + } finally { + stopSpinner(); + } + + const hasError = response.errors.length > 0; + let resultText = formatExtractResult(response.results[0], response.errors[0]); + + const truncation = truncateHead(resultText, { + maxLines: DEFAULT_MAX_LINES, + maxBytes: DEFAULT_MAX_BYTES, + }); + + resultText = truncation.content; + const truncated = truncation.truncated; + if (truncated) { + resultText += `\n\n[Output truncated: showing ${truncation.outputLines} of ${truncation.totalLines} lines `; + resultText += `(${formatSize(truncation.outputBytes)} of ${formatSize(truncation.totalBytes)}). `; + resultText += `Try a more specific objective to reduce content size.]`; + } + + const details: ExtractDetails = { + url, + error: hasError, + extractId: response.extract_id, + truncated, + }; + + return { + content: [{ type: "text", text: resultText }], + details, + }; + }, + + renderCall(args, theme) { + let text = theme.fg("toolTitle", theme.bold("extract ")); + text += theme.fg("accent", args.url.replace(/^https?:\/\//, "")); + return new Text(text, 0, 0); + }, + + renderResult(result, { expanded, isPartial }, theme) { + const details = result.details as ExtractDetails | undefined; + + if (isPartial) { + return renderSpinner(details, theme, "Extracting..."); + } + + if (!details || !details.extractId) { + const content = result.content[0]; + const msg = content?.type === "text" ? content.text : "No content"; + return new Text(theme.fg("dim", msg), 0, 0); + } + + if (!expanded) { + let text = ""; + if (details.error) { + text += "\n" + theme.fg("error", "error"); + } + text += `\n(${keyHint("app.tools.expand", "to expand")})`; + return new Text(text, 0, 0); + } + + const content = result.content[0]; + if (content?.type === "text" && content.text) { + const container = new Container(); + if (details.error) { + container.addChild(new Text(theme.fg("error", "error"), 0, 0)); + } + const mdTheme = getMarkdownTheme(); + container.addChild( + new Markdown(content.text, 1, 0, mdTheme, { color: (t) => theme.fg("dim", t) }), + ); + if (details.truncated) { + container.addChild(new Text(theme.fg("warning", "[truncated]"), 1, 0)); + } + container.addChild(new Spacer(1)); + container.addChild(new Text(`(${keyHint("app.tools.expand", "to collapse")})`, 0, 0)); + return container; + } + + let text = ""; + if (details.error) { + text += theme.fg("error", "error"); + } + return new Text(text, 0, 0); + }, + }); +} |