path: root/config/pi/extensions/parallel-extract.ts

/**
 * 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);
        },
    });
}