feat(ai): add packages/ai provider layer (Chunk #8)

Self-contained SSE-based provider layer for OpenAI-compatible endpoints
(Ollama/llama.cpp/LM Studio/vLLM), Anthropic, Gemini, and OpenRouter.

- budget.ts: re-export core token-budget engine (countTokens,
  allocateBudget, DEFAULT_BUDGET)
- client.ts: streamCompletion/collectStream/complete with budget
  enforcement via enforceBudget + measureBudget
- discover.ts: auto-detect local providers by probing common /v1/models
- adapters.ts: per-provider param mapping + streaming adapters + factory
- thinking.ts: thinking/reasoning-param normalization per provider
- http.ts: SSE parsing + 429 retry
- errors.ts: BudgetExceededError

35 unit/integration tests pass; typecheck and lint clean.
This commit is contained in:
Schwitzz
2026-10-28 18:59:54 -05:00
parent 5c28a94813
commit 4c4b83bb68
16 changed files with 1795 additions and 0 deletions
+25
View File
@@ -0,0 +1,25 @@
{
"name": "@pubscript/ai",
"version": "0.1.0",
"type": "module",
"main": "./src/index.ts",
"types": "./src/index.ts",
"exports": {
".": "./src/index.ts"
},
"scripts": {
"dev": "echo \"@pubscript/ai has no dev server\" (exit 0)",
"build": "tsc --noEmit",
"test": "vitest run",
"lint": "eslint .",
"typecheck": "tsc --noEmit"
},
"dependencies": {
"@pubscript/core": "workspace:*"
},
"devDependencies": {
"@types/node": "^22.10.2",
"typescript": "^5.7.2",
"vitest": "^2.1.8"
}
}
+162
View File
@@ -0,0 +1,162 @@
import { describe, expect, it } from "vitest";
import {
allocateBudget,
countTokens,
DEFAULT_BUDGET,
enforceBudget,
measureBudget,
} from "./index.js";
import type { Message } from "./index.js";
describe("countTokens", () => {
it("counts tokens for a known string", () => {
// "hello world" is 2 words ≈ 3-4 tokens; assert a stable lower bound.
const n = countTokens("hello world");
expect(n).toBeGreaterThan(0);
expect(Number.isInteger(n)).toBe(true);
});
it("returns 0 for an empty string", () => {
expect(countTokens("")).toBe(0);
});
it("scales with input length", () => {
const short = countTokens("a");
const long = countTokens("a".repeat(1000));
expect(long).toBeGreaterThan(short);
});
});
describe("allocateBudget", () => {
it("splits a 128k context window into per-slot caps", () => {
const budget = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 128_000,
});
expect(budget.system).toBe(12_800);
expect(budget.context).toBe(51_200);
expect(budget.history).toBe(38_400);
expect(budget.output).toBe(25_600);
expect(budget.total).toBe(128_000);
});
it("scales the output budget with context window size", () => {
const small = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 8_000,
});
const large = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 128_000,
});
expect(large.output).toBeGreaterThan(small.output);
});
it("never allocates more than half the context window to output", () => {
const budget = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 16_000,
});
expect(budget.output).toBeLessThanOrEqual(16_000 / 2);
});
});
describe("measureBudget", () => {
it("breaks down tokens by dimension", () => {
const messages: Message[] = [
{ role: "system", content: "You are a helpful editor." },
{ role: "user", content: "Summarize this chapter." },
];
const measured = measureBudget(
messages.find((m) => m.role === "system")!.content,
"retrieved context",
messages.filter((m) => m.role !== "system"),
DEFAULT_BUDGET.output,
);
expect(measured.system).toBeGreaterThan(0);
expect(measured.context).toBeGreaterThan(0);
expect(measured.history).toBeGreaterThan(0);
expect(measured.output).toBe(DEFAULT_BUDGET.output);
});
});
describe("enforceBudget", () => {
it("passes when all dimensions are within budget", () => {
const budget = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 128_000,
});
const measured = {
system: 100,
context: 200,
history: 300,
output: 500,
};
expect(() => enforceBudget(measured, budget)).not.toThrow();
});
it("throws when the context dimension overflows", () => {
const budget = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 8_000,
});
const measured = {
system: 100,
context: budget.context + 1,
history: 100,
output: 100,
};
expect(() => enforceBudget(measured, budget)).toThrow(/context/);
});
it("throws when the total overflows even if each dimension fits", () => {
// Caps sum to more than the total, so each dimension can fit individually
// while the aggregate still exceeds the context window.
const budget = allocateBudget({
system: 30,
context: 30,
history: 30,
output: 30,
total: 1_000,
});
const measured = { system: 300, context: 300, history: 300, output: 300 };
expect(() => enforceBudget(measured, budget)).toThrow(/total/);
});
it("reports all overflowing dimensions", () => {
const budget = allocateBudget({
system: 10,
context: 40,
history: 30,
output: 20,
total: 8_000,
});
const measured = {
system: budget.system + 1,
context: budget.context + 1,
history: 100,
output: 100,
};
expect(() => enforceBudget(measured, budget)).toThrow(/system/);
expect(() => enforceBudget(measured, budget)).toThrow(/context/);
});
});
+15
View File
@@ -0,0 +1,15 @@
/**
* Token budget re-exports.
*
* The budget engine lives in `packages/core` (the single source of truth for
* token accounting). This module re-exports it so the AI layer has a single
* import surface, per CHUNK#8.md §3.
*/
export {
countTokens,
allocateBudget,
DEFAULT_BUDGET,
} from "@pubscript/core";
export type { Budget } from "@pubscript/core";
+201
View File
@@ -0,0 +1,201 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import {
BudgetExceededError,
collectStream,
complete,
streamCompletion,
UnknownProviderError,
} from "./index.js";
import type { ResolvedProvider } from "./index.js";
/** Build a `text/event-stream` Response from a list of SSE event payloads. */
function sseResponse(events: string[], status = 200): Response {
const encoder = new TextEncoder();
const body = new ReadableStream<Uint8Array<ArrayBuffer>>({
start(controller) {
for (const event of events) {
controller.enqueue(encoder.encode(`data: ${event}\n\n`));
}
controller.enqueue(encoder.encode("data: [DONE]\n\n"));
controller.close();
},
});
return new Response(body, {
status,
headers: { "content-type": "text/event-stream" },
});
}
/** Build a JSON Response (for non-streaming completions). */
function jsonResponse(body: unknown, status = 200): Response {
return new Response(JSON.stringify(body), {
status,
headers: { "content-type": "application/json" },
});
}
const config: ResolvedProvider = {
id: "test-openai",
kind: "openai-compatible",
baseUrl: "https://api.openai.com/v1",
apiKey: "sk-test",
model: "gpt-4o",
contextWindow: 128_000,
};
let savedFetch: typeof fetch;
beforeEach(() => {
savedFetch = globalThis.fetch;
});
afterEach(() => {
globalThis.fetch = savedFetch;
vi.restoreAllMocks();
});
describe("streamCompletion", () => {
it("yields text deltas in order", async () => {
globalThis.fetch = (() =>
sseResponse([
JSON.stringify({ choices: [{ delta: { content: "Hello" } }] }),
JSON.stringify({ choices: [{ delta: { content: " world" } }] }),
])) as unknown as typeof fetch;
const chunks: Array<{ type: string; text?: string }> = [];
for await (const chunk of streamCompletion(
config,
[{ role: "user", content: "hi" }],
)) {
chunks.push(chunk);
}
const textChunks = chunks.filter((c) => c.type === "text");
expect(textChunks.map((c) => c.text)).toEqual(["Hello", " world"]);
});
it("parses reasoning deltas separately from text", async () => {
globalThis.fetch = (() =>
sseResponse([
JSON.stringify({ choices: [{ delta: { reasoning_content: "let me think" } }] }),
JSON.stringify({ choices: [{ delta: { content: "the answer" } }] }),
])) as unknown as typeof fetch;
const reasoning: string[] = [];
const text: string[] = [];
for await (const chunk of streamCompletion(
config,
[{ role: "user", content: "hi" }],
)) {
if (chunk.type === "reasoning" && chunk.reasoning) {
reasoning.push(chunk.reasoning);
} else if (chunk.type === "text" && chunk.text) {
text.push(chunk.text);
}
}
expect(reasoning).toEqual(["let me think"]);
expect(text).toEqual(["the answer"]);
});
it("tracks usage on the final done chunk", async () => {
globalThis.fetch = (() =>
sseResponse([
JSON.stringify({ choices: [{ delta: { content: "ok" } }] }),
JSON.stringify({
choices: [{ delta: {} }],
usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 },
}),
])) as unknown as typeof fetch;
const result = await collectStream(
config,
[{ role: "user", content: "hi" }],
);
expect(result.text).toBe("ok");
expect(result.usage).toEqual({ promptTokens: 12, completionTokens: 3, total: 15 });
});
});
describe("collectStream", () => {
it("assembles full text, reasoning, and usage", async () => {
globalThis.fetch = (() =>
sseResponse([
JSON.stringify({ choices: [{ delta: { reasoning_content: "hmm" } }] }),
JSON.stringify({ choices: [{ delta: { content: "Final" } }] }),
JSON.stringify({ choices: [{ delta: { content: " answer" } }] }),
JSON.stringify({
choices: [{ delta: {} }],
usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 },
}),
])) as unknown as typeof fetch;
const result = await collectStream(
config,
[{ role: "user", content: "q" }],
);
expect(result.text).toBe("Final answer");
expect(result.reasoning).toBe("hmm");
expect(result.usage).toEqual({ promptTokens: 5, completionTokens: 2, total: 7 });
});
});
describe("complete (non-streaming)", () => {
it("returns the full text and usage", async () => {
globalThis.fetch = (() =>
jsonResponse({
choices: [{ message: { content: "full text" } }],
usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 },
})) as unknown as typeof fetch;
const result = await complete(
config,
[{ role: "user", content: "hi" }],
);
expect(result.text).toBe("full text");
expect(result.usage).toEqual({ promptTokens: 1, completionTokens: 2, total: 3 });
});
});
describe("retry on 429", () => {
it("retries a rate-limited request and succeeds", async () => {
let calls = 0;
globalThis.fetch = ((() => {
calls += 1;
if (calls === 1) {
return jsonResponse({ error: { message: "rate limited" } }, 429);
}
return sseResponse([JSON.stringify({ choices: [{ delta: { content: "recovered" } }] })]);
})) as unknown as typeof fetch;
const result = await collectStream(
config,
[{ role: "user", content: "hi" }],
);
expect(result.text).toBe("recovered");
expect(calls).toBe(2);
});
});
describe("budget enforcement", () => {
it("throws BudgetExceededError before any network call", async () => {
const tinyConfig: ResolvedProvider = { ...config, contextWindow: 8 };
const fetchSpy = vi.fn();
globalThis.fetch = fetchSpy as unknown as typeof fetch;
await expect(
collectStream(tinyConfig, [{ role: "system", content: "x".repeat(200) }]),
).rejects.toBeInstanceOf(BudgetExceededError);
expect(fetchSpy).not.toHaveBeenCalled();
});
});
describe("provider resolution", () => {
it("throws UnknownProviderError for an unrecognized kind", async () => {
const badConfig = { ...config, kind: "nope" } as unknown as ResolvedProvider;
await expect(
collectStream(badConfig, [{ role: "user", content: "hi" }]),
).rejects.toBeInstanceOf(UnknownProviderError);
});
});
+171
View File
@@ -0,0 +1,171 @@
/**
* The completion client: the public entry point for talking to a provider.
*
* It enforces the token budget (via `packages/core`), resolves the adapter for
* a provider, and exposes:
* - `streamCompletion`: an async iterator over `StreamChunk`s.
* - `collectStream`: collect all chunks + usage into a single result.
* - `complete`: a non-streaming completion (used for tool use).
*
* See CHUNK#8.md §3.
*/
import {
allocateBudget,
countTokens,
DEFAULT_BUDGET,
} from "@pubscript/core";
import { BudgetExceededError } from "./errors.js";
import { getProvider } from "./providers/adapters.js";
import type {
CallOptions,
CompleteResult,
CompletionUsage,
Message,
ProviderAdapter,
ResolvedProvider,
StreamChunk,
} from "./providers/types.js";
/** The measured token breakdown used by the budget engine. */
export interface MeasuredBudget {
system: number;
context: number;
history: number;
output: number;
}
/** The assembled result of a collected stream. */
export interface StreamResult {
text: string;
reasoning?: string;
usage: CompletionUsage;
}
/** Measure the token breakdown of a prompt against the budget engine. */
export function measureBudget(
system: string,
context: string,
history: Message[],
output: number,
): MeasuredBudget {
return {
system: countTokens(system),
context: countTokens(context),
history: countTokens(history.map((m) => `${m.role}: ${m.content}`).join("\n")),
output,
};
}
/**
* Enforce a measured budget against the configured token budget.
*
* @throws {BudgetExceededError} if any dimension overflows.
*/
export function enforceBudget(
measured: MeasuredBudget,
budget: ReturnType<typeof allocateBudget>,
): void {
const overflows: string[] = [];
if (measured.system > budget.system) {
overflows.push(`system (${measured.system}/${budget.system})`);
}
if (measured.context > budget.context) {
overflows.push(`context (${measured.context}/${budget.context})`);
}
if (measured.history > budget.history) {
overflows.push(`history (${measured.history}/${budget.history})`);
}
if (measured.output > budget.output) {
overflows.push(`output (${measured.output}/${budget.output})`);
}
const total =
measured.system + measured.context + measured.history + measured.output;
if (total > budget.total) {
overflows.push(`total (${total}/${budget.total})`);
}
if (overflows.length > 0) {
throw new BudgetExceededError(
`Prompt exceeds context budget; overflowed: ${overflows.join(", ")}`,
);
}
}
/** The output budget to reserve, based on the thinking config. */
function outputBudget(opts: CallOptions): number {
const thinking = opts.thinking;
return typeof thinking?.budget === "number" ? thinking.budget : DEFAULT_BUDGET.output;
}
/**
* Enforce the token budget for a set of messages + context, then resolve the
* provider adapter.
*
* @throws {BudgetExceededError} if the prompt exceeds the budget.
*/
function getAdapter(
config: ResolvedProvider,
messages: Message[],
opts: CallOptions,
): ProviderAdapter {
const output = outputBudget(opts);
const system = messages
.filter((m) => m.role === "system")
.map((m) => m.content)
.join("\n");
const history = messages.filter((m) => m.role !== "system");
const measured = measureBudget(system, opts.contextText ?? "", history, output);
enforceBudget(measured, allocateBudget({ ...DEFAULT_BUDGET, total: config.contextWindow }));
return getProvider(config);
}
/**
* Stream a completion, enforcing the budget first. Yields `StreamChunk`s.
*/
export async function* streamCompletion(
config: ResolvedProvider,
messages: Message[],
opts: CallOptions = {},
): AsyncIterableIterator<StreamChunk> {
const adapter = getAdapter(config, messages, opts);
yield* adapter.stream(messages, opts);
}
/**
* Collect a full completion stream into a single result.
*/
export async function collectStream(
config: ResolvedProvider,
messages: Message[],
opts: CallOptions = {},
): Promise<StreamResult> {
const adapter = getAdapter(config, messages, opts);
let text = "";
let reasoning: string | undefined;
let usage: CompletionUsage = { promptTokens: 0, completionTokens: 0, total: 0 };
for await (const chunk of adapter.stream(messages, opts)) {
if (chunk.type === "text" && chunk.text) {
text += chunk.text;
} else if (chunk.type === "reasoning" && chunk.reasoning) {
reasoning = (reasoning ?? "") + chunk.reasoning;
} else if (chunk.type === "done" && chunk.usage) {
usage = chunk.usage;
}
}
return { text, ...(reasoning ? { reasoning } : {}), usage };
}
/**
* Run a non-streaming completion (used for tool use and offline calls).
*/
export async function complete(
config: ResolvedProvider,
messages: Message[],
opts: CallOptions = {},
): Promise<CompleteResult> {
const adapter = getAdapter(config, messages, opts);
return adapter.complete(messages, opts);
}
+81
View File
@@ -0,0 +1,81 @@
import { describe, expect, it } from "vitest";
import { discoverLocalProviders } from "./index.js";
/** Build a mock fetch that returns a models catalog for a given URL. */
function mockFetch(
catalog: Record<string, unknown>,
): typeof fetch {
return ((url: string) => {
if (catalog[url]) {
return Promise.resolve(
new Response(JSON.stringify(catalog[url]), {
status: 200,
headers: { "content-type": "application/json" },
}),
);
}
return Promise.reject(new Error("not found"));
}) as unknown as typeof fetch;
}
describe("discoverLocalProviders", () => {
it("detects a running Ollama server", async () => {
const results = await discoverLocalProviders({
timeoutMs: 200,
fetchImpl: mockFetch({
"http://localhost:11434/v1/models": {
object: "list",
data: [
{ id: "llama3.1", owned_by: "ollama" },
{ id: "mistral", owned_by: "ollama" },
],
},
}),
});
expect(results).toHaveLength(1);
expect(results[0]!.endpoint).toBe("http://localhost:11434/v1/models");
expect(results[0]!.kind).toBe("openai-compatible");
expect(results[0]!.models.map((m) => m.id)).toEqual(["llama3.1", "mistral"]);
});
it("returns an empty array when no endpoints respond", async () => {
const results = await discoverLocalProviders({
timeoutMs: 100,
fetchImpl: mockFetch({}),
});
expect(results).toEqual([]);
});
it("ignores endpoints that return no models", async () => {
const results = await discoverLocalProviders({
timeoutMs: 200,
fetchImpl: mockFetch({
"http://localhost:1234/v1/models": { object: "list", data: [] },
}),
});
expect(results).toEqual([]);
});
it("probes multiple endpoints and collects all responders", async () => {
const results = await discoverLocalProviders({
timeoutMs: 200,
fetchImpl: mockFetch({
"http://localhost:11434/v1/models": {
object: "list",
data: [{ id: "llama3.1" }],
},
"http://localhost:1234/v1/models": {
object: "list",
data: [{ id: "qwen2.5" }],
},
}),
});
expect(results).toHaveLength(2);
const ids = results.flatMap((r) => r.models.map((m) => m.id));
expect(ids).toContain("llama3.1");
expect(ids).toContain("qwen2.5");
});
});
+97
View File
@@ -0,0 +1,97 @@
/**
* Local provider discovery.
*
* Scans a set of well-known local endpoints for OpenAI-compatible model
* catalogs, so a user can point PubScript at a running Ollama / llama.cpp /
* LM Studio server without hand-entering config. See CHUNK#8.md §3.
*/
import type { ProviderKind } from "./providers/types.js";
/** A model as reported by a provider's `/v1/models` endpoint. */
export interface DetectedModel {
id: string;
owned_by?: string;
}
/** A detected local provider and its available models. */
export interface DetectedProvider {
endpoint: string;
kind: ProviderKind;
models: DetectedModel[];
}
/** A well-known local endpoint to probe. */
interface LocalEndpoint {
url: string;
kind: ProviderKind;
}
/**
* The local endpoints probed during discovery. Each returns an OpenAI-style
* `{ data: [{ id, owned_by }] }` catalog.
*/
export const LOCAL_ENDPOINTS: LocalEndpoint[] = [
{ url: "http://localhost:11434/v1/models", kind: "openai-compatible" }, // Ollama
{ url: "http://localhost:1234/v1/models", kind: "openai-compatible" }, // llama.cpp / LM Studio
{ url: "http://localhost:8080/v1/models", kind: "openai-compatible" }, // generic
];
/** The shape of a provider's `/v1/models` response. */
interface ModelsResponse {
data?: Array<{ id?: string; owned_by?: string }>;
}
/**
* Discover local OpenAI-compatible providers.
*
* @param options.timeoutMs Per-endpoint timeout in milliseconds.
* @param options.fetchImpl Injectable fetch (for testing).
*/
export async function discoverLocalProviders(
options: { timeoutMs?: number; fetchImpl?: typeof fetch } = {},
): Promise<DetectedProvider[]> {
const timeoutMs = options.timeoutMs ?? 800;
const fetchImpl = options.fetchImpl ?? fetch;
const results: DetectedProvider[] = [];
await Promise.all(
LOCAL_ENDPOINTS.map(async (endpoint) => {
try {
const controller = new AbortController();
const timer = setTimeout(() => controller.abort(), timeoutMs);
try {
const response = await fetchImpl(endpoint.url, {
signal: controller.signal,
});
if (!response.ok) {
return;
}
const json = (await response.json()) as ModelsResponse;
const models = (json.data ?? [])
.filter((model) => typeof model.id === "string")
.map((model) => ({
id: model.id as string,
...(typeof model.owned_by === "string"
? { owned_by: model.owned_by }
: {}),
}));
if (models.length > 0) {
results.push({
endpoint: endpoint.url,
kind: endpoint.kind,
models,
});
}
} finally {
clearTimeout(timer);
}
} catch {
// Endpoint unreachable or timed out; skip it.
}
}),
);
return results;
}
+49
View File
@@ -0,0 +1,49 @@
/**
* Error types for the AI provider layer.
*
* Errors are classified so callers can react appropriately: auth failures are
* terminal (the key is wrong), rate-limit/overloaded errors are retryable, and
* budget violations are a client-side guard.
*
* See CHUNK#8.md §3.
*/
/** Base class for all provider-layer errors. */
export class ProviderError extends Error {
constructor(message: string) {
super(message);
this.name = "ProviderError";
}
}
/** The provider rejected the request due to missing/invalid credentials. */
export class ProviderAuthError extends ProviderError {
constructor(message = "Provider authentication failed (invalid API key).") {
super(message);
this.name = "ProviderAuthError";
}
}
/** The provider rate-limited or overloaded the request; safe to retry. */
export class ProviderRateLimitError extends ProviderError {
constructor(message = "Provider rate limit exceeded; retrying.") {
super(message);
this.name = "ProviderRateLimitError";
}
}
/** The request exceeded the configured token budget. */
export class BudgetExceededError extends ProviderError {
constructor(message: string) {
super(message);
this.name = "BudgetExceededError";
}
}
/** No provider adapter is registered for the given kind. */
export class UnknownProviderError extends ProviderError {
constructor(kind: string) {
super(`Unknown provider kind: ${kind}`);
this.name = "UnknownProviderError";
}
}
+170
View File
@@ -0,0 +1,170 @@
/**
* Low-level HTTP + SSE helpers for the provider layer.
*
* This module provides:
* - `readSSEEvents`: parse a `text/event-stream` `ReadableStream` into
* concatenated `data:` payloads (one per event).
* - `fetchWithRetry`: a `fetch` wrapper that retries on 429/5xx with
* exponential backoff and surfaces auth errors immediately.
*
* It is intentionally dependency-free and browser-safe (uses the global
* `fetch`/`TextDecoder`/`ReadableStream`). See CHUNK#8.md §3.
*/
import {
ProviderAuthError,
ProviderError,
ProviderRateLimitError,
} from "./errors.js";
/** Read a `text/event-stream` body and yield each event's `data:` payload. */
export async function* readSSEEvents(
stream: ReadableStream<Uint8Array<ArrayBuffer>>,
): AsyncGenerator<string> {
const reader = stream.getReader();
const decoder = new TextDecoder();
let buffer = "";
const dataLines: string[] = [];
const flush = (): string | null => {
if (dataLines.length === 0) {
return null;
}
const value = dataLines.join("\n");
dataLines.length = 0;
return value;
};
for (;;) {
const { done, value } = await reader.read();
if (value) {
buffer += decoder.decode(value, { stream: true });
}
let newlineIndex: number;
while ((newlineIndex = buffer.indexOf("\n")) >= 0) {
let line = buffer.slice(0, newlineIndex);
buffer = buffer.slice(newlineIndex + 1);
if (line.endsWith("\r")) {
line = line.slice(0, -1);
}
if (line === "") {
// Blank line terminates an SSE event.
const payload = flush();
if (payload !== null) {
yield payload;
}
} else if (line.startsWith("data:")) {
dataLines.push(line.slice(5).replace(/^ /, ""));
}
// Ignore `event:`, `id:`, `retry:` lines.
}
if (done) {
break;
}
}
// Flush any trailing data without a terminating blank line.
if (buffer.length > 0) {
if (buffer.startsWith("data:")) {
dataLines.push(buffer.slice(5).replace(/^ /, ""));
}
const payload = flush();
if (payload !== null) {
yield payload;
}
}
}
/** Parse a JSON SSE event, returning `null` for malformed input. */
export function parseSSEEvent<T = unknown>(event: string): T | null {
try {
return JSON.parse(event) as T;
} catch {
return null;
}
}
/** Sleep for `ms` milliseconds. */
export function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
/** Build the retry backoff delay for a given attempt (0-based). */
function backoffDelay(attempt: number, maxMs: number): number {
const base = 250 * 2 ** attempt;
return Math.min(base, maxMs);
}
/**
* `fetch` wrapper that retries on 429/5xx with exponential backoff. Auth
* errors (401/403) are thrown immediately; other non-retryable errors are
* surfaced as-is.
*/
export async function fetchWithRetry(
url: string,
init: RequestInit,
options: { retries?: number; maxDelayMs?: number; signal?: AbortSignal } = {},
): Promise<Response> {
const retries = options.retries ?? 2;
const maxDelayMs = options.maxDelayMs ?? 5000;
let lastError: unknown;
for (let attempt = 0; attempt <= retries; attempt++) {
if (attempt > 0) {
await sleep(backoffDelay(attempt - 1, maxDelayMs));
}
try {
const response = await fetch(url, init);
if (response.ok) {
return response;
}
const detail = await safeText(response);
if (response.status === 401 || response.status === 403) {
throw new ProviderAuthError(
`Provider rejected credentials (${response.status}).`,
);
}
if (response.status === 429 || response.status >= 500) {
lastError = new ProviderRateLimitError(
`Provider rate-limited or overloaded (${response.status}).`,
);
if (attempt < retries) {
continue;
}
throw lastError;
}
throw new ProviderError(
`Provider request failed (${response.status}): ${detail}`,
);
} catch (error) {
if (error instanceof Error && error.name === "AbortError") {
throw error;
}
lastError = error;
if (attempt >= retries) {
throw error;
}
}
}
// Unreachable, but satisfies the type checker.
throw lastError instanceof Error
? lastError
: new ProviderError("Request failed after retries.");
}
/** Read a response body as text, tolerating errors. */
async function safeText(response: Response): Promise<string> {
try {
return await response.text();
} catch {
return "(no body)";
}
}
+72
View File
@@ -0,0 +1,72 @@
/**
* `@pubscript/ai` — the AI provider layer for PubScript.
*
* A self-contained, SSE-based provider abstraction that talks to any
* OpenAI-compatible endpoint with streaming, thinking/reasoning params, and
* the token-budget engine from `packages/core`. See CHUNK#8.md.
*
* Public API:
* - `streamCompletion`, `collectStream`, `complete` — the completion client.
* - `measureBudget`, `enforceBudget` — budget helpers.
* - `discoverLocalProviders` — local provider discovery.
* - `getProvider` — resolve a provider config into an adapter.
* - `normalizeThinking` — thinking-param mapping.
* - Re-exports of the token budget engine from `packages/core`.
* - All types and error classes.
*/
// Completion client
export {
streamCompletion,
collectStream,
complete,
measureBudget,
enforceBudget,
} from "./client.js";
export type { MeasuredBudget, StreamResult } from "./client.js";
// Provider adapters + factory
export { getProvider } from "./providers/adapters.js";
// Thinking-param mapping
export { normalizeThinking, reasoningBudgetTokens } from "./thinking.js";
// Local provider discovery
export {
discoverLocalProviders,
LOCAL_ENDPOINTS,
} from "./discover.js";
export type { DetectedModel, DetectedProvider } from "./discover.js";
// Token budget engine (re-exported from core)
export {
countTokens,
allocateBudget,
DEFAULT_BUDGET,
} from "./budget.js";
export type { Budget } from "./budget.js";
// Types
export type {
ProviderKind,
ThinkingConfig,
ThinkingEffort,
ProviderConfig,
ResolvedProvider,
Message,
CallOptions,
CompletionUsage,
StreamEventType,
StreamChunk,
CompleteResult,
ProviderAdapter,
} from "./providers/types.js";
// Errors
export {
ProviderError,
ProviderAuthError,
ProviderRateLimitError,
BudgetExceededError,
UnknownProviderError,
} from "./errors.js";
+424
View File
@@ -0,0 +1,424 @@
/**
* Provider adapters: translate the provider-agnostic contract into concrete
* wire formats for OpenAI-compatible, Anthropic, Google (Gemini), and
* OpenRouter endpoints.
*
* Each adapter implements `ProviderAdapter` (see `providers/types.ts`): a
* `stream()` method that yields `StreamChunk`s and a `complete()` method for
* non-streaming calls. See CHUNK#8.md §3.
*/
import { ProviderAuthError, UnknownProviderError } from "../errors.js";
import { fetchWithRetry, readSSEEvents } from "../http.js";
import { normalizeThinking } from "../thinking.js";
import type {
CallOptions,
CompletionUsage,
Message,
ProviderAdapter,
ResolvedProvider,
} from "./types.js";
/** Normalize a raw provider usage object into our canonical shape. */
function normalizeUsage(usage: Record<string, unknown>): CompletionUsage {
const promptTokens = (usage.prompt_tokens as number) ?? 0;
const completionTokens =
(usage.completion_tokens as number) ??
(usage.completion_tokens_details as number) ??
0;
const total =
(usage.total_tokens as number) ??
promptTokens + completionTokens;
return { promptTokens, completionTokens, total };
}
/** Build the auth + content-type headers for a provider. */
function buildHeaders(config: ResolvedProvider): Record<string, string> {
const headers: Record<string, string> = {
"content-type": "application/json",
};
if (config.kind === "anthropic") {
if (config.apiKey) {
headers["x-api-key"] = config.apiKey;
}
if (config.apiVersion) {
headers["anthropic-api-version"] = config.apiVersion;
}
} else if (config.apiKey) {
headers["authorization"] = `Bearer ${config.apiKey}`;
}
if (config.kind === "openrouter") {
headers["HTTP-Referer"] = "https://pubscript.dev";
headers["X-Title"] = "PubScript";
}
if (config.apiVersion && config.kind !== "anthropic") {
headers["openai-api-version"] = config.apiVersion;
}
Object.assign(headers, config.headers);
return headers;
}
/** Join a base URL and a path, avoiding double slashes. */
function joinUrl(baseUrl: string, path: string): string {
return `${baseUrl.replace(/\/+$/, "")}/${path.replace(/^\/+/, "")}`;
}
/** POST a JSON body and return the raw Response. */
async function post(
config: ResolvedProvider,
path: string,
body: Record<string, unknown>,
opts: CallOptions,
): Promise<Response> {
const url = joinUrl(config.baseUrl, path);
return fetchWithRetry(url, {
method: "POST",
headers: buildHeaders(config),
body: JSON.stringify(body),
signal: opts.signal ?? null,
});
}
// ---------------------------------------------------------------------------
// OpenAI-compatible adapter
// ---------------------------------------------------------------------------
function createOpenAICompatibleAdapter(
config: ResolvedProvider,
): ProviderAdapter {
const path = "/chat/completions";
return {
kind: "openai-compatible",
async* stream(messages: Message[], opts: CallOptions) {
const body = {
model: config.model,
messages: messages.map((m) => ({ role: m.role, content: m.content })),
stream: true,
temperature: opts.temperature ?? config.temperature ?? 0.7,
max_tokens: opts.maxTokens ?? config.maxTokens ?? 4096,
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
...opts.extraBody,
};
const response = await post(config, path, body, opts);
for await (const event of readSSEEvents(response.body!)) {
if (event === "[DONE]") {
yield { type: "done" };
return;
}
const json = event ? JSON.parse(event) : null;
if (!json) {
continue;
}
const choice = json.choices?.[0] as
| { delta?: { content?: string; reasoning_content?: string } }
| undefined;
const delta = choice?.delta ?? {};
if (delta.reasoning_content) {
yield { type: "reasoning", reasoning: delta.reasoning_content };
}
if (delta.content) {
yield { type: "text", text: delta.content };
}
if (json.usage) {
yield { type: "done", usage: normalizeUsage(json.usage) };
}
}
},
async complete(messages: Message[], opts: CallOptions) {
const body = {
model: config.model,
messages: messages.map((m) => ({ role: m.role, content: m.content })),
stream: false,
temperature: opts.temperature ?? config.temperature ?? 0.7,
max_tokens: opts.maxTokens ?? config.maxTokens ?? 4096,
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
...opts.extraBody,
};
const response = await post(config, path, body, opts);
const json = (await response.json()) as {
choices?: Array<{ message?: { content?: string } }>;
usage?: Record<string, unknown>;
};
const text = json.choices?.[0]?.message?.content ?? "";
return {
text,
usage: json.usage ? normalizeUsage(json.usage) : { promptTokens: 0, completionTokens: 0, total: 0 },
};
},
};
}
// ---------------------------------------------------------------------------
// Anthropic adapter
// ---------------------------------------------------------------------------
function createAnthropicAdapter(config: ResolvedProvider): ProviderAdapter {
const path = "/v1/messages";
return {
kind: "anthropic",
async* stream(messages: Message[], opts: CallOptions) {
const system = messages
.filter((m) => m.role === "system")
.map((m) => m.content)
.join("\n");
const nonSystem = messages.filter((m) => m.role !== "system");
const body = {
model: config.model,
max_tokens: opts.maxTokens ?? config.maxTokens ?? 4096,
temperature: opts.temperature ?? config.temperature ?? 0.7,
stream: true,
anthropic_version: config.apiVersion ?? "2023-06-01",
...(system ? { system } : {}),
messages: nonSystem.map((m) => ({
role: m.role as "user" | "assistant",
content: m.content,
})),
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
...opts.extraBody,
};
const response = await post(config, path, body, opts);
for await (const event of readSSEEvents(response.body!)) {
if (event === "[DONE]") {
yield { type: "done" };
return;
}
const json = event ? JSON.parse(event) : null;
if (!json) {
continue;
}
if (json.type === "content_block_delta") {
const delta = json.delta as { type?: string; text?: string; thinking?: string };
if (delta.type === "text_delta" && delta.text) {
yield { type: "text", text: delta.text };
} else if (delta.type === "thinking_delta" && delta.thinking) {
yield { type: "reasoning", reasoning: delta.thinking };
} else if (delta.type === "input_json_delta" && delta.thinking) {
yield { type: "reasoning", reasoning: delta.thinking };
}
} else if (json.type === "message_delta") {
const usage = json.delta?.usage as { output_tokens?: number } | undefined;
if (usage?.output_tokens) {
yield {
type: "done",
usage: {
promptTokens: json.usage?.input_tokens ?? 0,
completionTokens: usage.output_tokens,
total: (json.usage?.input_tokens ?? 0) + usage.output_tokens,
},
};
}
}
}
},
async complete(messages: Message[], opts: CallOptions) {
const system = messages
.filter((m) => m.role === "system")
.map((m) => m.content)
.join("\n");
const nonSystem = messages.filter((m) => m.role !== "system");
const body = {
model: config.model,
max_tokens: opts.maxTokens ?? config.maxTokens ?? 4096,
temperature: opts.temperature ?? config.temperature ?? 0.7,
stream: false,
anthropic_version: config.apiVersion ?? "2023-06-01",
...(system ? { system } : {}),
messages: nonSystem.map((m) => ({
role: m.role as "user" | "assistant",
content: m.content,
})),
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
...opts.extraBody,
};
const response = await post(config, path, body, opts);
const json = (await response.json()) as {
content?: Array<{ type?: string; text?: string; thinking?: string }>;
usage?: { input_tokens?: number; output_tokens?: number };
};
const text = (json.content ?? [])
.filter((block) => block.type === "text")
.map((block) => block.text ?? "")
.join("");
const reasoning = (json.content ?? [])
.filter((block) => block.type === "thinking")
.map((block) => block.thinking ?? "")
.join("");
const usage = json.usage
? {
promptTokens: json.usage.input_tokens ?? 0,
completionTokens: json.usage.output_tokens ?? 0,
total: (json.usage.input_tokens ?? 0) + (json.usage.output_tokens ?? 0),
}
: { promptTokens: 0, completionTokens: 0, total: 0 };
return { text, ...(reasoning ? { reasoning } : {}), usage };
},
};
}
// ---------------------------------------------------------------------------
// Google (Gemini) adapter
// ---------------------------------------------------------------------------
/** Shape of a Gemini `streamGenerateContent` response. */
interface GeminiResponse {
candidates?: Array<{
content?: {
parts?: Array<{
text?: string;
thinking?: { reasoning_text?: string };
}>;
};
}>;
usageMetadata?: {
promptTokenCount?: number;
candidatesTokenCount?: number;
totalTokenCount?: number;
};
}
function createGoogleAdapter(config: ResolvedProvider): ProviderAdapter {
const path = "/v1beta/models/chat:streamGenerateContent";
return {
kind: "google",
async* stream(messages: Message[], opts: CallOptions) {
const body = {
contents: messages.map((m) => ({
role: m.role === "system" ? "user" : m.role,
parts: [{ text: m.content }],
})),
generationConfig: {
temperature: opts.temperature ?? config.temperature ?? 0.7,
maxOutputTokens: opts.maxTokens ?? config.maxTokens ?? 4096,
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
},
...opts.extraBody,
};
const response = await post(config, path, body, opts);
for await (const event of readSSEEvents(response.body!)) {
if (event === "[DONE]") {
yield { type: "done" };
return;
}
const json = event ? JSON.parse(event) : null;
if (!json) {
continue;
}
const candidates = json.candidates ?? [];
for (const candidate of candidates) {
const parts = candidate.content?.parts ?? [];
for (const part of parts) {
if (part.text) {
yield { type: "text", text: part.text };
}
if (part.thinking?.reasoning_text) {
yield { type: "reasoning", reasoning: part.thinking.reasoning_text };
}
}
}
if (json.usageMetadata) {
const meta = json.usageMetadata as {
promptTokenCount?: number;
candidatesTokenCount?: number;
totalTokenCount?: number;
};
yield {
type: "done",
usage: {
promptTokens: meta.promptTokenCount ?? 0,
completionTokens: meta.candidatesTokenCount ?? 0,
total: meta.totalTokenCount ?? 0,
},
};
}
}
},
async complete(messages: Message[], opts: CallOptions) {
const body = {
contents: messages.map((m) => ({
role: m.role === "system" ? "user" : m.role,
parts: [{ text: m.content }],
})),
generationConfig: {
temperature: opts.temperature ?? config.temperature ?? 0.7,
maxOutputTokens: opts.maxTokens ?? config.maxTokens ?? 4096,
...normalizeThinking(config.kind, opts.thinking ?? config.thinking),
},
...opts.extraBody,
};
const response = await post(config, path, body, opts);
const json = (await response.json()) as GeminiResponse;
const text = (json.candidates ?? [])
.flatMap((c) => c.content?.parts ?? [])
.filter((p) => p.text)
.map((p) => p.text)
.join("");
const reasoning = (json.candidates ?? [])
.flatMap((c) => c.content?.parts ?? [])
.filter((p) => p.thinking?.reasoning_text)
.map((p) => p.thinking?.reasoning_text)
.join("");
const meta = json.usageMetadata ?? {};
const usage = {
promptTokens: meta.promptTokenCount ?? 0,
completionTokens: meta.candidatesTokenCount ?? 0,
total: meta.totalTokenCount ?? 0,
};
return { text, ...(reasoning ? { reasoning } : {}), usage };
},
};
}
// ---------------------------------------------------------------------------
// OpenRouter adapter (OpenAI-compatible with extra headers)
// ---------------------------------------------------------------------------
function createOpenRouterAdapter(config: ResolvedProvider): ProviderAdapter {
// OpenRouter speaks the OpenAI chat-completions wire format; the only
// differences are the extra referer/title headers, handled in buildHeaders.
return createOpenAICompatibleAdapter({
...config,
baseUrl: config.baseUrl ?? "https://openrouter.ai/api/v1",
});
}
// ---------------------------------------------------------------------------
// Factory
// ---------------------------------------------------------------------------
/**
* Resolve a provider config into a concrete adapter.
*
* @throws {UnknownProviderError} if the kind is not recognized.
*/
export function getProvider(config: ResolvedProvider): ProviderAdapter {
switch (config.kind) {
case "openai-compatible":
return createOpenAICompatibleAdapter(config);
case "anthropic":
return createAnthropicAdapter(config);
case "google":
return createGoogleAdapter(config);
case "openrouter":
return createOpenRouterAdapter(config);
default:
throw new UnknownProviderError((config as { kind: string }).kind);
}
}
/** Re-export for parity with the auth-error surface used by adapters. */
export { ProviderAuthError };
+117
View File
@@ -0,0 +1,117 @@
/**
* Provider-agnostic types for the AI provider layer.
*
* These interfaces describe the *contract* between PubScript and any
* OpenAI-compatible endpoint, independent of the concrete provider. The
* adapters in `providers/adapters.ts` translate these into provider-specific
* wire formats; the client in `client.ts` orchestrates them.
*
* See CHUNK#8.md §3.
*/
/** The family of a provider, which drives the wire-format translation. */
export type ProviderKind =
| "openai-compatible"
| "anthropic"
| "google"
| "openrouter";
/** Reasoning budget expressed either in tokens or at a coarse effort level. */
export type ThinkingEffort = "low" | "medium" | "high";
/**
* Thinking/reasoning configuration. When enabled, the client maps `effort` or
* `budget` onto the provider-specific wire format (see `thinking.ts`).
*/
export interface ThinkingConfig {
enabled: boolean;
/** Reasoning budget in tokens. Maps to `reasoning_effort` when omitted. */
budget?: number;
/** Coarse effort level. Maps to `reasoning_effort` on OpenAI-style APIs. */
effort?: ThinkingEffort;
}
/** A provider configuration as stored in the project store. */
export interface ProviderConfig {
id: string;
kind: ProviderKind;
/** Base URL of the endpoint (e.g. `https://api.openai.com/v1`). */
baseUrl?: string;
/** API key. Never logged or serialized to disk in plaintext. */
apiKey?: string;
/** Model id to request. */
model: string;
temperature?: number;
maxTokens?: number;
thinking?: ThinkingConfig;
/** Context window in tokens, used by the budget engine. */
contextWindow: number;
/** Extra headers forwarded to the provider (e.g. OpenRouter referer). */
headers?: Record<string, string>;
/** `anthropic-api-version` or `openai` API version header. */
apiVersion?: string;
}
/** A provider config with a resolved, non-optional `baseUrl`. */
export interface ResolvedProvider extends ProviderConfig {
baseUrl: string;
}
/** A chat message. Only text content is carried in v1. */
export interface Message {
role: "system" | "user" | "assistant";
content: string;
}
/** Options for a single completion call. */
export interface CallOptions {
temperature?: number;
maxTokens?: number;
thinking?: ThinkingConfig;
signal?: AbortSignal;
/** Extra body params merged into the request (e.g. tools, metadata). */
extraBody?: Record<string, unknown>;
/** Retrieved context (world state) counted separately by the budget engine. */
contextText?: string;
}
/** Token usage for a single completion. */
export interface CompletionUsage {
promptTokens: number;
completionTokens: number;
total: number;
}
/** The kind of event emitted on a completion stream. */
export type StreamEventType = "text" | "reasoning" | "done";
/** A single chunk on a completion stream. */
export interface StreamChunk {
type: StreamEventType;
/** Text delta for `"text"`. */
text?: string;
/** Reasoning delta for `"reasoning"`. */
reasoning?: string;
/** Final usage, carried on the `"done"` chunk. */
usage?: CompletionUsage;
}
/** The fully-assembled result of a non-streaming completion. */
export interface CompleteResult {
text: string;
reasoning?: string;
usage: CompletionUsage;
}
/**
* The interface every provider adapter implements. Adapters translate the
* provider-agnostic `Message[]`/`CallOptions` into the provider's wire format
* and yield `StreamChunk`s.
*/
export interface ProviderAdapter {
readonly kind: ProviderKind;
/** Stream a completion, yielding text/reasoning deltas then a final `done`. */
stream(messages: Message[], opts: CallOptions): AsyncIterableIterator<StreamChunk>;
/** Run a non-streaming completion (used for tool use and offline calls). */
complete(messages: Message[], opts: CallOptions): Promise<CompleteResult>;
}
+97
View File
@@ -0,0 +1,97 @@
import { describe, expect, it } from "vitest";
import { normalizeThinking, reasoningBudgetTokens } from "./index.js";
import type { ThinkingConfig } from "./index.js";
describe("normalizeThinking", () => {
it("returns an empty object when thinking is disabled", () => {
expect(normalizeThinking("openai-compatible", { enabled: false })).toEqual({});
expect(normalizeThinking("openai-compatible", undefined)).toEqual({});
});
describe("openai-compatible", () => {
it("maps effort to reasoning_effort", () => {
const cfg: ThinkingConfig = { enabled: true, effort: "high" };
expect(normalizeThinking("openai-compatible", cfg)).toEqual({
reasoning_effort: "high",
});
});
it("defaults effort to medium", () => {
const cfg: ThinkingConfig = { enabled: true };
expect(normalizeThinking("openai-compatible", cfg)).toEqual({
reasoning_effort: "medium",
});
});
it("adds max_tokens when a budget is set", () => {
const cfg: ThinkingConfig = { enabled: true, budget: 4096 };
expect(normalizeThinking("openai-compatible", cfg)).toEqual({
reasoning_effort: "medium",
max_tokens: 4096,
});
});
});
describe("openrouter", () => {
it("maps effort to reasoning_effort", () => {
const cfg: ThinkingConfig = { enabled: true, effort: "low" };
expect(normalizeThinking("openrouter", cfg)).toEqual({
reasoning_effort: "low",
});
});
});
describe("anthropic", () => {
it("emits a thinking object and thinking_budget", () => {
const cfg: ThinkingConfig = { enabled: true, budget: 4096 };
const result = normalizeThinking("anthropic", cfg);
expect(result).toEqual({
thinking: { enabled: true, budget_tokens: 4096 },
thinking_budget: 4096,
});
});
it("caps the budget at the Anthropic maximum", () => {
const cfg: ThinkingConfig = { enabled: true, budget: 100_000 };
const result = normalizeThinking("anthropic", cfg) as {
thinking: { budget_tokens: number };
};
expect(result.thinking.budget_tokens).toBe(32768);
});
it("defaults the budget when unset", () => {
const cfg: ThinkingConfig = { enabled: true };
const result = normalizeThinking("anthropic", cfg) as {
thinking: { budget_tokens: number };
};
expect(result.thinking.budget_tokens).toBe(2048);
});
});
describe("google", () => {
it("emits reasoning_effort and thinking_budget", () => {
const cfg: ThinkingConfig = { enabled: true, effort: "high", budget: 1024 };
const result = normalizeThinking("google", cfg);
expect(result).toEqual({
reasoning_effort: "high",
thinking_budget: 1024,
});
});
});
});
describe("reasoningBudgetTokens", () => {
it("returns 0 when thinking is disabled", () => {
expect(reasoningBudgetTokens({ enabled: false })).toBe(0);
expect(reasoningBudgetTokens(undefined)).toBe(0);
});
it("returns the configured budget", () => {
expect(reasoningBudgetTokens({ enabled: true, budget: 4096 })).toBe(4096);
});
it("defaults to 2048 when unset", () => {
expect(reasoningBudgetTokens({ enabled: true })).toBe(2048);
});
});
+90
View File
@@ -0,0 +1,90 @@
/**
* Thinking/reasoning parameter mapping.
*
* Different providers express the reasoning budget differently:
*
* - OpenAI-style (`openai-compatible`, `openrouter`): `reasoning_effort`
* (`"low" | "medium" | "high"`), optionally `max_tokens` for the reasoning
* budget.
* - Anthropic: a `thinking` object `{ enabled: boolean, budget_tokens: number }`
* plus `thinking_budget` (must be <= 32768 and > 0).
* - Google (Gemini): `reasoning_effort` + `thinking_budget`.
*
* This module normalizes the provider-agnostic `ThinkingConfig` into the
* provider-specific wire format. See CHUNK#8.md §3.
*/
import type { ProviderKind, ThinkingConfig } from "./providers/types.js";
/** The maximum reasoning budget Anthropic accepts. */
const ANTHROPIC_MAX_THINKING_TOKENS = 32768;
/** Normalize a thinking config into the wire-format body fragment for a kind. */
export function normalizeThinking(
kind: ProviderKind,
thinking?: ThinkingConfig,
): Record<string, unknown> {
if (!thinking?.enabled) {
return {};
}
switch (kind) {
case "anthropic":
return normalizeAnthropicThinking(thinking);
case "google":
return normalizeGoogleThinking(thinking);
case "openai-compatible":
case "openrouter":
default:
return normalizeOpenAIThinking(thinking);
}
}
/** OpenAI-style: `reasoning_effort` + optional `max_tokens`. */
function normalizeOpenAIThinking(thinking: ThinkingConfig): Record<string, unknown> {
const effort = thinking.effort ?? "medium";
const body: Record<string, unknown> = {
reasoning_effort: effort,
};
if (typeof thinking.budget === "number") {
body.max_tokens = thinking.budget;
}
return body;
}
/** Anthropic: `thinking` object + `thinking_budget`. */
function normalizeAnthropicThinking(thinking: ThinkingConfig): Record<string, unknown> {
const budget =
typeof thinking.budget === "number"
? Math.min(thinking.budget, ANTHROPIC_MAX_THINKING_TOKENS)
: 2048;
return {
thinking: {
enabled: true,
budget_tokens: budget,
},
thinking_budget: budget,
};
}
/** Google: `reasoning_effort` + `thinking_budget`. */
function normalizeGoogleThinking(thinking: ThinkingConfig): Record<string, unknown> {
const effort = thinking.effort ?? "medium";
const budget = typeof thinking.budget === "number" ? thinking.budget : 2048;
return {
reasoning_effort: effort,
thinking_budget: budget,
};
}
/**
* Compute the reasoning budget (in tokens) to advertise to the budget engine,
* so the output budget accounts for the reasoning spend. Returns 0 when
* thinking is disabled.
*/
export function reasoningBudgetTokens(thinking?: ThinkingConfig): number {
if (!thinking?.enabled) {
return 0;
}
return typeof thinking.budget === "number" ? thinking.budget : 2048;
}
+8
View File
@@ -0,0 +1,8 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"outDir": "./dist",
"rootDir": "./src"
},
"include": ["src/**/*.ts"]
}
+16
View File
@@ -85,6 +85,22 @@ importers:
specifier: ^2.1.8 specifier: ^2.1.8
version: 2.1.9(@types/node@22.20.5)(jsdom@25.0.1) version: 2.1.9(@types/node@22.20.5)(jsdom@25.0.1)
packages/ai:
dependencies:
'@pubscript/core':
specifier: workspace:*
version: link:../core
devDependencies:
'@types/node':
specifier: ^22.10.2
version: 22.20.5
typescript:
specifier: ^5.7.2
version: 5.9.3
vitest:
specifier: ^2.1.8
version: 2.1.9(@types/node@22.20.5)(jsdom@25.0.1)
packages/core: packages/core:
dependencies: dependencies:
dexie: dexie: