@ultimat3/ai 1.0.0

package/src/provider.ts

@@ -0,0 +1,405 @@
+// The `Provider` interface, its two implementations — `AnthropicProvider` (the real Messages API,
+// streaming and not) and `EchoProvider` (deterministic, for tests and `x dev` without a key) —
+// and the money arithmetic over a call's reported usage. The REQUEST half only: ./models owns the
+// catalogue and the per-model rules, ./wire owns the response half.
+
+import type { Money } from '@ultimat3/money';
+import { AiKeyMissingError, AiTransportError } from './errors';
+import type { Effort, ModelId, ThinkingMode } from './models';
+import { DEFAULT_MODEL, MODEL_IDS, MODELS, reasoningBody } from './models';
+import { readSse } from './sse';
+import type { LlmTool, LlmToolCall } from './tools';
+import { MessageStream, parseStopDetails, parseStopReason, parseUsage } from './wire';
+
+export interface AiMessage {
+  readonly role: 'user' | 'assistant';
+  readonly content: string;
+}
+
+export interface GenerateRequest {
+  readonly model?: ModelId;
+  readonly system?: string;
+  readonly messages: readonly AiMessage[];
+  /** Enforced output ceiling. The model never sees it, so it can be cut off mid-answer. */
+  readonly maxTokens: number;
+  readonly effort?: Effort;
+  readonly thinking?: ThinkingMode;
+  readonly tools?: readonly LlmTool[];
+  readonly stopSequences?: readonly string[];
+}
+
+export interface TokenUsage {
+  readonly inputTokens: number;
+  readonly outputTokens: number;
+  readonly cacheReadTokens: number;
+  readonly cacheWriteTokens: number;
+}
+
+/**
+ * Why generation stopped. `refusal` is a successful HTTP response whose content may be
+ * empty or partial — callers must branch on this BEFORE reading `text`.
+ */
+export type StopReason =
+  | 'end_turn'
+  | 'max_tokens'
+  | 'stop_sequence'
+  | 'tool_use'
+  | 'pause_turn'
+  | 'refusal';
+
+/**
+ * Why a refusal happened. Only ever present when `stopReason` is `refusal`, and `category` is
+ * an open set — a closed union here would turn the provider adding a category into a parse
+ * failure. Carried rather than dropped because the category is the one thing that says whether
+ * a different model would answer, which is the only decision a caller can act on.
+ */
+export interface StopDetails {
+  readonly type: 'refusal';
+  readonly category: string | undefined;
+  readonly explanation: string | undefined;
+}
+
+export interface GenerateResult {
+  readonly model: ModelId;
+  readonly text: string;
+  readonly toolCalls: readonly LlmToolCall[];
+  readonly stopReason: StopReason;
+  /** Required, not optional: a provider that forgets it turns a refusal into a silent empty answer. */
+  readonly stopDetails: StopDetails | undefined;
+  readonly usage: TokenUsage;
+  /** Cost of this call, computed from `usage` and the model's prices. */
+  readonly cost: Money;
+}
+
+/** One streamed increment. `done` carries the assembled result. */
+export type StreamChunk =
+  | { readonly type: 'text'; readonly text: string }
+  | { readonly type: 'thinking'; readonly text: string }
+  | { readonly type: 'tool-call'; readonly call: LlmToolCall }
+  | { readonly type: 'done'; readonly result: GenerateResult };
+
+export interface Provider {
+  readonly name: string;
+  /** Models this provider can serve. The gateway routes by membership. */
+  readonly models: readonly ModelId[];
+  generate(request: GenerateRequest): Promise<GenerateResult>;
+  stream(request: GenerateRequest): AsyncIterable<StreamChunk>;
+}
+
+/**
+ * Cost of one call, in integer minor units. Rounded UP: a fraction of a cent that we round
+ * away is money the framework silently absorbs, and under-reporting spend defeats a budget.
+ */
+export function costOf(model: ModelId, usage: TokenUsage): Money {
+  const spec = MODELS[model];
+  // Cache reads are ~0.1x input; cache writes ~1.25x. Scaled by 10 to stay integral.
+  const inputUnits =
+    usage.inputTokens * 10 + usage.cacheReadTokens * 1 + Math.ceil(usage.cacheWriteTokens * 12.5);
+  const inputMinor = divideCeil(inputUnits * spec.inputPerMillion.minor, 10_000_000);
+  const outputMinor = divideCeil(usage.outputTokens * spec.outputPerMillion.minor, 1_000_000);
+  return { minor: inputMinor + outputMinor, currency: spec.inputPerMillion.currency };
+}
+
+function divideCeil(numerator: number, denominator: number): number {
+  return Math.ceil(numerator / denominator);
+}
+
+/** Total tokens a usage record accounts for — what a budget is debited by. */
+export function totalTokens(usage: TokenUsage): number {
+  return usage.inputTokens + usage.outputTokens + usage.cacheReadTokens + usage.cacheWriteTokens;
+}
+
+// ── Anthropic ────────────────────────────────────────────────────────────────
+
+export interface AnthropicProviderInput {
+  /** Reads `ANTHROPIC_API_KEY` when omitted. Absent at call time is a labelled throw. */
+  readonly apiKey?: string;
+  readonly baseUrl?: string;
+  /** Injectable so a test can assert the request body without a network. Defaults to `fetch`. */
+  readonly fetch?: typeof fetch;
+}
+
+const ANTHROPIC_VERSION = '2023-06-01';
+const API_KEY_ENV = 'ANTHROPIC_API_KEY';
+/** Enough of an error body to name the field that was wrong, not enough to fill a log. */
+const DETAIL_LIMIT = 300;
+
+/**
+ * Above this ceiling a non-streaming request sits on an open socket past the HTTP timeout and
+ * fails AFTER the completion was generated and billed. Every current model can be asked for
+ * eight times this, so the limit is the transport's, not the model's.
+ */
+export const STREAM_ONLY_MAX_TOKENS = 16_000;
+
+/** Whether this request has to go over the streaming transport to arrive at all. */
+export function requiresStreaming(request: GenerateRequest): boolean {
+  const model = request.model ?? DEFAULT_MODEL;
+  return Math.min(request.maxTokens, MODELS[model].maxOutput) > STREAM_ONLY_MAX_TOKENS;
+}
+
+/**
+ * The real Messages API shape. The request-surface rules are encoded rather than documented,
+ * because getting them wrong is a 400:
+ *   - `temperature`/`top_p`/`top_k` are REJECTED on every current model. Steer with the prompt.
+ *   - `thinking.budget_tokens` is REJECTED. Use `output_config.effort`.
+ *   - `effort` and adaptive thinking are per-model, and ./models owns which model takes which.
+ */
+export class AnthropicProvider implements Provider {
+  readonly name = 'anthropic';
+  readonly models = MODEL_IDS;
+  private readonly config: AnthropicProviderInput;
+
+  constructor(config: AnthropicProviderInput = {}) {
+    this.config = config;
+  }
+
+  /**
+   * Above `STREAM_ONLY_MAX_TOKENS` this runs the STREAMING transport and returns the assembled
+   * result, rather than refusing. Refusing would leak a transport limit into the API: `llm()`
+   * has no streaming path, so a declaration asking for a legal 64k completion would become
+   * undeclarable instead of merely awkward — and the caller wanting the whole answer at once is
+   * the same caller either way.
+   */
+  async generate(request: GenerateRequest): Promise<GenerateResult> {
+    if (requiresStreaming(request)) return this.assemble(request);
+    const response = await this.send({ ...this.body(request), stream: false });
+    const raw = (await response.json()) as Record<string, unknown>;
+    return parseMessage(request.model ?? DEFAULT_MODEL, raw);
+  }
+
+  /** Drive `stream()` to its `done` chunk. It throws on a cut stream, so a partial never lands. */
+  private async assemble(request: GenerateRequest): Promise<GenerateResult> {
+    for await (const chunk of this.stream(request)) {
+      if (chunk.type === 'done') return chunk.result;
+    }
+    throw new AiTransportError({
+      provider: this.name,
+      detail: 'the stream completed without a result',
+    });
+  }
+
+  /**
+   * Streaming is mandatory above ~16k `maxTokens`: a non-streaming request that large hits the
+   * HTTP timeout before the model finishes. The final `done` chunk carries the assembled
+   * result, so a consumer that only wants the answer can ignore every chunk before it.
+   */
+  async *stream(request: GenerateRequest): AsyncIterable<StreamChunk> {
+    const model = request.model ?? DEFAULT_MODEL;
+    const response = await this.send({ ...this.body(request), stream: true });
+    if (response.body === null) {
+      throw new AiTransportError({
+        provider: this.name,
+        status: response.status,
+        detail: 'a streaming response arrived with no body',
+      });
+    }
+    const message = new MessageStream();
+    for await (const frame of readSse(response.body)) {
+      for (const chunk of message.push(frame)) yield chunk;
+    }
+    // A connection cut mid-answer must fail, not resolve: the partial text reads as a complete
+    // answer, and `end_turn` would be a lie the caller has no way to detect.
+    if (!message.isComplete()) {
+      throw new AiTransportError({
+        provider: this.name,
+        detail: 'the stream ended before message_stop — the answer is truncated',
+      });
+    }
+    const state = message.state();
+    yield {
+      type: 'done',
+      result: { model, ...state, cost: costOf(model, state.usage) },
+    };
+  }
+
+  /** The request body. Pure and side-effect free so a test can assert it directly. */
+  body(request: GenerateRequest): Record<string, unknown> {
+    const model = request.model ?? DEFAULT_MODEL;
+    const body: Record<string, unknown> = {
+      model,
+      max_tokens: Math.min(request.maxTokens, MODELS[model].maxOutput),
+      messages: request.messages.map((m) => ({ role: m.role, content: m.content })),
+      ...reasoningBody(model, request.effort, request.thinking),
+    };
+    if (request.system !== undefined) body['system'] = request.system;
+    if (request.tools !== undefined && request.tools.length > 0) body['tools'] = request.tools;
+    if (request.stopSequences !== undefined) body['stop_sequences'] = request.stopSequences;
+    return body;
+  }
+
+  /**
+   * The one place a request leaves the process. A non-2xx becomes an `AiTransportError`
+   * carrying its status, because the gateway decides whether to retry from that status and a
+   * body parsed as if it were a message would read as an empty, successful answer.
+   */
+  private async send(body: Record<string, unknown>): Promise<Response> {
+    const apiKey = this.config.apiKey ?? Bun.env[API_KEY_ENV];
+    if (apiKey === undefined || apiKey === '') {
+      throw new AiKeyMissingError({ provider: this.name, envVar: API_KEY_ENV });
+    }
+    const doFetch = this.config.fetch ?? fetch;
+    const url = `${this.config.baseUrl ?? 'https://api.anthropic.com'}/v1/messages`;
+    const response = await doFetch(url, {
+      method: 'POST',
+      headers: {
+        'x-api-key': apiKey,
+        'anthropic-version': ANTHROPIC_VERSION,
+        'content-type': 'application/json',
+        accept: body['stream'] === true ? 'text/event-stream' : 'application/json',
+      },
+      body: JSON.stringify(body),
+    });
+    if (!response.ok) {
+      throw new AiTransportError({
+        provider: this.name,
+        status: response.status,
+        detail: await detailOf(response),
+      });
+    }
+    return response;
+  }
+}
+
+/** The provider's own message, when it sent one — it names the offending field, we name the fix. */
+async function detailOf(response: Response): Promise<string> {
+  const body = await response.text().catch(() => '');
+  try {
+    const parsed: unknown = JSON.parse(body);
+    if (typeof parsed === 'object' && parsed !== null) {
+      const error = (parsed as Record<string, unknown>)['error'];
+      if (typeof error === 'object' && error !== null) {
+        const message = (error as Record<string, unknown>)['message'];
+        if (typeof message === 'string') return message.slice(0, DETAIL_LIMIT);
+      }
+    }
+  } catch {
+    // Not JSON — a proxy or a gateway timeout page. The raw text is still the best evidence.
+  }
+  return body === '' ? response.statusText : body.slice(0, DETAIL_LIMIT);
+}
+
+/** Map a Messages API response onto `GenerateResult`. Exported so tests can drive it. */
+export function parseMessage(model: ModelId, raw: Record<string, unknown>): GenerateResult {
+  const content = Array.isArray(raw['content']) ? raw['content'] : [];
+  let text = '';
+  const toolCalls: LlmToolCall[] = [];
+  for (const block of content) {
+    if (typeof block !== 'object' || block === null) continue;
+    const b = block as Record<string, unknown>;
+    if (b['type'] === 'text' && typeof b['text'] === 'string') text += b['text'];
+    if (b['type'] === 'tool_use' && typeof b['id'] === 'string' && typeof b['name'] === 'string') {
+      toolCalls.push({
+        id: b['id'],
+        name: b['name'],
+        input: (b['input'] ?? {}) as Record<string, unknown>,
+      });
+    }
+  }
+  const usage = parseUsage(raw['usage']);
+  return {
+    model,
+    text,
+    toolCalls,
+    stopReason: parseStopReason(raw['stop_reason']),
+    stopDetails: parseStopDetails(raw['stop_details']),
+    usage,
+    cost: costOf(model, usage),
+  };
+}
+
+// ── Echo ─────────────────────────────────────────────────────────────────────
+
+export interface EchoProviderInput {
+  /** Fixed replies keyed by the last user message, for eval fixtures. */
+  readonly replies?: Readonly<Record<string, string>>;
+  /** Fallback when no key matches. Defaults to echoing the last user message. */
+  readonly fallback?: (prompt: string) => string;
+  readonly tokensPerCall?: number;
+}
+
+/**
+ * Deterministic provider. Same input, same output, same usage — which is what makes an eval
+ * suite a test rather than a sample. Token counts are derived from length, so a budget test
+ * can assert a refusal without a network.
+ */
+export class EchoProvider implements Provider {
+  readonly name = 'echo';
+  readonly models = MODEL_IDS;
+  private readonly config: EchoProviderInput;
+
+  constructor(config: EchoProviderInput = {}) {
+    this.config = config;
+  }
+
+  async generate(request: GenerateRequest): Promise<GenerateResult> {
+    const model = request.model ?? DEFAULT_MODEL;
+    const prompt = lastUserMessage(request.messages);
+    const text = this.config.replies?.[prompt] ?? this.config.fallback?.(prompt) ?? prompt;
+    const usage: TokenUsage = {
+      inputTokens: this.config.tokensPerCall ?? estimateTokens(request),
+      outputTokens: estimateTextTokens(text),
+      cacheReadTokens: 0,
+      cacheWriteTokens: 0,
+    };
+    return {
+      model,
+      text,
+      toolCalls: [],
+      stopReason: 'end_turn',
+      stopDetails: undefined,
+      usage,
+      cost: costOf(model, usage),
+    };
+  }
+
+  async *stream(request: GenerateRequest): AsyncIterable<StreamChunk> {
+    const result = await this.generate(request);
+    // One word per chunk: enough to exercise a consumer's assembly logic.
+    for (const word of result.text.split(' ')) {
+      if (word !== '') yield { type: 'text', text: `${word} ` };
+    }
+    yield { type: 'done', result };
+  }
+}
+
+function lastUserMessage(messages: readonly AiMessage[]): string {
+  for (let i = messages.length - 1; i >= 0; i -= 1) {
+    const message = messages[i];
+    if (message !== undefined && message.role === 'user') return message.content;
+  }
+  return '';
+}
+
+/**
+ * ~4 characters per token. Deliberately an ESTIMATE and never used for billing — the
+ * gateway's pre-flight budget check needs a number before the call exists, and the real
+ * count from `usage` replaces it afterwards.
+ */
+export function estimateTokens(request: GenerateRequest): number {
+  return estimateInputTokens(request) + request.maxTokens;
+}
+
+/** The prompt half alone — what the provider bills at the input rate. */
+export function estimateInputTokens(request: GenerateRequest): number {
+  const body = request.messages.map((m) => m.content).join(' ');
+  return estimateTextTokens(body) + estimateTextTokens(request.system ?? '');
+}
+
+/**
+ * Worst-case price of a request before it exists: the prompt at the input rate, the FULL
+ * `maxTokens` at the output rate. Deliberately pessimistic — a ceiling checked against an
+ * optimistic estimate is a ceiling one long completion walks through.
+ */
+export function estimateCost(request: GenerateRequest): Money {
+  const model = request.model ?? DEFAULT_MODEL;
+  return costOf(model, {
+    inputTokens: estimateInputTokens(request),
+    outputTokens: Math.min(request.maxTokens, MODELS[model].maxOutput),
+    cacheReadTokens: 0,
+    cacheWriteTokens: 0,
+  });
+}
+
+export function estimateTextTokens(text: string): number {
+  return Math.ceil(text.length / 4);
+}

package/src/rag.ts

@@ -0,0 +1,186 @@
+// The retrieval pipeline: chunk → embed → retrieve → rerank → assemble under a token budget.
+//
+// The budget is the point. Retrieval that returns "the top 20 chunks" and lets the caller
+// discover at request time that they don't fit is a truncation bug waiting to happen; the
+// assembler fills to a declared ceiling and reports what it dropped.
+
+import type { Embedder } from './embeddings';
+import { embedBatched, embedOne } from './embeddings';
+import { estimateTextTokens as estimateChunkTokens } from './provider';
+import type { SearchHit, VectorStore } from './vector';
+
+export interface Chunk {
+  readonly id: string;
+  readonly text: string;
+  readonly tokens: number;
+  readonly metadata: Readonly<Record<string, string>>;
+}
+
+export interface ChunkInput {
+  readonly id: string;
+  readonly text: string;
+  readonly metadata?: Readonly<Record<string, string>>;
+  /** Target chunk size in tokens. */
+  readonly size?: number;
+  /** Overlap in tokens. Without it a fact split across a boundary is retrievable by neither. */
+  readonly overlap?: number;
+}
+
+/**
+ * Token-aware chunker. Splits on paragraph, then sentence, then hard-wraps — so a chunk
+ * boundary lands at a meaning boundary whenever one is available within the budget.
+ */
+export function chunk(input: ChunkInput): readonly Chunk[] {
+  const size = input.size ?? 512;
+  const overlap = Math.min(input.overlap ?? 64, size - 1);
+  const units = splitUnits(input.text);
+  const chunks: Chunk[] = [];
+  let buffer: string[] = [];
+  let tokens = 0;
+
+  const flush = (): void => {
+    if (buffer.length === 0) return;
+    const text = buffer.join(' ').trim();
+    if (text !== '') {
+      chunks.push({
+        id: `${input.id}#${chunks.length}`,
+        text,
+        tokens: estimateChunkTokens(text),
+        metadata: { source: input.id, ...(input.metadata ?? {}) },
+      });
+    }
+    // Carry the tail forward as the overlap for the next chunk.
+    const carried: string[] = [];
+    let carriedTokens = 0;
+    for (let i = buffer.length - 1; i >= 0 && carriedTokens < overlap; i -= 1) {
+      const unit = buffer[i] ?? '';
+      carried.unshift(unit);
+      carriedTokens += estimateChunkTokens(unit);
+    }
+    buffer = carried;
+    tokens = carriedTokens;
+  };
+
+  for (const unit of units) {
+    const unitTokens = estimateChunkTokens(unit);
+    if (tokens + unitTokens > size && buffer.length > 0) flush();
+    buffer.push(unit);
+    tokens += unitTokens;
+  }
+  flush();
+  // The final flush leaves the overlap tail in the buffer; it is already covered.
+  return chunks;
+}
+
+function splitUnits(text: string): readonly string[] {
+  const units: string[] = [];
+  for (const paragraph of text.split(/\n{2,}/)) {
+    const trimmed = paragraph.trim();
+    if (trimmed === '') continue;
+    // Sentence split that keeps the terminator attached.
+    const sentences = trimmed.match(/[^.!?]+[.!?]*\s*/g) ?? [trimmed];
+    for (const sentence of sentences) {
+      const s = sentence.trim();
+      if (s !== '') units.push(s);
+    }
+  }
+  return units;
+}
+
+/** Index a document: chunk, embed, upsert. One call so no step is skipped by accident. */
+export async function indexDocument(input: {
+  readonly store: VectorStore;
+  readonly embedder: Embedder;
+  readonly document: ChunkInput;
+}): Promise<readonly Chunk[]> {
+  const chunks = chunk(input.document);
+  const vectors = await embedBatched(
+    input.embedder,
+    chunks.map((c) => c.text),
+  );
+  await input.store.upsert(
+    chunks.map((c, index) => ({
+      id: c.id,
+      vector: vectors[index] as Float32Array,
+      text: c.text,
+      metadata: c.metadata,
+    })),
+  );
+  return chunks;
+}
+
+/** A reranker reorders candidates with a stronger (and slower) signal than retrieval. */
+export interface Reranker {
+  readonly name: string;
+  rerank(input: {
+    readonly query: string;
+    readonly hits: readonly SearchHit[];
+    readonly k: number;
+  }): Promise<readonly SearchHit[]>;
+}
+
+/** Identity reranker — the honest default when no cross-encoder is configured. */
+export const passthroughReranker: Reranker = {
+  name: 'passthrough',
+  rerank: async ({ hits, k }) => hits.slice(0, k),
+};
+
+export interface RetrieveInput {
+  readonly store: VectorStore;
+  readonly embedder: Embedder;
+  readonly query: string;
+  readonly k?: number;
+  readonly filter?: Readonly<Record<string, string>>;
+  readonly reranker?: Reranker;
+}
+
+/** Hybrid retrieval then rerank. Hybrid by default because pure vector loses on exact terms. */
+export async function retrieve(input: RetrieveInput): Promise<readonly SearchHit[]> {
+  const k = input.k ?? 8;
+  const vector = await embedOne(input.embedder, input.query);
+  const hits = await input.store.hybrid({
+    query: input.query,
+    vector,
+    k: k * 3,
+    ...(input.filter !== undefined ? { filter: input.filter } : {}),
+  });
+  const reranker = input.reranker ?? passthroughReranker;
+  return reranker.rerank({ query: input.query, hits, k });
+}
+
+export interface AssembledContext {
+  readonly text: string;
+  readonly tokens: number;
+  readonly used: readonly string[];
+  /** Ids that did not fit. Reported, never silently dropped. */
+  readonly dropped: readonly string[];
+}
+
+/**
+ * Fill a context window in rank order until the budget is reached. Skips oversized hits
+ * rather than stopping, so one long chunk does not starve every shorter one behind it.
+ */
+export function assembleContext(input: {
+  readonly hits: readonly SearchHit[];
+  readonly maxTokens: number;
+  readonly separator?: string;
+}): AssembledContext {
+  const separator = input.separator ?? '\n\n---\n\n';
+  const separatorTokens = estimateChunkTokens(separator);
+  const parts: string[] = [];
+  const used: string[] = [];
+  const dropped: string[] = [];
+  let tokens = 0;
+
+  for (const hit of input.hits) {
+    const cost = estimateChunkTokens(hit.text) + (parts.length === 0 ? 0 : separatorTokens);
+    if (tokens + cost > input.maxTokens) {
+      dropped.push(hit.id);
+      continue;
+    }
+    parts.push(hit.text);
+    used.push(hit.id);
+    tokens += cost;
+  }
+  return { text: parts.join(separator), tokens, used, dropped };
+}