PubScript/packages/ai/src/client.ts

/**
 * 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);
}