diff --git a/packages/ai/package.json b/packages/ai/package.json new file mode 100644 index 0000000..42d474a --- /dev/null +++ b/packages/ai/package.json @@ -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" + } +} diff --git a/packages/ai/src/budget.test.ts b/packages/ai/src/budget.test.ts new file mode 100644 index 0000000..1dd4a50 --- /dev/null +++ b/packages/ai/src/budget.test.ts @@ -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/); + }); +}); diff --git a/packages/ai/src/budget.ts b/packages/ai/src/budget.ts new file mode 100644 index 0000000..5ce5528 --- /dev/null +++ b/packages/ai/src/budget.ts @@ -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"; diff --git a/packages/ai/src/client.test.ts b/packages/ai/src/client.test.ts new file mode 100644 index 0000000..880a165 --- /dev/null +++ b/packages/ai/src/client.test.ts @@ -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>({ + 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); + }); +}); diff --git a/packages/ai/src/client.ts b/packages/ai/src/client.ts new file mode 100644 index 0000000..13ca6eb --- /dev/null +++ b/packages/ai/src/client.ts @@ -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, +): 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 { + 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 { + 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 { + const adapter = getAdapter(config, messages, opts); + return adapter.complete(messages, opts); +} diff --git a/packages/ai/src/discover.test.ts b/packages/ai/src/discover.test.ts new file mode 100644 index 0000000..30ecc8a --- /dev/null +++ b/packages/ai/src/discover.test.ts @@ -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, +): 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"); + }); +}); diff --git a/packages/ai/src/discover.ts b/packages/ai/src/discover.ts new file mode 100644 index 0000000..9651183 --- /dev/null +++ b/packages/ai/src/discover.ts @@ -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 { + 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; +} diff --git a/packages/ai/src/errors.ts b/packages/ai/src/errors.ts new file mode 100644 index 0000000..b1a2322 --- /dev/null +++ b/packages/ai/src/errors.ts @@ -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"; + } +} diff --git a/packages/ai/src/http.ts b/packages/ai/src/http.ts new file mode 100644 index 0000000..43c0798 --- /dev/null +++ b/packages/ai/src/http.ts @@ -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>, +): AsyncGenerator { + 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(event: string): T | null { + try { + return JSON.parse(event) as T; + } catch { + return null; + } +} + +/** Sleep for `ms` milliseconds. */ +export function sleep(ms: number): Promise { + 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 { + 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 { + try { + return await response.text(); + } catch { + return "(no body)"; + } +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts new file mode 100644 index 0000000..fd817fd --- /dev/null +++ b/packages/ai/src/index.ts @@ -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"; diff --git a/packages/ai/src/providers/adapters.ts b/packages/ai/src/providers/adapters.ts new file mode 100644 index 0000000..8e8a414 --- /dev/null +++ b/packages/ai/src/providers/adapters.ts @@ -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): 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 { + const headers: Record = { + "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, + opts: CallOptions, +): Promise { + 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; + }; + 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 }; diff --git a/packages/ai/src/providers/types.ts b/packages/ai/src/providers/types.ts new file mode 100644 index 0000000..949e77f --- /dev/null +++ b/packages/ai/src/providers/types.ts @@ -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; + /** `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; + /** 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; + /** Run a non-streaming completion (used for tool use and offline calls). */ + complete(messages: Message[], opts: CallOptions): Promise; +} diff --git a/packages/ai/src/thinking.test.ts b/packages/ai/src/thinking.test.ts new file mode 100644 index 0000000..94c65be --- /dev/null +++ b/packages/ai/src/thinking.test.ts @@ -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); + }); +}); diff --git a/packages/ai/src/thinking.ts b/packages/ai/src/thinking.ts new file mode 100644 index 0000000..6df207e --- /dev/null +++ b/packages/ai/src/thinking.ts @@ -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 { + 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 { + const effort = thinking.effort ?? "medium"; + const body: Record = { + 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 { + 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 { + 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; +} diff --git a/packages/ai/tsconfig.json b/packages/ai/tsconfig.json new file mode 100644 index 0000000..98f36d9 --- /dev/null +++ b/packages/ai/tsconfig.json @@ -0,0 +1,8 @@ +{ + "extends": "../../tsconfig.base.json", + "compilerOptions": { + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src/**/*.ts"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index bb7ba2c..bb231fc 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -85,6 +85,22 @@ importers: specifier: ^2.1.8 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: dependencies: dexie: