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

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