my-pi/pi-web-search/clients/tavily.ts

/**
 * Tavily REST client: POST https://api.tavily.com/search|extract with a Bearer
 * key from auth.json (`tavily` api_key credential) — there is no keyless tier.
 */

import { getSearchApiKey } from "../auth.ts";

export class TavilyError extends Error {}

const TAVILY_BASE_URL = "https://api.tavily.com";
const SEARCH_TIMEOUT_MS = 45_000;
const EXTRACT_TIMEOUT_MS = 60_000;
const CLIENT_SOURCE = "my-pi-web-search";

export type FetchLike = (input: string | URL, init?: RequestInit) => Promise<Response>;

export interface TavilySearchParams {
	query: string;
	max_results?: number;
	search_depth?: "basic" | "advanced";
	topic?: "general" | "news" | "finance";
	include_domains?: string[];
	start_date?: string;
	end_date?: string;
}

export interface TavilySearchResult {
	title?: string;
	url?: string;
	content?: string;
	published_date?: string;
	[key: string]: unknown;
}

export function tavilyAvailable(resolve: (provider: "tavily") => string | undefined = getSearchApiKey): boolean {
	return Boolean(resolve("tavily"));
}

function withTimeout(signal: AbortSignal | undefined, timeoutMs: number): AbortSignal {
	const timeout = AbortSignal.timeout(timeoutMs);
	return signal ? AbortSignal.any([signal, timeout]) : timeout;
}

async function post(
	endpoint: string,
	body: Record<string, unknown>,
	timeoutMs: number,
	signal: AbortSignal | undefined,
	fetchImpl: FetchLike,
): Promise<Record<string, unknown>> {
	const apiKey = getSearchApiKey("tavily");
	if (!apiKey) {
		throw new TavilyError('Tavily is not configured: no "tavily" api_key credential in auth.json.');
	}
	let response: Response;
	try {
		response = await fetchImpl(`${TAVILY_BASE_URL}/${endpoint}`, {
			method: "POST",
			headers: {
				"Content-Type": "application/json",
				Authorization: `Bearer ${apiKey}`,
				"X-Client-Source": CLIENT_SOURCE,
			},
			body: JSON.stringify(body),
			signal: withTimeout(signal, timeoutMs),
		});
	} catch (error) {
		throw new TavilyError(`Could not reach the Tavily API: ${(error as Error).message}`);
	}
	if (!response.ok) {
		let detail = "";
		try {
			const data = (await response.json()) as Record<string, unknown>;
			detail = String(data.error ?? data.message ?? "");
		} catch {
			detail = "";
		}
		throw new TavilyError(`Tavily API error (${response.status})${detail ? `: ${detail}` : ""}`);
	}
	return (await response.json()) as Record<string, unknown>;
}

/** POST /search — broad discovery; returns the answer plus a compact result list. */
export async function tavilySearch(
	params: TavilySearchParams,
	signal?: AbortSignal,
	fetchImpl: FetchLike = fetch,
): Promise<{ answer?: string; results: TavilySearchResult[] }> {
	const payload: Record<string, unknown> = {
		query: params.query,
		max_results: params.max_results,
		search_depth: params.search_depth,
		topic: params.topic,
		include_domains: params.include_domains,
		start_date: params.start_date,
		end_date: params.end_date,
	};
	for (const key of Object.keys(payload)) {
		if (payload[key] === undefined) delete payload[key];
	}
	const data = await post("search", payload, SEARCH_TIMEOUT_MS, signal, fetchImpl);
	const rawResults = Array.isArray(data.results) ? (data.results as Record<string, unknown>[]) : [];
	return {
		answer: typeof data.answer === "string" ? data.answer : undefined,
		results: rawResults.map((result) => ({
			title: result.title,
			url: result.url,
			content: result.content,
			published_date: result.publishedDate ?? result.published_date,
		})),
	};
}

/** POST /extract — full page content as markdown for a batch of URLs. */
export async function tavilyExtract(
	urls: string[],
	signal?: AbortSignal,
	fetchImpl: FetchLike = fetch,
): Promise<{ results: Array<{ url?: string; content?: string }>; failed: string[] }> {
	const data = await post("extract", { urls, format: "markdown" }, EXTRACT_TIMEOUT_MS, signal, fetchImpl);
	const rawResults = Array.isArray(data.results) ? (data.results as Record<string, unknown>[]) : [];
	const rawFailed = Array.isArray(data.failed_results) ? (data.failed_results as Record<string, unknown>[]) : [];
	return {
		results: rawResults.map((result) => ({ url: result.url as string | undefined, content: result.raw_content as string | undefined })),
		failed: rawFailed.map((result) => String(result.url ?? "unknown")),
	};
}