@oh-my-pi/pi-ai 0.1.0

package/src/stream.ts

@@ -0,0 +1,340 @@
+import { supportsXhigh } from "./models";
+import { type AnthropicOptions, streamAnthropic } from "./providers/anthropic";
+import { type GoogleOptions, streamGoogle } from "./providers/google";
+import {
+	type GoogleGeminiCliOptions,
+	type GoogleThinkingLevel,
+	streamGoogleGeminiCli,
+} from "./providers/google-gemini-cli";
+import { type OpenAICompletionsOptions, streamOpenAICompletions } from "./providers/openai-completions";
+import { type OpenAIResponsesOptions, streamOpenAIResponses } from "./providers/openai-responses";
+import type {
+	Api,
+	AssistantMessage,
+	AssistantMessageEventStream,
+	Context,
+	KnownProvider,
+	Model,
+	OptionsForApi,
+	ReasoningEffort,
+	SimpleStreamOptions,
+} from "./types";
+
+/**
+ * Get API key for provider from known environment variables, e.g. OPENAI_API_KEY.
+ *
+ * Will not return API keys for providers that require OAuth tokens.
+ */
+export function getEnvApiKey(provider: KnownProvider): string | undefined;
+export function getEnvApiKey(provider: string): string | undefined;
+export function getEnvApiKey(provider: any): string | undefined {
+	// Fall back to environment variables
+	if (provider === "github-copilot") {
+		return process.env.COPILOT_GITHUB_TOKEN || process.env.GH_TOKEN || process.env.GITHUB_TOKEN;
+	}
+
+	// ANTHROPIC_OAUTH_TOKEN takes precedence over ANTHROPIC_API_KEY
+	if (provider === "anthropic") {
+		return process.env.ANTHROPIC_OAUTH_TOKEN || process.env.ANTHROPIC_API_KEY;
+	}
+
+	const envMap: Record<string, string> = {
+		openai: "OPENAI_API_KEY",
+		google: "GEMINI_API_KEY",
+		groq: "GROQ_API_KEY",
+		cerebras: "CEREBRAS_API_KEY",
+		xai: "XAI_API_KEY",
+		openrouter: "OPENROUTER_API_KEY",
+		zai: "ZAI_API_KEY",
+		mistral: "MISTRAL_API_KEY",
+	};
+
+	const envVar = envMap[provider];
+	return envVar ? process.env[envVar] : undefined;
+}
+
+export function stream<TApi extends Api>(
+	model: Model<TApi>,
+	context: Context,
+	options?: OptionsForApi<TApi>,
+): AssistantMessageEventStream {
+	const apiKey = options?.apiKey || getEnvApiKey(model.provider);
+	if (!apiKey) {
+		throw new Error(`No API key for provider: ${model.provider}`);
+	}
+	const providerOptions = { ...options, apiKey };
+
+	const api: Api = model.api;
+	switch (api) {
+		case "anthropic-messages":
+			return streamAnthropic(model as Model<"anthropic-messages">, context, providerOptions);
+
+		case "openai-completions":
+			return streamOpenAICompletions(model as Model<"openai-completions">, context, providerOptions as any);
+
+		case "openai-responses":
+			return streamOpenAIResponses(model as Model<"openai-responses">, context, providerOptions as any);
+
+		case "google-generative-ai":
+			return streamGoogle(model as Model<"google-generative-ai">, context, providerOptions);
+
+		case "google-gemini-cli":
+			return streamGoogleGeminiCli(
+				model as Model<"google-gemini-cli">,
+				context,
+				providerOptions as GoogleGeminiCliOptions,
+			);
+
+		default: {
+			// This should never be reached if all Api cases are handled
+			const _exhaustive: never = api;
+			throw new Error(`Unhandled API: ${_exhaustive}`);
+		}
+	}
+}
+
+export async function complete<TApi extends Api>(
+	model: Model<TApi>,
+	context: Context,
+	options?: OptionsForApi<TApi>,
+): Promise<AssistantMessage> {
+	const s = stream(model, context, options);
+	return s.result();
+}
+
+export function streamSimple<TApi extends Api>(
+	model: Model<TApi>,
+	context: Context,
+	options?: SimpleStreamOptions,
+): AssistantMessageEventStream {
+	const apiKey = options?.apiKey || getEnvApiKey(model.provider);
+	if (!apiKey) {
+		throw new Error(`No API key for provider: ${model.provider}`);
+	}
+
+	const providerOptions = mapOptionsForApi(model, options, apiKey);
+	return stream(model, context, providerOptions);
+}
+
+export async function completeSimple<TApi extends Api>(
+	model: Model<TApi>,
+	context: Context,
+	options?: SimpleStreamOptions,
+): Promise<AssistantMessage> {
+	const s = streamSimple(model, context, options);
+	return s.result();
+}
+
+function mapOptionsForApi<TApi extends Api>(
+	model: Model<TApi>,
+	options?: SimpleStreamOptions,
+	apiKey?: string,
+): OptionsForApi<TApi> {
+	const base = {
+		temperature: options?.temperature,
+		maxTokens: options?.maxTokens || Math.min(model.maxTokens, 32000),
+		signal: options?.signal,
+		apiKey: apiKey || options?.apiKey,
+	};
+
+	// Helper to clamp xhigh to high for providers that don't support it
+	const clampReasoning = (effort: ReasoningEffort | undefined) => (effort === "xhigh" ? "high" : effort);
+
+	switch (model.api) {
+		case "anthropic-messages": {
+			// Explicitly disable thinking when reasoning is not specified
+			if (!options?.reasoning) {
+				return { ...base, thinkingEnabled: false } satisfies AnthropicOptions;
+			}
+
+			const anthropicBudgets = {
+				minimal: 1024,
+				low: 2048,
+				medium: 8192,
+				high: 16384,
+			};
+
+			return {
+				...base,
+				thinkingEnabled: true,
+				thinkingBudgetTokens: anthropicBudgets[clampReasoning(options.reasoning)!],
+			} satisfies AnthropicOptions;
+		}
+
+		case "openai-completions":
+			return {
+				...base,
+				reasoningEffort: supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning),
+			} satisfies OpenAICompletionsOptions;
+
+		case "openai-responses":
+			return {
+				...base,
+				reasoningEffort: supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning),
+			} satisfies OpenAIResponsesOptions;
+
+		case "google-generative-ai": {
+			// Explicitly disable thinking when reasoning is not specified
+			// This is needed because Gemini has "dynamic thinking" enabled by default
+			if (!options?.reasoning) {
+				return { ...base, thinking: { enabled: false } } satisfies GoogleOptions;
+			}
+
+			const googleModel = model as Model<"google-generative-ai">;
+			const effort = clampReasoning(options.reasoning)!;
+
+			// Gemini 3 models use thinkingLevel exclusively instead of thinkingBudget.
+			// https://ai.google.dev/gemini-api/docs/thinking#set-budget
+			if (isGemini3ProModel(googleModel) || isGemini3FlashModel(googleModel)) {
+				return {
+					...base,
+					thinking: {
+						enabled: true,
+						level: getGemini3ThinkingLevel(effort, googleModel),
+					},
+				} satisfies GoogleOptions;
+			}
+
+			return {
+				...base,
+				thinking: {
+					enabled: true,
+					budgetTokens: getGoogleBudget(googleModel, effort),
+				},
+			} satisfies GoogleOptions;
+		}
+
+		case "google-gemini-cli": {
+			if (!options?.reasoning) {
+				return { ...base, thinking: { enabled: false } } satisfies GoogleGeminiCliOptions;
+			}
+
+			const effort = clampReasoning(options.reasoning)!;
+
+			// Gemini 3 models use thinkingLevel instead of thinkingBudget
+			if (model.id.includes("3-pro") || model.id.includes("3-flash")) {
+				return {
+					...base,
+					thinking: {
+						enabled: true,
+						level: getGeminiCliThinkingLevel(effort, model.id),
+					},
+				} satisfies GoogleGeminiCliOptions;
+			}
+
+			// Gemini 2.x models use thinkingBudget
+			const budgets: Record<ClampedReasoningEffort, number> = {
+				minimal: 1024,
+				low: 2048,
+				medium: 8192,
+				high: 16384,
+			};
+
+			return {
+				...base,
+				thinking: {
+					enabled: true,
+					budgetTokens: budgets[effort],
+				},
+			} satisfies GoogleGeminiCliOptions;
+		}
+
+		default: {
+			// Exhaustiveness check
+			const _exhaustive: never = model.api;
+			throw new Error(`Unhandled API in mapOptionsForApi: ${_exhaustive}`);
+		}
+	}
+}
+
+type ClampedReasoningEffort = Exclude<ReasoningEffort, "xhigh">;
+
+function isGemini3ProModel(model: Model<"google-generative-ai">): boolean {
+	// Covers gemini-3-pro, gemini-3-pro-preview, and possible other prefixed ids in the future
+	return model.id.includes("3-pro");
+}
+
+function isGemini3FlashModel(model: Model<"google-generative-ai">): boolean {
+	// Covers gemini-3-flash, gemini-3-flash-preview, and possible other prefixed ids in the future
+	return model.id.includes("3-flash");
+}
+
+function getGemini3ThinkingLevel(
+	effort: ClampedReasoningEffort,
+	model: Model<"google-generative-ai">,
+): GoogleThinkingLevel {
+	if (isGemini3ProModel(model)) {
+		// Gemini 3 Pro only supports LOW/HIGH (for now)
+		switch (effort) {
+			case "minimal":
+			case "low":
+				return "LOW";
+			case "medium":
+			case "high":
+				return "HIGH";
+		}
+	}
+	// Gemini 3 Flash supports all four levels
+	switch (effort) {
+		case "minimal":
+			return "MINIMAL";
+		case "low":
+			return "LOW";
+		case "medium":
+			return "MEDIUM";
+		case "high":
+			return "HIGH";
+	}
+}
+
+function getGeminiCliThinkingLevel(effort: ClampedReasoningEffort, modelId: string): GoogleThinkingLevel {
+	if (modelId.includes("3-pro")) {
+		// Gemini 3 Pro only supports LOW/HIGH (for now)
+		switch (effort) {
+			case "minimal":
+			case "low":
+				return "LOW";
+			case "medium":
+			case "high":
+				return "HIGH";
+		}
+	}
+	// Gemini 3 Flash supports all four levels
+	switch (effort) {
+		case "minimal":
+			return "MINIMAL";
+		case "low":
+			return "LOW";
+		case "medium":
+			return "MEDIUM";
+		case "high":
+			return "HIGH";
+	}
+}
+
+function getGoogleBudget(model: Model<"google-generative-ai">, effort: ClampedReasoningEffort): number {
+	// See https://ai.google.dev/gemini-api/docs/thinking#set-budget
+	if (model.id.includes("2.5-pro")) {
+		const budgets: Record<ClampedReasoningEffort, number> = {
+			minimal: 128,
+			low: 2048,
+			medium: 8192,
+			high: 32768,
+		};
+		return budgets[effort];
+	}
+
+	if (model.id.includes("2.5-flash")) {
+		// Covers 2.5-flash-lite as well
+		const budgets: Record<ClampedReasoningEffort, number> = {
+			minimal: 128,
+			low: 2048,
+			medium: 8192,
+			high: 24576,
+		};
+		return budgets[effort];
+	}
+
+	// Unknown model - use dynamic
+	return -1;
+}

package/src/types.ts

@@ -0,0 +1,218 @@
+import type { AnthropicOptions } from "./providers/anthropic";
+import type { GoogleOptions } from "./providers/google";
+import type { GoogleGeminiCliOptions } from "./providers/google-gemini-cli";
+import type { OpenAICompletionsOptions } from "./providers/openai-completions";
+import type { OpenAIResponsesOptions } from "./providers/openai-responses";
+import type { AssistantMessageEventStream } from "./utils/event-stream";
+
+export type { AssistantMessageEventStream } from "./utils/event-stream";
+
+export type Api =
+	| "openai-completions"
+	| "openai-responses"
+	| "anthropic-messages"
+	| "google-generative-ai"
+	| "google-gemini-cli";
+
+export interface ApiOptionsMap {
+	"anthropic-messages": AnthropicOptions;
+	"openai-completions": OpenAICompletionsOptions;
+	"openai-responses": OpenAIResponsesOptions;
+	"google-generative-ai": GoogleOptions;
+	"google-gemini-cli": GoogleGeminiCliOptions;
+}
+
+// Compile-time exhaustiveness check - this will fail if ApiOptionsMap doesn't have all KnownApi keys
+type _CheckExhaustive =
+	ApiOptionsMap extends Record<Api, StreamOptions>
+		? Record<Api, StreamOptions> extends ApiOptionsMap
+			? true
+			: ["ApiOptionsMap is missing some KnownApi values", Exclude<Api, keyof ApiOptionsMap>]
+		: ["ApiOptionsMap doesn't extend Record<KnownApi, StreamOptions>"];
+const _exhaustive: _CheckExhaustive = true;
+
+// Helper type to get options for a specific API
+export type OptionsForApi<TApi extends Api> = ApiOptionsMap[TApi];
+
+export type KnownProvider =
+	| "anthropic"
+	| "google"
+	| "google-gemini-cli"
+	| "google-antigravity"
+	| "openai"
+	| "github-copilot"
+	| "xai"
+	| "groq"
+	| "cerebras"
+	| "openrouter"
+	| "zai"
+	| "mistral";
+export type Provider = KnownProvider | string;
+
+export type ReasoningEffort = "minimal" | "low" | "medium" | "high" | "xhigh";
+
+// Base options all providers share
+export interface StreamOptions {
+	temperature?: number;
+	maxTokens?: number;
+	signal?: AbortSignal;
+	apiKey?: string;
+}
+
+// Unified options with reasoning passed to streamSimple() and completeSimple()
+export interface SimpleStreamOptions extends StreamOptions {
+	reasoning?: ReasoningEffort;
+}
+
+// Generic StreamFunction with typed options
+export type StreamFunction<TApi extends Api> = (
+	model: Model<TApi>,
+	context: Context,
+	options: OptionsForApi<TApi>,
+) => AssistantMessageEventStream;
+
+export interface TextContent {
+	type: "text";
+	text: string;
+	textSignature?: string; // e.g., for OpenAI responses, the message ID
+}
+
+export interface ThinkingContent {
+	type: "thinking";
+	thinking: string;
+	thinkingSignature?: string; // e.g., for OpenAI responses, the reasoning item ID
+}
+
+export interface ImageContent {
+	type: "image";
+	data: string; // base64 encoded image data
+	mimeType: string; // e.g., "image/jpeg", "image/png"
+}
+
+export interface ToolCall {
+	type: "toolCall";
+	id: string;
+	name: string;
+	arguments: Record<string, any>;
+	thoughtSignature?: string; // Google-specific: opaque signature for reusing thought context
+}
+
+export interface Usage {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	totalTokens: number;
+	cost: {
+		input: number;
+		output: number;
+		cacheRead: number;
+		cacheWrite: number;
+		total: number;
+	};
+}
+
+export type StopReason = "stop" | "length" | "toolUse" | "error" | "aborted";
+
+export interface UserMessage {
+	role: "user";
+	content: string | (TextContent | ImageContent)[];
+	timestamp: number; // Unix timestamp in milliseconds
+}
+
+export interface AssistantMessage {
+	role: "assistant";
+	content: (TextContent | ThinkingContent | ToolCall)[];
+	api: Api;
+	provider: Provider;
+	model: string;
+	usage: Usage;
+	stopReason: StopReason;
+	errorMessage?: string;
+	timestamp: number; // Unix timestamp in milliseconds
+}
+
+export interface ToolResultMessage<TDetails = any> {
+	role: "toolResult";
+	toolCallId: string;
+	toolName: string;
+	content: (TextContent | ImageContent)[]; // Supports text and images
+	details?: TDetails;
+	isError?: boolean;
+	timestamp: number; // Unix timestamp in milliseconds
+}
+
+export type Message = UserMessage | AssistantMessage | ToolResultMessage;
+
+import type { TSchema } from "@sinclair/typebox";
+
+export interface Tool<TParameters extends TSchema = TSchema> {
+	name: string;
+	description: string;
+	parameters: TParameters;
+}
+
+export interface Context {
+	systemPrompt?: string;
+	messages: Message[];
+	tools?: Tool[];
+}
+
+export type AssistantMessageEvent =
+	| { type: "start"; partial: AssistantMessage }
+	| { type: "text_start"; contentIndex: number; partial: AssistantMessage }
+	| { type: "text_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
+	| { type: "text_end"; contentIndex: number; content: string; partial: AssistantMessage }
+	| { type: "thinking_start"; contentIndex: number; partial: AssistantMessage }
+	| { type: "thinking_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
+	| { type: "thinking_end"; contentIndex: number; content: string; partial: AssistantMessage }
+	| { type: "toolcall_start"; contentIndex: number; partial: AssistantMessage }
+	| { type: "toolcall_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
+	| { type: "toolcall_end"; contentIndex: number; toolCall: ToolCall; partial: AssistantMessage }
+	| { type: "done"; reason: Extract<StopReason, "stop" | "length" | "toolUse">; message: AssistantMessage }
+	| { type: "error"; reason: Extract<StopReason, "aborted" | "error">; error: AssistantMessage };
+
+/**
+ * Compatibility settings for openai-completions API.
+ * Use this to override URL-based auto-detection for custom providers.
+ */
+export interface OpenAICompat {
+	/** Whether the provider supports the `store` field. Default: auto-detected from URL. */
+	supportsStore?: boolean;
+	/** Whether the provider supports the `developer` role (vs `system`). Default: auto-detected from URL. */
+	supportsDeveloperRole?: boolean;
+	/** Whether the provider supports `reasoning_effort`. Default: auto-detected from URL. */
+	supportsReasoningEffort?: boolean;
+	/** Which field to use for max tokens. Default: auto-detected from URL. */
+	maxTokensField?: "max_completion_tokens" | "max_tokens";
+	/** Whether tool results require the `name` field. Default: auto-detected from URL. */
+	requiresToolResultName?: boolean;
+	/** Whether a user message after tool results requires an assistant message in between. Default: auto-detected from URL. */
+	requiresAssistantAfterToolResult?: boolean;
+	/** Whether thinking blocks must be converted to text blocks with <thinking> delimiters. Default: auto-detected from URL. */
+	requiresThinkingAsText?: boolean;
+	/** Whether tool call IDs must be normalized to Mistral format (exactly 9 alphanumeric chars). Default: auto-detected from URL. */
+	requiresMistralToolIds?: boolean;
+}
+
+// Model interface for the unified model system
+export interface Model<TApi extends Api> {
+	id: string;
+	name: string;
+	api: TApi;
+	provider: Provider;
+	baseUrl: string;
+	reasoning: boolean;
+	input: ("text" | "image")[];
+	cost: {
+		input: number; // $/million tokens
+		output: number; // $/million tokens
+		cacheRead: number; // $/million tokens
+		cacheWrite: number; // $/million tokens
+	};
+	contextWindow: number;
+	maxTokens: number;
+	headers?: Record<string, string>;
+	/** Compatibility overrides for openai-completions API. If not set, auto-detected from baseUrl. */
+	compat?: TApi extends "openai-completions" ? OpenAICompat : never;
+}

package/src/utils/event-stream.ts

@@ -0,0 +1,82 @@
+import type { AssistantMessage, AssistantMessageEvent } from "../types";
+
+// Generic event stream class for async iteration
+export class EventStream<T, R = T> implements AsyncIterable<T> {
+	private queue: T[] = [];
+	private waiting: ((value: IteratorResult<T>) => void)[] = [];
+	private done = false;
+	private finalResultPromise: Promise<R>;
+	private resolveFinalResult!: (result: R) => void;
+
+	constructor(
+		private isComplete: (event: T) => boolean,
+		private extractResult: (event: T) => R,
+	) {
+		this.finalResultPromise = new Promise((resolve) => {
+			this.resolveFinalResult = resolve;
+		});
+	}
+
+	push(event: T): void {
+		if (this.done) return;
+
+		if (this.isComplete(event)) {
+			this.done = true;
+			this.resolveFinalResult(this.extractResult(event));
+		}
+
+		// Deliver to waiting consumer or queue it
+		const waiter = this.waiting.shift();
+		if (waiter) {
+			waiter({ value: event, done: false });
+		} else {
+			this.queue.push(event);
+		}
+	}
+
+	end(result?: R): void {
+		this.done = true;
+		if (result !== undefined) {
+			this.resolveFinalResult(result);
+		}
+		// Notify all waiting consumers that we're done
+		while (this.waiting.length > 0) {
+			const waiter = this.waiting.shift()!;
+			waiter({ value: undefined as any, done: true });
+		}
+	}
+
+	async *[Symbol.asyncIterator](): AsyncIterator<T> {
+		while (true) {
+			if (this.queue.length > 0) {
+				yield this.queue.shift()!;
+			} else if (this.done) {
+				return;
+			} else {
+				const result = await new Promise<IteratorResult<T>>((resolve) => this.waiting.push(resolve));
+				if (result.done) return;
+				yield result.value;
+			}
+		}
+	}
+
+	result(): Promise<R> {
+		return this.finalResultPromise;
+	}
+}
+
+export class AssistantMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
+	constructor() {
+		super(
+			(event) => event.type === "done" || event.type === "error",
+			(event) => {
+				if (event.type === "done") {
+					return event.message;
+				} else if (event.type === "error") {
+					return event.error;
+				}
+				throw new Error("Unexpected event type for final result");
+			},
+		);
+	}
+}

package/src/utils/json-parse.ts

@@ -0,0 +1,28 @@
+import { parse as partialParse } from "partial-json";
+
+/**
+ * Attempts to parse potentially incomplete JSON during streaming.
+ * Always returns a valid object, even if the JSON is incomplete.
+ *
+ * @param partialJson The partial JSON string from streaming
+ * @returns Parsed object or empty object if parsing fails
+ */
+export function parseStreamingJson<T = any>(partialJson: string | undefined): T {
+	if (!partialJson || partialJson.trim() === "") {
+		return {} as T;
+	}
+
+	// Try standard parsing first (fastest for complete JSON)
+	try {
+		return JSON.parse(partialJson) as T;
+	} catch {
+		// Try partial-json for incomplete JSON
+		try {
+			const result = partialParse(partialJson);
+			return (result ?? {}) as T;
+		} catch {
+			// If all parsing fails, return empty object
+			return {} as T;
+		}
+	}
+}