diff --git a/examples/ai-sdk-wrapper-usage.ts b/examples/ai-sdk-wrapper-usage.ts new file mode 100644 index 0000000..cd00a12 --- /dev/null +++ b/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); diff --git a/examples/bun.lock b/examples/bun.lock index 0f31057..0cb74e5 100644 --- a/examples/bun.lock +++ b/examples/bun.lock @@ -5,7 +5,11 @@ "": { "name": "@scrawn/sdkexamples", "dependencies": { + "@ai-sdk/google": "^3.0.79", + "@ai-sdk/openai": "^3.0.65", + "@scrawn/analytics": "link:@scrawn/analytics", "@scrawn/core": "link:@scrawn/core", + "ai": "^6.0.191", "dotenv": "^17.2.3", "express": "^4.18.2", }, @@ -16,8 +20,24 @@ }, }, "packages": { + "@ai-sdk/gateway": ["@ai-sdk/gateway@3.0.120", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27", "@vercel/oidc": "3.2.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-MYKAeD2q7/sa1ZdqtL2tw0Me0B8Tok6Q/fhkJDhJl39dG8u+VBlWO9yk9lcdm784bM418o1EKObo4aOxs6+18Q=="], + + "@ai-sdk/google": ["@ai-sdk/google@3.0.79", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-QWVAvYeA7JzEX2wkSyXOWv/I9PD9kvTzdykkSTLi+Eu8RyJ6gA0tdPIGa8esEtOcHE//G5vy6FTB70qQw8l/uw=="], + + "@ai-sdk/openai": ["@ai-sdk/openai@3.0.65", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ZlVoWH+zrdiYDiUt6n/xvfCsk33mzsB81TUQkBRVx79rxU1FKZqVH9J/QCtEpSLqx0cUzjvtIw9l9p7EbUv+dw=="], + + "@ai-sdk/provider": ["@ai-sdk/provider@3.0.10", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-Q3BZ27qfpYqnCYGvE3vt+Qi6LGOF9R5Nmzn+9JoM1lCRsD9mYaIhfJLkSunN48nfGXJ6n+XNV0J/XVpqGQl7Dw=="], + + "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.27", "", { "dependencies": { "@ai-sdk/provider": "3.0.10", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.8" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-ubkAJ+xODouwtmN1tYlvTPphH1hPOBfZaEQe8U7skGvFAnIRs9PPpsq57bC2+Ky/MB4yzhd6YOsxTAx9sGpazw=="], + + "@opentelemetry/api": ["@opentelemetry/api@1.9.1", "", {}, "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q=="], + + "@scrawn/analytics": ["@scrawn/analytics@link:@scrawn/analytics", {}], + "@scrawn/core": ["@scrawn/core@link:@scrawn/core", {}], + "@standard-schema/spec": ["@standard-schema/spec@1.1.0", "", {}, "sha512-l2aFy5jALhniG5HgqrD6jXLi/rUWrKvqN/qJx6yoJsgKhblVd+iqqU4RCXavm/jPityDo5TCvKMnpjKnOriy0w=="], + "@types/body-parser": ["@types/body-parser@1.19.6", "", { "dependencies": { "@types/connect": "*", "@types/node": "*" } }, "sha512-HLFeCYgz89uk22N5Qg3dvGvsv46B8GLvKKo1zKG4NybA8U2DiEO3w9lqGg29t/tfLRJpJ6iQxnVw4OnB7MoM9g=="], "@types/connect": ["@types/connect@3.4.38", "", { "dependencies": { "@types/node": "*" } }, "sha512-K6uROf1LD88uDQqJCktA4yzL1YYAK6NgfsI0v/mTgyPKWsX1CnJ0XPSDhViejru1GcRkLWb8RlzFYJRqGUbaug=="], @@ -40,8 +60,12 @@ "@types/serve-static": ["@types/serve-static@1.15.10", "", { "dependencies": { "@types/http-errors": "*", "@types/node": "*", "@types/send": "<1" } }, "sha512-tRs1dB+g8Itk72rlSI2ZrW6vZg0YrLI81iQSTkMmOqnqCaNr/8Ek4VwWcN5vZgCYWbg/JJSGBlUaYGAOP73qBw=="], + "@vercel/oidc": ["@vercel/oidc@3.2.0", "", {}, "sha512-UycprH3T6n3jH0k44NHMa7pnFHGu/N05MjojYr+Mc6I7obkoLIJujSWwin1pCvdy/eOxrI/l3uDLQsmcrOb4ug=="], + "accepts": ["accepts@1.3.8", "", { "dependencies": { "mime-types": "~2.1.34", "negotiator": "0.6.3" } }, "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw=="], + "ai": ["ai@6.0.191", "", { "dependencies": { "@ai-sdk/gateway": "3.0.120", "@ai-sdk/provider": "3.0.10", "@ai-sdk/provider-utils": "4.0.27", "@opentelemetry/api": "^1.9.0" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-zAxvjKebQE7YkSyyNIl0OM7i6/zygnKeF+yNUjD4nWOelYrG+LpDd6RnH6mjySI4zUpZ7o4wbnmAy8jc6u98vQ=="], + "array-flatten": ["array-flatten@1.1.1", "", {}, "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg=="], "body-parser": ["body-parser@1.20.3", "", { "dependencies": { "bytes": "3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", "destroy": "1.2.0", "http-errors": "2.0.0", "iconv-lite": "0.4.24", "on-finished": "2.4.1", "qs": "6.13.0", "raw-body": "2.5.2", "type-is": "~1.6.18", "unpipe": "1.0.0" } }, "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g=="], @@ -84,6 +108,8 @@ "etag": ["etag@1.8.1", "", {}, "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg=="], + "eventsource-parser": ["eventsource-parser@3.0.8", "", {}, "sha512-70QWGkr4snxr0OXLRWsFLeRBIRPuQOvt4s8QYjmUlmlkyTZkRqS7EDVRZtzU3TiyDbXSzaOeF0XUKy8PchzukQ=="], + "express": ["express@4.21.2", "", { "dependencies": { "accepts": "~1.3.8", "array-flatten": "1.1.1", "body-parser": "1.20.3", "content-disposition": "0.5.4", "content-type": "~1.0.4", "cookie": "0.7.1", "cookie-signature": "1.0.6", "debug": "2.6.9", "depd": "2.0.0", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "etag": "~1.8.1", "finalhandler": "1.3.1", "fresh": "0.5.2", "http-errors": "2.0.0", "merge-descriptors": "1.0.3", "methods": "~1.1.2", "on-finished": "2.4.1", "parseurl": "~1.3.3", "path-to-regexp": "0.1.12", "proxy-addr": "~2.0.7", "qs": "6.13.0", "range-parser": "~1.2.1", "safe-buffer": "5.2.1", "send": "0.19.0", "serve-static": "1.16.2", "setprototypeof": "1.2.0", "statuses": "2.0.1", "type-is": "~1.6.18", "utils-merge": "1.0.1", "vary": "~1.1.2" } }, "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA=="], "finalhandler": ["finalhandler@1.3.1", "", { "dependencies": { "debug": "2.6.9", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "on-finished": "2.4.1", "parseurl": "~1.3.3", "statuses": "2.0.1", "unpipe": "~1.0.0" } }, "sha512-6BN9trH7bp3qvnrRyzsBz+g3lZxTNZTbVO2EV1CS0WIcDbawYVdYvGflME/9QP0h0pYlCDBCTjYa9nZzMDpyxQ=="], @@ -112,6 +138,8 @@ "ipaddr.js": ["ipaddr.js@1.9.1", "", {}, "sha512-0KI/607xoxSToH7GjN1FfSbLoU0+btTicjsQSWQlh/hZykN8KpmMf7uYwPW3R+akZ6R/w18ZlXSHBYXiYUPO3g=="], + "json-schema": ["json-schema@0.4.0", "", {}, "sha512-es94M3nTIfsEPisRafak+HDLfHXnKBhV3vU5eqPcS3flIWqcxJWgXHXiey3YrpaNsanY5ei1VoYEbOzijuq9BA=="], + "math-intrinsics": ["math-intrinsics@1.1.0", "", {}, "sha512-/IXtbwEk5HTPyEwyKX6hGkYXxM9nbj64B+ilVJnC/R6B0pH5G4V3b0pVbL7DBj4tkhBAppbQUlf6F6Xl9LHu1g=="], "media-typer": ["media-typer@0.3.0", "", {}, "sha512-dq+qelQ9akHpcOl/gUVRTxVIOkAJ1wR3QAvb4RsVjS8oVoFjDGTc679wJYmUmknUF5HwMLOgb5O+a3KxfWapPQ=="], @@ -178,6 +206,8 @@ "vary": ["vary@1.1.2", "", {}, "sha512-BNGbWLfd0eUPabhkXUVm0j8uuvREyTh5ovRa/dyow/BqAbZJyC+5fU+IzQOzmAKzYqYRAISoRhdQr3eIZ/PXqg=="], + "zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="], + "@types/serve-static/@types/send": ["@types/send@0.17.6", "", { "dependencies": { "@types/mime": "^1", "@types/node": "*" } }, "sha512-Uqt8rPBE8SY0RK8JB1EzVOIZ32uqy8HwdxCnoCOsYrvnswqmFZ/k+9Ikidlk/ImhsdvBsloHbAlewb2IEBV/Og=="], "send/encodeurl": ["encodeurl@1.0.2", "", {}, "sha512-TPJXq8JqFaVYm2CWmPvnP2Iyo4ZSM7/QKcSmuMLDObfpH5fi7RUGmd/rTDf+rut/saiDiQEeVTNgAmJEdAOx0w=="], diff --git a/examples/package.json b/examples/package.json index 059eaa8..2bcef74 100644 --- a/examples/package.json +++ b/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" } } diff --git a/packages/scrawn/proto b/packages/scrawn/proto index 8e70696..d759346 160000 --- a/packages/scrawn/proto +++ b/packages/scrawn/proto @@ -1 +1 @@ -Subproject commit 8e706964fedccfe107152d2ec0063acae7f2b35d +Subproject commit d759346676331bc1bb15f6d40f6ef3dd1aeb8706 diff --git a/packages/scrawn/src/core/ai/track.ts b/packages/scrawn/src/core/ai/track.ts new file mode 100644 index 0000000..90cdf7c --- /dev/null +++ b/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( + userId: string, + model: ModelInfo, + usage: LanguageModelUsage, + overrides: BillableCallParams, + defaults: { + inputDebit: DebitField; + outputDebit: DebitField; + inputCacheDebit: DebitField; + outputCacheDebit: DebitField; + provider?: string; + } +): AITokenUsagePayload { + 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, + }; +} diff --git a/packages/scrawn/src/core/ai/types.ts b/packages/scrawn/src/core/ai/types.ts new file mode 100644 index 0000000..f07621e --- /dev/null +++ b/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 { + /** Default billing for input tokens (required). */ + inputDebit: DebitField; + /** Default billing for output tokens (required). */ + outputDebit: DebitField; + /** Default billing for cached input tokens. Falls back to inputDebit if not set. */ + inputCacheDebit?: DebitField; + /** Default billing for cached output tokens. Falls back to outputDebit if not set. */ + outputCacheDebit?: DebitField; + /** 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 { + /** The user ID to bill against. If omitted, billing is skipped. */ + userId?: string; + /** Override input token billing for this specific call. */ + inputDebit?: DebitField; + /** Override output token billing for this specific call. */ + outputDebit?: DebitField; + /** Override cached input token billing for this specific call. */ + inputCacheDebit?: DebitField; + /** Override cached output token billing for this specific call. */ + outputCacheDebit?: DebitField; + /** 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> = 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; +}; diff --git a/packages/scrawn/src/core/ai/wrap.ts b/packages/scrawn/src/core/ai/wrap.ts new file mode 100644 index 0000000..509d08d --- /dev/null +++ b/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; +}; + +/** + * 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( + sdk: AISDKModule, + biller: Scrawn, + opts: BillableAIOptions +): Record { + const proxied: Record = { ...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; + const userId = params.userId as string | undefined; + const billing: BillableCallParams = 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; + }) => { + 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( + params: Record +): BillableCallParams { + return { + inputDebit: params.inputDebit as BillableCallParams["inputDebit"], + outputDebit: params.outputDebit as BillableCallParams["outputDebit"], + inputCacheDebit: + params.inputCacheDebit as BillableCallParams["inputCacheDebit"], + outputCacheDebit: + params.outputCacheDebit as BillableCallParams["outputCacheDebit"], + }; +} + +function chainHandlers( + first: (e: { model: ModelInfo; usage: Record }) => void, + second: + | ((e: { model: ModelInfo; usage: Record }) => void) + | undefined +): (e: { model: ModelInfo; usage: Record }) => void { + if (!second) return first; + return (e: { model: ModelInfo; usage: Record }) => { + first(e); + second(e); + }; +} diff --git a/packages/scrawn/src/core/grpc/requestBuilder.ts b/packages/scrawn/src/core/grpc/requestBuilder.ts index e0f4686..88925ec 100644 --- a/packages/scrawn/src/core/grpc/requestBuilder.ts +++ b/packages/scrawn/src/core/grpc/requestBuilder.ts @@ -1,6 +1,7 @@ import * as grpc from "@grpc/grpc-js"; import type { GrpcCallOptions } from "./types.js"; import type { GrpcCallContext } from "./callContext.js"; +import { initClient, buildCallOptions, getRequestMetadata } from "./utils.js"; export class RequestBuilder< C extends { new (...args: any[]): any; serviceName: string } @@ -47,10 +48,7 @@ export class RequestBuilder< this.ctx.logCallStart(); try { - const client = new this.ctx.ClientConstructor( - this.ctx.target, - this.ctx.credentials - ) as grpc.Client & Record; + const client = initClient(this.ctx); const method = client[this.ctx.methodName] as ( request: unknown, metadata: grpc.Metadata, @@ -58,16 +56,13 @@ export class RequestBuilder< callback: (error: grpc.ServiceError | null, response: TResponse) => void ) => void; - const callOptions: grpc.CallOptions = {}; - if (this.options.deadline !== undefined) { - callOptions.deadline = this.options.deadline; - } + const callOptions = buildCallOptions(this.options); const response = await new Promise((resolve, reject) => { method.call( client, this.payload, - this.options.metadata ?? this.ctx.getMetadata(), + getRequestMetadata(this.options, this.ctx), callOptions, (error, response) => { if (error) { diff --git a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts index bac9d2c..0c37129 100644 --- a/packages/scrawn/src/core/grpc/streamRequestBuilder.ts +++ b/packages/scrawn/src/core/grpc/streamRequestBuilder.ts @@ -1,6 +1,7 @@ import * as grpc from "@grpc/grpc-js"; import type { GrpcCallOptions } from "./types.js"; import type { GrpcCallContext } from "./callContext.js"; +import { initClient, buildCallOptions, getRequestMetadata } from "./utils.js"; export class StreamRequestBuilder< C extends { new (...args: any[]): any; serviceName: string } @@ -37,25 +38,19 @@ export class StreamRequestBuilder< this.ctx.logCallStart(); try { - const client = new this.ctx.ClientConstructor( - this.ctx.target, - this.ctx.credentials - ) as grpc.Client & Record; + const client = initClient(this.ctx); const method = client[this.ctx.methodName] as ( metadata: grpc.Metadata, options: grpc.CallOptions, callback: (error: grpc.ServiceError | null, response: TResponse) => void ) => grpc.ClientWritableStream; - const callOptions: grpc.CallOptions = {}; - if (this.options.deadline !== undefined) { - callOptions.deadline = this.options.deadline; - } + const callOptions = buildCallOptions(this.options); const response = await new Promise((resolve, reject) => { const stream = method.call( client, - this.options.metadata ?? this.ctx.getMetadata(), + getRequestMetadata(this.options, this.ctx), callOptions, (error, result) => { if (error) { diff --git a/packages/scrawn/src/core/grpc/utils.ts b/packages/scrawn/src/core/grpc/utils.ts new file mode 100644 index 0000000..01fe71b --- /dev/null +++ b/packages/scrawn/src/core/grpc/utils.ts @@ -0,0 +1,34 @@ +import * as grpc from "@grpc/grpc-js"; +import type { GrpcCallOptions } from "./types.js"; +import type { GrpcCallContext } from "./callContext.js"; + +/** + * Shared client initializer used by both RequestBuilder and StreamRequestBuilder. + */ +export function initClient< + C extends { new (...args: any[]): any; serviceName: string } +>(ctx: GrpcCallContext): grpc.Client & Record { + return new ctx.ClientConstructor(ctx.target, ctx.credentials) as grpc.Client & + Record; +} + +/** + * Shared call-options builder from optional GrpcCallOptions. + */ +export function buildCallOptions(options: GrpcCallOptions): grpc.CallOptions { + const callOptions: grpc.CallOptions = {}; + if (options.deadline !== undefined) { + callOptions.deadline = options.deadline; + } + return callOptions; +} + +/** + * Shared metadata getter — prefers per-request metadata, falls back to context metadata. + */ +export function getRequestMetadata( + options: GrpcCallOptions, + ctx: { getMetadata(): grpc.Metadata } +): grpc.Metadata { + return options.metadata ?? ctx.getMetadata(); +} diff --git a/packages/scrawn/src/core/pricing/types.ts b/packages/scrawn/src/core/pricing/types.ts index cb4bf96..4f7427c 100644 --- a/packages/scrawn/src/core/pricing/types.ts +++ b/packages/scrawn/src/core/pricing/types.ts @@ -19,15 +19,6 @@ */ export type OpType = "ADD" | "SUB" | "MUL" | "DIV"; -/** - * Intellisense hint type for tag names. - * Tag names must be ALL CAPS with underscores only (e.g., PREMIUM_CALL, FEE, INPUT_RATE). - * No lowercase, digits, or hyphens allowed. - * - * This is a branded type that provides IDE hints while remaining compatible with `string`. - */ -export type TagName = Uppercase & { readonly __brand?: "TagName" }; - /** * A literal amount in cents (must be an integer). */ diff --git a/packages/scrawn/src/core/scrawn.ts b/packages/scrawn/src/core/scrawn.ts index 93c8f36..7ed87be 100644 --- a/packages/scrawn/src/core/scrawn.ts +++ b/packages/scrawn/src/core/scrawn.ts @@ -8,6 +8,7 @@ import type { AITokenUsagePayload, EventConsumerErrorCallback, RetryContext, + DebitField, } from "./types/event.js"; import type { AuthRegistry, @@ -59,6 +60,15 @@ import { prettyPrintExpr, tag as _tag, } from "./pricing/index.js"; +import { createBillableAI } from "./ai/wrap.js"; +import type { BillableAIOptions } from "./ai/types.js"; +import type { WithUserId } from "./ai/types.js"; +import { buildAIPayload } from "./ai/track.js"; +import type { + LanguageModelUsage, + ModelInfo, + BillableCallParams, +} from "./ai/types.js"; import { ScrawnConfig } from "../config.js"; import { randomUUID } from "node:crypto"; @@ -144,6 +154,24 @@ export class Scrawn< return error; } + /** Shared: formats Zod issues into a ScrawnValidationError and notifies the callback. */ + private formatValidationError( + message: string, + issues: import("zod").ZodIssue[], + onError?: EventConsumerErrorCallback + ): ScrawnValidationError { + const error = new ScrawnValidationError(message, { + details: { + errors: issues.map((e) => ({ + field: e.path.join("."), + message: e.message, + })), + }, + }); + this.notifyValidationError(error, onError); + return error; + } + /** * Creates a new Scrawn SDK instance. * @@ -388,15 +416,11 @@ export class Scrawn< .map((e) => `${e.path.join(".")}: ${e.message}`) .join(", "); log.error(`Invalid payload for basicUsageEventConsumer: ${errors}`); - const error = new ScrawnValidationError("Payload validation failed", { - details: { - errors: validationResult.error.issues.map((e) => ({ - field: e.path.join("."), - message: e.message, - })), - }, - }); - this.notifyValidationError(error, options?.onError); + this.formatValidationError( + "Payload validation failed", + validationResult.error.issues, + options?.onError + ); return; } @@ -555,15 +579,11 @@ export class Scrawn< log.error( `Invalid payload extracted in middlewareEventConsumer: ${errors}` ); - const error = new ScrawnValidationError("Payload validation failed", { - details: { - errors: validationResult.error.issues.map((e) => ({ - field: e.path.join("."), - message: e.message, - })), - }, - }); - this.notifyValidationError(error, config.onError); + this.formatValidationError( + "Payload validation failed", + validationResult.error.issues, + config.onError + ); return next(); } @@ -931,27 +951,12 @@ export class Scrawn< const responsePromise = (async (): Promise< StreamEventResponse | undefined > => { - try { - log.info("Starting AI token usage stream (return mode)"); - - const response = await this.grpcClient - .newStreamCall(EventServiceClient, "streamEvents") - .addMetadata("authorization", `Bearer ${creds.apiKey}`) - .stream(transformedStream); - - log.info( - `AI token stream completed: ${response.eventsProcessed} events processed` - ); - return response; - } catch (error) { - log.error( - `Failed to stream AI token usage: ${ - error instanceof Error ? error.message : "Unknown error" - }` - ); - this.notifyEventConsumerError(error, onError); - return undefined; - } + const result = await this.performAIStreamCall( + creds.apiKey, + transformedStream, + onError + ); + return result; })(); return { response: responsePromise, stream: userStream }; @@ -960,12 +965,24 @@ export class Scrawn< // Default: fire-and-forget mode const transformedStream = this.transformAITokenStream(stream, onError); + return this.performAIStreamCall(creds.apiKey, transformedStream, onError); + } + + /** + * Shared: performs a gRPC streaming call for AI token events. + * Used by both return-mode and fire-and-forget branches of aiTokenStreamConsumer. + */ + private async performAIStreamCall( + apiKey: string, + transformedStream: AsyncIterable, + onError?: EventConsumerErrorCallback + ): Promise { try { log.info("Starting AI token usage stream"); const response = await this.grpcClient .newStreamCall(EventServiceClient, "streamEvents") - .addMetadata("authorization", `Bearer ${creds.apiKey}`) + .addMetadata("authorization", `Bearer ${apiKey}`) .stream(transformedStream); log.info( @@ -1041,18 +1058,11 @@ export class Scrawn< .map((e) => `${e.path.join(".")}: ${e.message}`) .join(", "); log.error(`Invalid AI token usage payload, skipping: ${errors}`); - const error = new ScrawnValidationError( + this.formatValidationError( "AI token usage payload validation failed", - { - details: { - errors: validationResult.error.issues.map((e) => ({ - field: e.path.join("."), - message: e.message, - })), - }, - } + validationResult.error.issues, + onError ); - this.notifyValidationError(error, onError); continue; } @@ -1180,6 +1190,93 @@ export class Scrawn< yield request; } } + + /** + * Wraps the Vercel AI SDK with automatic per-step billing. + * + * Returns the AI SDK with `streamText`, `generateText`, `streamObject`, + * and `generateObject` patched to accept a `userId` parameter and + * automatically track token usage. All original AI SDK types pass + * through unchanged — returns the same module shape you passed in, + * with billing injected. + * + * User callbacks (`onStepFinish`, `onFinish`) are chained alongside billing. + * + * @param sdk - The Vercel AI SDK module (import * as ai from "ai") + * @param opts - Default billing configuration for all calls + * + * @example + * ```typescript + * import * as ai from "ai"; + * + * const aii = biller.ai(ai, { + * inputDebit: { tag: "AI_INPUT" }, + * outputDebit: { tag: "AI_OUTPUT" }, + * }); + * + * const result = await aii.streamText({ + * userId: "user-123", + * model: openai("gpt-4o-mini"), + * prompt: "Write a story.", + * }); + * // result.text → Promise (preserved from AI SDK) + * ``` + */ + ai>( + sdk: TSDK, + opts: BillableAIOptions + ): WithUserId { + return createBillableAI(sdk, this, opts) as WithUserId; + } + + /** + * Manually track AI token usage from an event callback. + * + * Converts a Vercel AI SDK `onStepFinish` or `onFinish` event into + * an `AITokenUsagePayload` and streams it to the backend (fire-and-forget). + * + * Use this for manual control when you don't want the full `biller.ai()` wrapper. + * + * @param userId - The user ID to bill against + * @param model - Model info (modelId + provider from the event) + * @param usage - Token usage from the event (event.usage or event.totalUsage) + * @param overrides - Override billing per-call (optional) + * @param defaults - Fallback debit config (required — use same as biller.ai opts) + * + * @example + * ```typescript + * const result = await ai.streamText({ + * model: openai("gpt-4o"), + * prompt: "Hello", + * onStepFinish: event => { + * biller.trackAI("user-123", event.model, event.usage, {}, { + * inputDebit: { tag: "AI_INPUT" }, + * outputDebit: { tag: "AI_OUTPUT" }, + * }); + * }, + * }); + * ``` + */ + trackAI( + userId: string, + model: ModelInfo, + usage: LanguageModelUsage, + overrides: BillableCallParams, + defaults: { + inputDebit: DebitField; + outputDebit: DebitField; + inputCacheDebit: DebitField; + outputCacheDebit: DebitField; + provider?: string; + } + ): void { + const payload = buildAIPayload(userId, model, usage, overrides, defaults); + this.aiTokenStreamConsumer( + (async function* () { + yield payload; + })() + ); + } } /** diff --git a/packages/scrawn/src/core/types/event.ts b/packages/scrawn/src/core/types/event.ts index 5c6629a..50e93c8 100644 --- a/packages/scrawn/src/core/types/event.ts +++ b/packages/scrawn/src/core/types/event.ts @@ -16,23 +16,28 @@ const ALL_EXPR_KINDS = [ "exprRef", ]; +/** + * Custom zod schema for PriceExpr validation (allows token placeholders). + * Used in AI token usage payloads where inputTokens()/outputTokens() are valid. + */ +/** + * Shared validator: checks the value is a valid PriceExpr structure (kind exists, expr validates). + * Does NOT check for token placeholders — that's the caller's responsibility. + */ +function isValidPriceExpr(val: unknown): val is PriceExpr { + if (val === null || val === undefined || typeof val !== "object") { + return false; + } + const expr = val as PriceExpr; + return ALL_EXPR_KINDS.includes(expr.kind) && isValidExpr(expr); +} + /** * Custom zod schema for PriceExpr validation (allows token placeholders). * Used in AI token usage payloads where inputTokens()/outputTokens() are valid. */ const PriceExprSchema = z.custom>( - (val): val is PriceExpr => { - if (val === null || val === undefined || typeof val !== "object") { - return false; - } - const expr = val as PriceExpr; - // Check that it has a valid kind (including token placeholders) - if (!ALL_EXPR_KINDS.includes(expr.kind)) { - return false; - } - // Use the validation function - return isValidExpr(expr); - }, + (val): val is PriceExpr => isValidPriceExpr(val), { message: "Must be a valid pricing expression (use tag(), add(), sub(), mul(), div(), amount(), inputTokens(), or outputTokens())", @@ -46,23 +51,8 @@ const PriceExprSchema = z.custom>( */ const PriceExprNoTokensSchema = z.custom>( (val): val is PriceExpr => { - if (val === null || val === undefined || typeof val !== "object") { - return false; - } - const expr = val as PriceExpr; - // Check that it has a valid kind - if (!ALL_EXPR_KINDS.includes(expr.kind)) { - return false; - } - // Use the validation function - if (!isValidExpr(expr)) { - return false; - } - // Reject token placeholders in SDK call context - if (containsTokenExpr(expr)) { - return false; - } - return true; + if (!isValidPriceExpr(val)) return false; + return !containsTokenExpr(val as PriceExpr); }, { message: diff --git a/packages/scrawn/src/index.ts b/packages/scrawn/src/index.ts index ef994f3..f938db8 100644 --- a/packages/scrawn/src/index.ts +++ b/packages/scrawn/src/index.ts @@ -61,3 +61,12 @@ export type { // Export central configuration export { ScrawnConfig, scrawnConfig } from "./config.js"; export type { ScrawnCLIConfig } from "./config.js"; + +// Export AI SDK wrapper types +export type { + BillableAIOptions, + BillableCallParams, + LanguageModelUsage, + ModelInfo, + WithUserId, +} from "./core/ai/types.ts";