Feat/ai sdk wrapper

examples/ai-sdk-wrapper-usage.ts

@@ -0,0 +1,22 @@
+import * as ai from "ai";
+import { google } from "@ai-sdk/google";
+import { biller } from "./scrawn/biller.js";
+import { config } from "dotenv";
+config({ path: ".env.local" });
+
+async function main() {
+  const aii = biller.ai(ai, {
+    inputDebit: { tag: "PREMIUM_CALL" },
+    outputDebit: { tag: "EXTRA_FEE" },
+  });
+
+  const result = await aii.streamText({
+    userId: "c0971bcb-b901-4c3e-a191-c9a97871c39f",
+    model: google("gemini-2.5-flash"),
+    prompt: "Write a 2 sentence story about a robot.",
+  });
+
+  console.log(`Generated: "${await result.text}"\n`);
+}
+
+main().catch(console.error);

examples/package.json

@@ -13,9 +13,12 @@
     "@types/node": "^24.10.0"
   },
   "dependencies": {
-    "dotenv": "^17.2.3",
-    "express": "^4.18.2",
+    "@ai-sdk/google": "^3.0.79",
+    "@ai-sdk/openai": "^3.0.65",
+    "@scrawn/analytics": "link:@scrawn/analytics",
     "@scrawn/core": "link:@scrawn/core",
-    "@scrawn/analytics": "link:@scrawn/analytics"
+    "ai": "^6.0.191",
+    "dotenv": "^17.2.3",
+    "express": "^4.18.2"
   }
 }

packages/scrawn/src/core/ai/track.ts

@@ -0,0 +1,38 @@
+import type { AITokenUsagePayload, DebitField } from "../types/event.js";
+import type {
+  BillableCallParams,
+  LanguageModelUsage,
+  ModelInfo,
+} from "./types.js";
+
+/**
+ * Builds an AITokenUsagePayload from an AI SDK step/finish event.
+ * Falls back to regular debit pricing for cache tokens if not specified.
+ */
+export function buildAIPayload<TTag extends string = string>(
+  userId: string,
+  model: ModelInfo,
+  usage: LanguageModelUsage,
+  overrides: BillableCallParams<TTag>,
+  defaults: {
+    inputDebit: DebitField<TTag>;
+    outputDebit: DebitField<TTag>;
+    inputCacheDebit: DebitField<TTag>;
+    outputCacheDebit: DebitField<TTag>;
+    provider?: string;
+  }
+): AITokenUsagePayload<TTag> {
+  return {
+    userId,
+    model: model.modelId,
+    inputTokens: usage.inputTokens,
+    outputTokens: usage.outputTokens,
+    inputDebit: overrides.inputDebit ?? defaults.inputDebit,
+    outputDebit: overrides.outputDebit ?? defaults.outputDebit,
+    provider: overrides.provider ?? defaults.provider ?? model.provider,
+    inputCacheTokens: usage.inputCachedTokens,
+    inputCacheDebit: overrides.inputCacheDebit ?? defaults.inputCacheDebit,
+    outputCacheTokens: usage.outputCachedTokens,
+    outputCacheDebit: overrides.outputCacheDebit ?? defaults.outputCacheDebit,
+  };
+}

packages/scrawn/src/core/ai/types.ts

@@ -0,0 +1,100 @@
+import type { DebitField } from "../types/event.js";
+
+/**
+ * Configuration for the biller.ai() wrapper.
+ * Default billing settings applied automatically to all AI SDK calls.
+ */
+export interface BillableAIOptions<TTag extends string = string> {
+  /** Default billing for input tokens (required). */
+  inputDebit: DebitField<TTag>;
+  /** Default billing for output tokens (required). */
+  outputDebit: DebitField<TTag>;
+  /** Default billing for cached input tokens. Falls back to inputDebit if not set. */
+  inputCacheDebit?: DebitField<TTag>;
+  /** Default billing for cached output tokens. Falls back to outputDebit if not set. */
+  outputCacheDebit?: DebitField<TTag>;
+  /** Default provider override. If not set, auto-detected from the model's provider. */
+  provider?: string;
+}
+
+/**
+ * Additional fields injected by the AI SDK wrapper into function params.
+ * All original AI SDK params are preserved; only userId is added.
+ */
+export interface BillableCallParams<TTag extends string = string> {
+  /** The user ID to bill against. If omitted, billing is skipped. */
+  userId?: string;
+  /** Override input token billing for this specific call. */
+  inputDebit?: DebitField<TTag>;
+  /** Override output token billing for this specific call. */
+  outputDebit?: DebitField<TTag>;
+  /** Override cached input token billing for this specific call. */
+  inputCacheDebit?: DebitField<TTag>;
+  /** Override cached output token billing for this specific call. */
+  outputCacheDebit?: DebitField<TTag>;
+  /** Override provider for this specific call. */
+  provider?: string;
+}
+
+/**
+ * Language model usage as returned by Vercel AI SDK event listeners.
+ * Matches the shape of OnStepFinishEvent.usage and OnFinishEvent.totalUsage.
+ */
+export interface LanguageModelUsage {
+  /** Number of input (prompt) tokens consumed. */
+  inputTokens: number;
+  /** Number of output (completion) tokens consumed. */
+  outputTokens: number;
+  /** Total tokens consumed (inputTokens + outputTokens). */
+  totalTokens: number;
+  /** Cached input tokens (e.g., prompt caching). */
+  inputCachedTokens?: number;
+  /** Cached output tokens. */
+  outputCachedTokens?: number;
+}
+
+/**
+ * Minimal subset of the AI SDK model info needed for billing.
+ * Comes from OnStepFinishEvent.model or OnFinishEvent.model.
+ */
+export interface ModelInfo {
+  /** Model ID, e.g. "gpt-4o-mini". */
+  modelId: string;
+  /** Provider name, e.g. "openai", "anthropic". */
+  provider: string;
+}
+
+// ── Billable AI SDK type mapping ──
+
+/** Keys in the AI SDK that we wrap with billing. */
+type BillableKeys =
+  | "streamText"
+  | "generateText"
+  | "streamObject"
+  | "generateObject";
+
+/**
+ * Given the original AI SDK type TSDK, replaces the wrapped function
+ * signatures with widened versions that accept `userId`. Preserves the
+ * original param types via conditional inference and the original return
+ * types unchanged.
+ *
+ * All other AI SDK properties pass through via Omit.
+ */
+export type WithUserId<TSDK extends Record<string, unknown>> = Omit<
+  TSDK,
+  BillableKeys
+> & {
+  streamText: TSDK extends { streamText: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  generateText: TSDK extends { generateText: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  streamObject: TSDK extends { streamObject: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+  generateObject: TSDK extends { generateObject: (params: infer P) => infer R }
+    ? (params: P & BillableCallParams) => R
+    : never;
+};

packages/scrawn/src/core/ai/wrap.ts

@@ -0,0 +1,136 @@
+import type { Scrawn } from "../scrawn.js";
+import type {
+  BillableAIOptions,
+  BillableCallParams,
+  ModelInfo,
+} from "./types.js";
+
+/** AI SDK function names that accept event callbacks and should be wrapped. */
+const BILLABLE_FNS = [
+  "streamText",
+  "generateText",
+  "streamObject",
+  "generateObject",
+] as const;
+
+type BillableFnName = (typeof BILLABLE_FNS)[number];
+
+/** An AI SDK module shape — duck-typed for flexibility. */
+type AISDKModule = {
+  [K in BillableFnName]?: (...args: any[]) => Promise<unknown>;
+};
+
+/**
+ * Returns a proxied AI SDK module. Each text generation function is wrapped to:
+ * 1. Accept a `userId` field (and optional billing overrides)
+ * 2. Auto-inject `onStepFinish` to track billing on every step
+ * 3. Chain the user's own `onStepFinish`/`onFinish` after billing
+ *
+ * The returned object has the same types as the original AI SDK,
+ * with the billable params injected.
+ */
+export function createBillableAI<TTag extends string>(
+  sdk: AISDKModule,
+  biller: Scrawn<TTag>,
+  opts: BillableAIOptions<TTag>
+): Record<string, unknown> {
+  const proxied: Record<string, unknown> = { ...sdk };
+
+  for (const fnName of BILLABLE_FNS) {
+    const original = sdk[fnName as BillableFnName];
+    if (typeof original !== "function") continue;
+
+    proxied[fnName] = (...args: unknown[]): unknown => {
+      const params = (args[0] ?? {}) as Record<string, unknown>;
+      const userId = params.userId as string | undefined;
+      const billing: BillableCallParams<TTag> = extractBillingParams(params);
+
+      if (userId === undefined || userId.trim() === "") {
+        // No userId — pass through to original unchanged
+        return original.apply(sdk, args);
+      }
+
+      const { onStepFinish: userStep, ...rest } = params;
+      const billingParams = { ...rest };
+
+      const defaults = {
+        inputDebit: opts.inputDebit,
+        outputDebit: opts.outputDebit,
+        inputCacheDebit: opts.inputCacheDebit ?? opts.inputDebit,
+        outputCacheDebit: opts.outputCacheDebit ?? opts.outputDebit,
+        provider: opts.provider,
+      };
+
+      // Inject onStepFinish for per-step billing using biller.trackAI
+      const billingStep = (event: {
+        model: ModelInfo;
+        usage: Record<string, unknown>;
+      }) => {
+        if (!event.usage) return;
+
+        biller.trackAI(
+          userId,
+          {
+            modelId: event.model?.modelId ?? "unknown",
+            provider: event.model?.provider ?? "unknown",
+          },
+          {
+            inputTokens: (event.usage.inputTokens as number) ?? 0,
+            outputTokens: (event.usage.outputTokens as number) ?? 0,
+            totalTokens: (event.usage.totalTokens as number) ?? 0,
+            inputCachedTokens: event.usage.inputCachedTokens as
+              | number
+              | undefined,
+            outputCachedTokens: event.usage.outputCachedTokens as
+              | number
+              | undefined,
+          },
+          billing,
+          defaults
+        );
+      };
+
+      // Chain billing + user callbacks
+      if (
+        typeof userStep === "function" ||
+        userStep === undefined ||
+        userStep === null
+      ) {
+        billingParams.onStepFinish = chainHandlers(
+          billingStep,
+          userStep as ((e: unknown) => void) | undefined
+        );
+      }
+
+      return original.call(sdk, billingParams);
+    };
+  }
+
+  return proxied;
+}
+
+function extractBillingParams<TTag extends string>(
+  params: Record<string, unknown>
+): BillableCallParams<TTag> {
+  return {
+    inputDebit: params.inputDebit as BillableCallParams<TTag>["inputDebit"],
+    outputDebit: params.outputDebit as BillableCallParams<TTag>["outputDebit"],
+    inputCacheDebit:
+      params.inputCacheDebit as BillableCallParams<TTag>["inputCacheDebit"],
+    outputCacheDebit:
+      params.outputCacheDebit as BillableCallParams<TTag>["outputCacheDebit"],
+  };
+}
+
+function chainHandlers(
+  first: (e: { model: ModelInfo; usage: Record<string, unknown> }) => void,
+  second:
+    | ((e: { model: ModelInfo; usage: Record<string, unknown> }) => void)
+    | undefined
+): (e: { model: ModelInfo; usage: Record<string, unknown> }) => void {
+  if (!second) return first;
+  return (e: { model: ModelInfo; usage: Record<string, unknown> }) => {
+    first(e);
+    second(e);
+  };
+}