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
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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/);
|
||||
});
|
||||
});
|
||||
@@ -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";
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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";
|
||||
}
|
||||
}
|
||||
@@ -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)";
|
||||
}
|
||||
}
|
||||
@@ -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";
|
||||
@@ -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 };
|
||||
@@ -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>;
|
||||
}
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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;
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
{
|
||||
"extends": "../../tsconfig.base.json",
|
||||
"compilerOptions": {
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src"
|
||||
},
|
||||
"include": ["src/**/*.ts"]
|
||||
}
|
||||
Reference in New Issue
Block a user