packages ai

AI

Active contributors: Mario Zechner, kt, Armin Ronacher

Purpose

packages/ai is @earendil-works/pi-ai, the LLM provider abstraction shared by the whole repo. It turns a provider-agnostic Model plus a Context into a streaming AssistantMessageEventStream, with provider-specific code isolated behind one interface. It also owns the model catalog, OAuth login, the MCP client, and token and cost accounting.

Everything the rest of the codebase talks to a language model through lives here:

• Streaming primitives: stream, streamSimple, complete, completeSimple (packages/ai/src/stream.ts) over a push-based AssistantMessageEventStream (packages/ai/src/utils/event-stream.ts).
• A provider registry that maps an Api identifier to stream and streamSimple functions (packages/ai/src/api-registry.ts), populated lazily by packages/ai/src/providers/register-builtins.ts.
• Nine built-in APIs (anthropic-messages, openai-responses, openai-completions, openai-codex-responses, azure-openai-responses, google-generative-ai, google-vertex, mistral-conversations, bedrock-converse-stream) implemented under packages/ai/src/providers/, plus a test-only faux provider.
• A generated model registry of 1172 models across 31 providers (packages/ai/src/models.generated.ts, packages/ai/src/models.ts), cost calculation, and thinking-level support helpers.
• Credential handling: env API key detection (packages/ai/src/env-api-keys.ts), OAuth login flows for Anthropic, GitHub Copilot, and OpenAI Codex (packages/ai/src/utils/oauth/), and a generic OAuth 2.1 provider for MCP servers (packages/ai/src/mcp/oauth.ts).
• Stream plumbing shared by all providers: failure classification, JSON repair, Unicode sanitization, context-overflow detection, and tool-call validation (packages/ai/src/utils/).

The package is browser-safe at the entry level: heavy or node-only imports (AWS SDK, OAuth callback servers) are deferred so bundlers and the web build do not pull them in.

Directory layout

packages/ai/
├── package.json              # exports, subpath exports, pi-ai bin
├── src/
│   ├── index.ts              # public entry: re-exports types, stream fns, utils, OAuth types
│   ├── stream.ts             # stream / complete / streamSimple / completeSimple
│   ├── types.ts              # Api, Provider, StreamOptions, Model, Message, AssistantMessageEvent
│   ├── api-registry.ts       # ApiProvider registry: registerApiProvider / getApiProvider
│   ├── models.ts             # runtime model registry, cost calc, thinking-level helpers
│   ├── models.generated.ts   # generated catalog, 1172 models / 31 providers (never edit by hand)
│   ├── env-api-keys.ts       # env key detection, Vertex ADC and Bedrock ambient auth
│   ├── session-resources.ts  # session-scoped cleanup hooks (used by kernel, fork server)
│   ├── cache-pricing.ts      # Anthropic cache read/write cost multipliers
│   ├── log.ts                # injectable structured logger (setLogSink)
│   ├── cli.ts                # pi-ai CLI: OAuth login and provider list
│   ├── oauth.ts              # re-exports utils/oauth
│   ├── mcp.ts                # re-exports mcp/
│   ├── bedrock-provider.ts   # node-only Bedrock module for the Bun override
│   ├── providers/            # one implementation per API, plus shared helpers
│   ├── mcp/                  # MCP catalog + OAuth 2.1 provider
│   └── utils/                # event-stream, stream-failure, overflow, validation, ...
└── scripts/
    └── generate-models.ts    # regenerates models.generated.ts from upstream catalogs

Key abstractions

Type	Path	Description
Model	packages/ai/src/types.ts	Provider-agnostic model descriptor: id, name, api, provider, baseUrl, reasoning, thinkingLevelMap, input modalities, cost, contextWindow, maxTokens, optional compat overrides
StreamOptions	packages/ai/src/types.ts	Options every provider accepts: apiKey, signal, temperature, maxTokens, cacheRetention, sessionId, onPayload, onResponse, headers, timeoutMs, maxRetries, serviceTier
Context	packages/ai/src/types.ts	The request body: systemPrompt, messages, tools
Message	packages/ai/src/types.ts	UserMessage / AssistantMessage / ToolResultMessage union passed into Context
AssistantMessageEvent	packages/ai/src/types.ts	Stream event union: start, text_*, thinking_*, toolcall_* deltas, and the terminal done / error
AssistantMessage	packages/ai/src/types.ts	Terminal result with content blocks, usage, stopReason, errorMessage, and redacted diagnostics
AssistantMessageEventStream	packages/ai/src/utils/event-stream.ts	Push-based AsyncIterable that always terminates with done or error and exposes result()
StreamFunction	packages/ai/src/types.ts	The provider contract: (model, context, options) => AssistantMessageEventStream; failures are encoded in the stream, never thrown
ApiProvider	packages/ai/src/api-registry.ts	Registry record pairing an Api with its stream and streamSimple functions
Tool	packages/ai/src/types.ts	TypeBox schema tool definition; arguments are validated and coerced by validateToolCall (packages/ai/src/utils/validation.ts)
OAuthProviderInterface	packages/ai/src/utils/oauth/types.ts	Login, refresh, and getApiKey contract for subscription providers
McpCatalogEntry	packages/ai/src/mcp/catalog.ts	Built-in MCP server entry (Linear, Notion) with OAuth config

How it works

The stream pipeline is the core of the package. Consumers call stream or streamSimple, the registry resolves the provider for the model's api, the provider module is lazy-loaded, and events are pushed into a single AssistantMessageEventStream. The stream contract (packages/ai/src/types.ts): providers push start before partial updates, then terminate with done carrying the final message, or error carrying a message with stopReason "error" or "aborted" plus errorMessage. result() resolves to that terminal message.

graph TD
    CALL[stream / streamSimple / complete] --> REG[api-registry.ts<br/>getApiProvider by model.api]
    REG -->|provider not loaded yet| LAZY[register-builtins.ts<br/>import provider module on first call]
    LAZY --> PROV[Provider stream function<br/>e.g. anthropic.ts]
    PROV -->|build output message, transform messages| API[Upstream API<br/>SSE / WebSocket / SDK]
    API --> EVTS[Push AssistantMessageEvent<br/>text_delta, thinking_delta, toolcall_delta, done/error]
    EVTS --> STREAM[AssistantMessageEventStream<br/>utils/event-stream.ts]
    STREAM -->|for await| LOOP[Consumer<br/>packages/agent agent-loop]
    STREAM -->|result| FINAL[AssistantMessage<br/>stopReason, usage, errorMessage]
    PROV --> FAIL[utils/stream-failure.ts<br/>classify refusal, safety, rate limit, auth]
    FAIL --> DIAG[AssistantMessageDiagnostic<br/>persisted on the message]

Loading

The registry (packages/ai/src/api-registry.ts) is a Map<string, RegisteredApiProvider>. registerApiProvider wraps the provider's stream and streamSimple with a check that model.api matches the registered api. register-builtins.ts calls it for all nine built-in APIs at module load, but the provider implementations themselves are loaded lazily: each loader holds a module promise created by import("./anthropic.js"), so loading @earendil-works/pi-ai never eagerly pulls in provider SDKs. A module-load failure is converted into an error event stream (createLazyLoadErrorMessage in packages/ai/src/providers/register-builtins.ts) rather than a thrown exception.

Integration points

• packages/agent/src/agent-loop.ts is the primary stream consumer. streamAssistantResponse iterates the AssistantMessageEventStream, forwards text_delta / thinking_delta / toolcall_delta as message_update events, and takes the terminal message from the done / error event.
• packages/coding-agent/src/core/sdk.ts wraps streamSimple with API-key resolution, per-provider retry settings, and extension hooks (onPayload, onResponse, transformContext).
• packages/coding-agent/src/core/model-registry.ts imports getModels / getProviders, registers custom providers and models from ~/.prime/agent/models.json via registerApiProvider, and manages OAuth and MCP providers (resetOAuthProviders, registerBuiltinMcpOAuthProviders).
• packages/coding-agent/src/core/logging.ts installs a setLogSink so pi-ai log entries land in the session JSONL.
• packages/coding-agent/src/core/kernel/index.ts and packages/coding-agent/src/core/kernel/fork-server.ts register registerSessionResourceCleanup hooks that stop when the session ends.
• packages/coding-agent/src/bun/register-bedrock.ts injects the node-only Bedrock implementation via setBedrockProviderModule for Bun runtimes.
• Extensions receive raw assistantMessageEvent deltas through packages/coding-agent/src/core/extensions/types.ts; telemetry (packages/coding-agent/src/core/telemetry.ts) counts text_delta tokens.

Entry points for modification

• Add or change a provider: implement stream / streamSimple in a file under packages/ai/src/providers/, register it lazily in packages/ai/src/providers/register-builtins.ts, add an env key to packages/ai/src/env-api-keys.ts, and a subpath export in packages/ai/package.json. See Providers.
• Add or update models: edit packages/ai/scripts/generate-models.ts, never packages/ai/src/models.generated.ts directly. See Models.
• Change the stream event protocol: update the AssistantMessageEvent union in packages/ai/src/types.ts, AssistantMessageEventStream in packages/ai/src/utils/event-stream.ts, and every provider that emits events.
• Change credential handling: packages/ai/src/env-api-keys.ts for keys, packages/ai/src/utils/oauth/ for OAuth flows, packages/ai/src/mcp/oauth.ts for MCP OAuth.
• Change failure reporting: packages/ai/src/utils/stream-failure.ts (classification and user-facing messages) and packages/ai/src/utils/diagnostics.ts (structured diagnostics on the message).

Key source files

File	Role
packages/ai/src/stream.ts	Public streaming entry points
packages/ai/src/types.ts	All shared types: Model, Message, events, options, compat settings
packages/ai/src/api-registry.ts	Provider registry
packages/ai/src/providers/register-builtins.ts	Lazy provider registration
packages/ai/src/utils/event-stream.ts	AssistantMessageEventStream implementation
packages/ai/src/models.ts	Runtime model registry, cost and thinking-level helpers
packages/ai/src/env-api-keys.ts	Credential detection
packages/ai/src/utils/stream-failure.ts	Failure classification and reporting
packages/ai/src/utils/validation.ts	Tool-call argument validation and coercion

Related pages

• Providers: each provider implementation, auth, and quirks
• Models: the model registry and generation pipeline
• Coding agent overview: the main consumer of the streaming API
• Core agent runtime: the agent loop that iterates the event stream
• Provider login and model selection: user-facing login and model pickers
• MCP client and catalog: MCP servers over pi-ai OAuth
• OAuth login flows: subscription login flows
• Glossary: term definitions
• Patterns and conventions: repo-wide rules, including the models.generated.ts rule