@oh-my-pi/pi-coding-agent 12.6.0 → 12.7.0

package/src/web/search/providers/zai.ts

@@ -0,0 +1,352 @@
+/**
+ * Z.AI Web Search Provider
+ *
+ * Calls Z.AI's remote MCP server (`webSearchPrime`) and adapts results into
+ * the unified SearchResponse shape used by the web search tool.
+ */
+import { getEnvApiKey } from "@oh-my-pi/pi-ai";
+import { getAgentDbPath } from "@oh-my-pi/pi-utils/dirs";
+import { AgentStorage } from "../../../session/agent-storage";
+
+import type { SearchResponse, SearchSource } from "../../../web/search/types";
+import { SearchProviderError } from "../../../web/search/types";
+import type { SearchParams } from "./base";
+import { SearchProvider } from "./base";
+
+const ZAI_MCP_URL = "https://api.z.ai/api/mcp/web_search_prime/mcp";
+const ZAI_TOOL_NAME = "webSearchPrime";
+const DEFAULT_NUM_RESULTS = 10;
+
+export interface ZaiSearchParams {
+	query: string;
+	num_results?: number;
+}
+
+interface ZaiSearchResult {
+	title?: string;
+	content?: string;
+	link?: string;
+	url?: string;
+	media?: string;
+	publish_date?: string;
+	publishedDate?: string;
+}
+
+interface ZaiWebSearchResponse {
+	id?: string;
+	request_id?: string;
+	requestId?: string;
+	search_result?: ZaiSearchResult[];
+	results?: ZaiSearchResult[];
+}
+
+interface JsonRpcError {
+	code?: number;
+	message?: string;
+}
+
+interface JsonRpcPayload {
+	result?: unknown;
+	error?: JsonRpcError;
+}
+
+function asRecord(value: unknown): Record<string, unknown> | null {
+	return typeof value === "object" && value !== null ? (value as Record<string, unknown>) : null;
+}
+
+function asString(value: unknown): string | undefined {
+	return typeof value === "string" && value.trim().length > 0 ? value : undefined;
+}
+
+/**
+ * Finds Z.AI API credentials from environment or saved auth storage.
+ * Priority: ZAI_API_KEY env var, then credentials stored under provider "zai".
+ */
+export async function findApiKey(): Promise<string | null> {
+	const envKey = getEnvApiKey("zai");
+	if (envKey) return envKey;
+
+	try {
+		const storage = await AgentStorage.open(getAgentDbPath());
+		const records = storage.listAuthCredentials("zai");
+		for (const record of records) {
+			const credential = record.credential;
+			if (credential.type === "api_key" && credential.key.trim().length > 0) {
+				return credential.key;
+			}
+			if (credential.type === "oauth" && credential.access.trim().length > 0) {
+				return credential.access;
+			}
+		}
+	} catch {
+		return null;
+	}
+
+	return null;
+}
+
+function dateToAgeSeconds(dateStr: string | undefined): number | undefined {
+	if (!dateStr) return undefined;
+	try {
+		const date = new Date(dateStr);
+		if (Number.isNaN(date.getTime())) return undefined;
+		return Math.floor((Date.now() - date.getTime()) / 1000);
+	} catch {
+		return undefined;
+	}
+}
+
+async function callZaiTool(apiKey: string, args: Record<string, unknown>): Promise<unknown> {
+	const response = await fetch(ZAI_MCP_URL, {
+		method: "POST",
+		headers: {
+			Authorization: `Bearer ${apiKey}`,
+			"Content-Type": "application/json",
+			Accept: "application/json, text/event-stream",
+		},
+		body: JSON.stringify({
+			jsonrpc: "2.0",
+			id: crypto.randomUUID(),
+			method: "tools/call",
+			params: {
+				name: ZAI_TOOL_NAME,
+				arguments: args,
+			},
+		}),
+	});
+
+	if (!response.ok) {
+		const errorText = await response.text();
+		throw new SearchProviderError("zai", `Z.AI MCP error (${response.status}): ${errorText}`, response.status);
+	}
+
+	const rawText = await response.text();
+
+	const parsedMessages: unknown[] = [];
+	for (const line of rawText.split("\n")) {
+		const trimmed = line.trim();
+		if (!trimmed.startsWith("data:")) continue;
+		const data = trimmed.slice(5).trim();
+		if (!data) continue;
+		try {
+			parsedMessages.push(JSON.parse(data));
+		} catch {
+			// Ignore non-JSON data events.
+		}
+	}
+
+	if (parsedMessages.length === 0) {
+		try {
+			parsedMessages.push(JSON.parse(rawText));
+		} catch {
+			throw new SearchProviderError("zai", "Failed to parse Z.AI MCP response", 500);
+		}
+	}
+
+	const parsed = parsedMessages[parsedMessages.length - 1];
+	const parsedRecord = asRecord(parsed);
+	const directErrorCode = typeof parsedRecord?.code === "number" ? parsedRecord.code : undefined;
+	const directErrorSuccess = parsedRecord?.success;
+	const directErrorMessage =
+		asString(parsedRecord?.msg) ?? asString(parsedRecord?.message) ?? asString(parsedRecord?.error_message);
+	if (directErrorSuccess === false && directErrorMessage) {
+		throw new SearchProviderError(
+			"zai",
+			`Z.AI API error${directErrorCode ? ` (${directErrorCode})` : ""}: ${directErrorMessage}`,
+			directErrorCode,
+		);
+	}
+
+	const payload = parsed as JsonRpcPayload;
+	if (payload.error) {
+		const status = typeof payload.error.code === "number" ? payload.error.code : 400;
+		throw new SearchProviderError(
+			"zai",
+			`Z.AI MCP error${payload.error.code ? ` (${payload.error.code})` : ""}: ${payload.error.message ?? "Unknown error"}`,
+			status,
+		);
+	}
+
+	const resultRecord = asRecord(payload.result);
+	if (resultRecord?.isError === true) {
+		const content = Array.isArray(resultRecord.content) ? resultRecord.content : [];
+		const errorText = content
+			.map(item => asString(asRecord(item)?.text))
+			.filter((text): text is string => text !== undefined)
+			.join("\n")
+			.trim();
+		const statusMatch = errorText.match(/MCP error\s*(-?\d+)/i);
+		const statusCode = statusMatch ? Math.abs(Number.parseInt(statusMatch[1], 10)) : 400;
+		throw new SearchProviderError("zai", errorText || "Z.AI MCP tool call failed", statusCode);
+	}
+
+	if (payload.result !== undefined) {
+		return payload.result;
+	}
+
+	return parsed;
+}
+
+async function callZaiSearch(apiKey: string, params: ZaiSearchParams): Promise<unknown> {
+	const count = params.num_results ?? DEFAULT_NUM_RESULTS;
+	const attempts: Record<string, unknown>[] = [
+		{ query: params.query, count },
+		{ search_query: params.query, count },
+		{ search_query: params.query, search_engine: "search-prime", count },
+	];
+
+	let lastError: unknown;
+	for (let i = 0; i < attempts.length; i++) {
+		try {
+			return await callZaiTool(apiKey, attempts[i]);
+		} catch (error) {
+			lastError = error;
+			const isLastAttempt = i === attempts.length - 1;
+			if (isLastAttempt) {
+				throw error;
+			}
+			if (!(error instanceof SearchProviderError)) {
+				throw error;
+			}
+			const message = error.message.toLowerCase();
+			const looksLikeArgError =
+				error.status === 400 ||
+				message.includes("invalid") ||
+				message.includes("argument") ||
+				message.includes("search_query") ||
+				message.includes("query");
+			if (!looksLikeArgError) {
+				throw error;
+			}
+		}
+	}
+
+	throw lastError ?? new SearchProviderError("zai", "Z.AI search failed", 500);
+}
+
+function getSearchResults(value: unknown): ZaiSearchResult[] {
+	if (Array.isArray(value)) {
+		return value as ZaiSearchResult[];
+	}
+	const obj = asRecord(value);
+	if (!obj) return [];
+
+	const searchResult = obj.search_result;
+	if (Array.isArray(searchResult)) return searchResult as ZaiSearchResult[];
+
+	const results = obj.results;
+	if (Array.isArray(results)) return results as ZaiSearchResult[];
+
+	return [];
+}
+
+function parseSearchPayload(rawResult: unknown): {
+	results: ZaiSearchResult[];
+	answer?: string;
+	requestId?: string;
+} {
+	const candidates: unknown[] = [rawResult];
+	const textParts: string[] = [];
+
+	const root = asRecord(rawResult);
+	if (root) {
+		if (root.structuredContent) candidates.push(root.structuredContent);
+		if (root.data) candidates.push(root.data);
+		if (root.result) candidates.push(root.result);
+
+		const content = root.content;
+		if (Array.isArray(content)) {
+			for (const part of content) {
+				const partObj = asRecord(part);
+				const text = asString(partObj?.text);
+				if (!text) continue;
+				textParts.push(text);
+				try {
+					candidates.push(JSON.parse(text));
+				} catch {
+					// Not JSON payload; keep as fallback answer text.
+				}
+			}
+		}
+	}
+
+	for (const candidate of candidates) {
+		const results = getSearchResults(candidate);
+		if (results.length > 0) {
+			const obj = asRecord(candidate) as ZaiWebSearchResponse | null;
+			return {
+				results,
+				answer: textParts.length > 0 ? textParts.join("\n\n") : undefined,
+				requestId: obj?.request_id ?? obj?.requestId ?? obj?.id,
+			};
+		}
+	}
+
+	return {
+		results: [],
+		answer: textParts.length > 0 ? textParts.join("\n\n") : undefined,
+	};
+}
+
+function toSources(results: ZaiSearchResult[]): SearchSource[] {
+	const sources: SearchSource[] = [];
+	for (const result of results) {
+		const url = asString(result.link) ?? asString(result.url);
+		if (!url) continue;
+
+		const publishedDate = asString(result.publish_date) ?? asString(result.publishedDate);
+		sources.push({
+			title: asString(result.title) ?? url,
+			url,
+			snippet: asString(result.content),
+			publishedDate,
+			ageSeconds: dateToAgeSeconds(publishedDate),
+			author: asString(result.media),
+		});
+	}
+	return sources;
+}
+
+/** Execute Z.AI web search via remote MCP endpoint. */
+export async function searchZai(params: ZaiSearchParams): Promise<SearchResponse> {
+	const apiKey = await findApiKey();
+	if (!apiKey) {
+		throw new Error("Z.AI credentials not found. Set ZAI_API_KEY or login with 'omp /login zai'.");
+	}
+
+	const rawResult = await callZaiSearch(apiKey, params);
+	const payload = parseSearchPayload(rawResult);
+	let sources = toSources(payload.results);
+
+	if (params.num_results && sources.length > params.num_results) {
+		sources = sources.slice(0, params.num_results);
+	}
+
+	return {
+		provider: "zai",
+		answer: payload.answer,
+		sources,
+		requestId: payload.requestId,
+	};
+}
+
+/** Search provider for Z.AI web search MCP. */
+export class ZaiProvider extends SearchProvider {
+	readonly id = "zai";
+	readonly label = "Z.AI";
+
+	async isAvailable(): Promise<boolean> {
+		try {
+			return !!(await findApiKey());
+		} catch {
+			return false;
+		}
+	}
+
+	search(params: SearchParams): Promise<SearchResponse> {
+		return searchZai({
+			query: params.query,
+			num_results: params.numSearchResults ?? params.limit,
+		});
+	}
+}

package/src/web/search/types.ts

@@ -5,7 +5,7 @@
  */
 
 /** Supported web search providers */
-export type SearchProviderId = "exa" | "jina" | "anthropic" | "perplexity" | "gemini" | "codex";
+export type SearchProviderId = "exa" | "jina" | "zai" | "anthropic" | "perplexity" | "gemini" | "codex";
 
 /** Source returned by search (all providers) */
 export interface SearchSource {