Add ollama-chat API provider supporting: - Streaming chat completions via Ollama Cloud (ollama.com) API - API key authentication via OLLAMA_API_KEY env var - Thinking/reasoning model support with configurable effort levels - Tool calling with streaming JSON argument assembly - Dynamic model discovery via /api/tags and /api/show metadata - Context window detection from model info - Login flow via pi login ollama-cloud
97 lines
2.2 KiB
TypeScript
97 lines
2.2 KiB
TypeScript
/**
|
|
* Custom API provider registry.
|
|
*
|
|
* Allows extensions to register streaming functions for custom API types
|
|
* (e.g., "vertex-claude-api") that are not built into stream.ts.
|
|
*/
|
|
import type {
|
|
Api,
|
|
AssistantMessageEventStream,
|
|
Context,
|
|
KnownApi,
|
|
Model,
|
|
SimpleStreamOptions,
|
|
StreamOptions,
|
|
} from "./types";
|
|
|
|
const BUILTIN_APIS = new Set<KnownApi>([
|
|
"openai-completions",
|
|
"openai-responses",
|
|
"openai-codex-responses",
|
|
"azure-openai-responses",
|
|
"anthropic-messages",
|
|
"bedrock-converse-stream",
|
|
"google-generative-ai",
|
|
"google-gemini-cli",
|
|
"google-vertex",
|
|
"ollama-chat",
|
|
"cursor-agent",
|
|
]);
|
|
|
|
export type CustomStreamFn = (
|
|
model: Model<Api>,
|
|
context: Context,
|
|
options?: StreamOptions,
|
|
) => AssistantMessageEventStream;
|
|
export type CustomStreamSimpleFn = (
|
|
model: Model<Api>,
|
|
context: Context,
|
|
options?: SimpleStreamOptions,
|
|
) => AssistantMessageEventStream;
|
|
|
|
export interface RegisteredCustomApi {
|
|
stream: CustomStreamFn;
|
|
streamSimple: CustomStreamSimpleFn;
|
|
sourceId?: string;
|
|
}
|
|
|
|
const customApiRegistry = new Map<string, RegisteredCustomApi>();
|
|
|
|
function assertCustomApiName(api: string): void {
|
|
if (BUILTIN_APIS.has(api as KnownApi)) {
|
|
throw new Error(`Cannot register custom API "${api}": built-in API names are reserved.`);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Register a custom API streaming function.
|
|
*/
|
|
export function registerCustomApi(
|
|
api: string,
|
|
streamSimple: CustomStreamSimpleFn,
|
|
sourceId?: string,
|
|
stream?: CustomStreamFn,
|
|
): void {
|
|
assertCustomApiName(api);
|
|
customApiRegistry.set(api, {
|
|
stream: stream ?? ((model, context, options) => streamSimple(model, context, options as SimpleStreamOptions)),
|
|
streamSimple,
|
|
sourceId,
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Get a custom API provider by API identifier.
|
|
*/
|
|
export function getCustomApi(api: string): RegisteredCustomApi | undefined {
|
|
return customApiRegistry.get(api);
|
|
}
|
|
|
|
/**
|
|
* Remove all custom APIs registered by a specific source (e.g., extension path).
|
|
*/
|
|
export function unregisterCustomApis(sourceId: string): void {
|
|
for (const [api, entry] of customApiRegistry.entries()) {
|
|
if (entry.sourceId === sourceId) {
|
|
customApiRegistry.delete(api);
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Clear all custom API registrations.
|
|
*/
|
|
export function clearCustomApis(): void {
|
|
customApiRegistry.clear();
|
|
}
|