diff --git a/.gitignore b/.gitignore index 58a0b8b47..9e04e6a10 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,7 @@ supabase.sh tests/ !server/tests !packages/mcp/tests +!packages/ai-sdk/tests !apps/leaf/tests !vite/tests .secrets diff --git a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx index bb33e85f0..e9bb2b92a 100644 --- a/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx +++ b/apps/docs/mintlify/documentation/external-providers/ai-sdk.mdx @@ -33,16 +33,16 @@ bun add @useautumn/ai-sdk #### 2. Wrap your model -Use `withTokenTracking` to wrap any AI SDK language model. It intercepts generate and stream calls, reads the token usage from the response, and reports it to Autumn automatically. +Use `withAutumn` to wrap any AI SDK language model. It intercepts generate and stream calls, reads the token usage from the response, and reports it to Autumn automatically. ```typescript import { Autumn } from "autumn-js"; import { anthropic } from "@ai-sdk/anthropic"; -import { withTokenTracking } from "@useautumn/ai-sdk"; +import { withAutumn } from "@useautumn/ai-sdk"; const autumn = new Autumn({ secretKey: "am_sk_test_1234" }); -const model = withTokenTracking({ +const model = withAutumn({ autumn, model: anthropic("claude-sonnet-4-5-20250514"), customerId: "user_123", @@ -89,7 +89,7 @@ import { createOpenRouter } from "@openrouter/ai-sdk-provider"; const openrouter = createOpenRouter(); -const model = withTokenTracking({ +const model = withAutumn({ autumn, model: openrouter("anthropic/claude-opus-4-6"), customerId: "user_123", @@ -116,12 +116,12 @@ const model = withTokenTracking({ import { Autumn } from "autumn-js"; import { openai } from "@ai-sdk/openai"; import { generateText } from "ai"; -import { withTokenTracking } from "@useautumn/ai-sdk"; +import { withAutumn } from "@useautumn/ai-sdk"; const autumn = new Autumn({ secretKey: process.env.AUTUMN_SECRET_KEY! }); async function chat(customerId: string, message: string) { - const model = withTokenTracking({ + const model = withAutumn({ autumn, model: openai("gpt-4o"), customerId, diff --git a/bun.lock b/bun.lock index f8d001f5a..6d4d8dcdf 100644 --- a/bun.lock +++ b/bun.lock @@ -1,5 +1,6 @@ { "lockfileVersion": 1, + "configVersion": 0, "workspaces": { "": { "name": "autumn", @@ -7099,6 +7100,8 @@ "@typespec/ts-http-runtime/https-proxy-agent": ["https-proxy-agent@7.0.6", "", { "dependencies": { "agent-base": "^7.1.2", "debug": "4" } }, "sha512-vK9P5/iUfdl95AI+JVyUuIcVtd4ofvtrOr3HNtM2yxC9bnMbEdp3x01OhQNnjb8IJYi38VlTE3mBXwcfvywuSw=="], + "@useautumn/ai-sdk/typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="], + "@useautumn/sdk/zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="], "@vercel/sdk/zod": ["zod@4.4.3", "", {}, "sha512-ytENFjIJFl2UwYglde2jchW2Hwm4GJFLDiSXWdTrJQBIN9Fcyp7n4DhxJEiWNAJMV1/BqWfW/kkg71UDcHJyTQ=="], diff --git a/packages/ai-sdk/package.json b/packages/ai-sdk/package.json index f8adf149d..e7f964477 100644 --- a/packages/ai-sdk/package.json +++ b/packages/ai-sdk/package.json @@ -20,6 +20,7 @@ }, "scripts": { "ts": "tsgo --noEmit --skipLibCheck", + "test": "bun test tests/unit", "build": "rm -rf dist && tsup", "prepublishOnly": "bun run build" }, diff --git a/packages/ai-sdk/src/index.ts b/packages/ai-sdk/src/index.ts index ce2f5a9ff..e655582ff 100644 --- a/packages/ai-sdk/src/index.ts +++ b/packages/ai-sdk/src/index.ts @@ -4,6 +4,7 @@ import { type LanguageModelUsage, wrapLanguageModel, } from "ai"; +// @ts-expect-error autumn-js types resolve in consuming projects; this package only needs the peer type. import type { Autumn } from "autumn-js"; // Standalone published package: must not import from the internal @autumn/shared workspace. @@ -143,7 +144,7 @@ export const withAutumn = ({ const trackUsage = async (usage: AnyUsage) => { try { const pools = normalizeUsage(usage); - // @ts-expect-error trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. + // @ts-ignore trackTokens is generated from OpenAPI; local autumn-js types may not include it yet. await autumn.balances.trackTokens({ customerId, modelId: modelName, diff --git a/packages/ai-sdk/tests/unit/index.test.ts b/packages/ai-sdk/tests/unit/index.test.ts new file mode 100644 index 000000000..fd28f5863 --- /dev/null +++ b/packages/ai-sdk/tests/unit/index.test.ts @@ -0,0 +1,144 @@ +import { describe, expect, test } from "bun:test"; +import type { LanguageModelV3, LanguageModelV3Usage } from "@ai-sdk/provider"; +import { generateText, streamText } from "ai"; +import { withAutumn } from "../../src/index.js"; + +type TrackTokensParams = { + customerId: string; + modelId: string; + inputTokens: number; + outputTokens: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + reasoningTokens?: number; + featureId?: string; + entityId?: string; + properties?: Record; +}; + +const usage: LanguageModelV3Usage = { + inputTokens: { + total: 13, + noCache: 10, + cacheRead: 2, + cacheWrite: 1, + }, + outputTokens: { + total: 7, + text: 5, + reasoning: 2, + }, +}; + +const finishReason = { unified: "stop" as const, raw: "stop" }; + +const createAutumn = () => { + const calls: TrackTokensParams[] = []; + + return { + calls, + autumn: { + balances: { + trackTokens: async (params: TrackTokensParams) => { + calls.push(params); + }, + }, + }, + }; +}; + +const createModel = (): LanguageModelV3 => ({ + specificationVersion: "v3", + provider: "openai", + modelId: "gpt-test", + supportedUrls: {}, + async doGenerate() { + return { + content: [{ type: "text", text: "hello" }], + finishReason, + usage, + warnings: [], + }; + }, + async doStream() { + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "text-start", id: "text-1" }); + controller.enqueue({ + type: "text-delta", + id: "text-1", + delta: "hello", + }); + controller.enqueue({ type: "text-end", id: "text-1" }); + controller.enqueue({ type: "finish", finishReason, usage }); + controller.close(); + }, + }), + }; + }, +}); + +describe("withAutumn", () => { + test("tracks token usage from generateText", async () => { + const { autumn, calls } = createAutumn(); + + const model = withAutumn({ + autumn, + model: createModel(), + customerId: "cus_test", + featureId: "ai_credits", + entityId: "entity_test", + properties: { source: "test" }, + }); + + const result = await generateText({ model, prompt: "Say hello" }); + + expect(result.text).toBe("hello"); + expect(calls).toEqual([ + { + customerId: "cus_test", + modelId: "openai/gpt-test", + inputTokens: 10, + outputTokens: 5, + cacheReadTokens: 2, + cacheWriteTokens: 1, + reasoningTokens: 2, + featureId: "ai_credits", + entityId: "entity_test", + properties: { source: "test" }, + }, + ]); + }); + + test("tracks token usage from streamText when the stream finishes", async () => { + const { autumn, calls } = createAutumn(); + + const model = withAutumn({ + autumn, + model: createModel(), + customerId: "cus_stream", + providerId: "custom-openai", + }); + + const result = streamText({ model, prompt: "Say hello" }); + const chunks: string[] = []; + + for await (const chunk of result.textStream) { + chunks.push(chunk); + } + + expect(chunks.join("")).toBe("hello"); + expect(calls).toEqual([ + { + customerId: "cus_stream", + modelId: "custom-openai/gpt-test", + inputTokens: 10, + outputTokens: 5, + cacheReadTokens: 2, + cacheWriteTokens: 1, + reasoningTokens: 2, + }, + ]); + }); +}); diff --git a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts index a67702701..2c53db6f5 100644 --- a/packages/atmn/src/lib/transforms/sdkToApi/feature.ts +++ b/packages/atmn/src/lib/transforms/sdkToApi/feature.ts @@ -12,7 +12,7 @@ export interface ApiFeatureParams { credit_cost: number; }>; model_markups?: Record; diff --git a/packages/openapi/tsconfig.json b/packages/openapi/tsconfig.json index 6d08724fc..fcbd4453d 100644 --- a/packages/openapi/tsconfig.json +++ b/packages/openapi/tsconfig.json @@ -7,6 +7,7 @@ "moduleResolution": "bundler", "target": "ES2020", "noEmit": true, + "types": ["node", "bun"], "paths": { "@autumn/shared": ["../../shared/index.ts"], "@api/*": ["../../shared/api/*"], @@ -15,6 +16,5 @@ } }, "include": ["./**/*"], - "types": ["node"], "exclude": ["node_modules", "dist"] } diff --git a/server/src/init.ts b/server/src/init.ts index 707bd94f9..254807594 100644 --- a/server/src/init.ts +++ b/server/src/init.ts @@ -45,7 +45,6 @@ import { import { preWarmOrgRedisConnections } from "./external/redis/orgRedisPool.js"; import { createHonoApp } from "./initHono.js"; import { otelSdk } from "./instrumentation.js"; -import { initializeDatabaseFunctions } from "./db/initializeDatabaseFunctions.js"; import { checkEnvVars } from "./utils/initUtils.js"; import { startMemoryMonitor } from "./utils/memoryMonitor.js"; @@ -67,7 +66,6 @@ const init = async ({ startupStartedAt }: { startupStartedAt: number }) => { void preWarmOrgRedisConnections({ db }).catch((error) => { logger.warn("[OrgRedis] Warmup failed", { error }); }); - await initializeDatabaseFunctions(); await startAllEdgeConfigPolling({ logger }); await Promise.all([primeRedisMonitor(), primeRedisV2Monitor()]); diff --git a/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts b/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts index 81a9e2d3f..a58047276 100644 --- a/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts +++ b/shared/utils/productUtils/priceUtils/convertPrice/priceToStripePrepaidV2Tiers.ts @@ -1,11 +1,10 @@ import type { Organization } from "@models/orgModels/orgTable"; import type { Entitlement } from "@models/productModels/entModels/entModels"; -import type { UsagePriceConfig } from "@models/productModels/priceModels/priceConfig/usagePriceConfig"; import type { Price } from "@models/productModels/priceModels/priceModels"; import { orgToCurrency } from "@utils/orgUtils/convertOrgUtils"; import { - isFinalTier, isNotFinalTier, + isPrepaidPrice, } from "@utils/productUtils/priceUtils/classifyPriceUtils"; import { atmnToStripeAmountDecimal } from "@utils/productUtils/priceUtils/convertAmountUtils"; import { Decimal } from "decimal.js"; @@ -33,8 +32,13 @@ export const priceToStripePrepaidV2Tiers = ({ price: Price; entitlement: Entitlement; org: Organization; -}) => { - const config = price.config as UsagePriceConfig; +}): Stripe.PriceCreateParams.Tier[] => { + if (!isPrepaidPrice(price)) { + throw new Error( + `priceToStripePrepaidV2Tiers requires a prepaid price, got price ${price.id}`, + ); + } + const config = price.config; const tiers: Stripe.PriceCreateParams.Tier[] = []; @@ -47,9 +51,8 @@ export const priceToStripePrepaidV2Tiers = ({ }); } - for (let i = 0; i < config.usage_tiers.length; i++) { - const tier = config.usage_tiers[i]; - const atmnUnitAmount = new Decimal(tier.amount).div( + for (const tier of config.usage_tiers) { + const atmnUnitAmount = new Decimal(tier.amount ?? 0).div( config.billing_units ?? 1, ); @@ -58,14 +61,14 @@ export const priceToStripePrepaidV2Tiers = ({ currency: orgToCurrency({ org }), }); - let upTo = tier.to; - if (isNotFinalTier(tier) && entitlement.allowance) { - upTo = tier.to + entitlement.allowance; + let upTo: Stripe.PriceCreateParams.Tier["up_to"] = "inf"; + if (isNotFinalTier(tier)) { + upTo = entitlement.allowance ? tier.to + entitlement.allowance : tier.to; } const stripeTier: Stripe.PriceCreateParams.Tier = { unit_amount_decimal: stripeUnitAmountDecimal, - up_to: isFinalTier(tier) ? "inf" : upTo, + up_to: upTo, }; if (tier.flat_amount) { @@ -79,13 +82,13 @@ export const priceToStripePrepaidV2Tiers = ({ } // Divide all tiers by billing units - const dividedTiers = tiers.map((tier, index: number) => ({ + return tiers.map((tier, index) => ({ ...tier, up_to: - index === tiers.length - 1 + index === tiers.length - 1 || tier.up_to === "inf" ? "inf" - : new Decimal(tier.up_to ?? 0) + : new Decimal(tier.up_to) .div(config.billing_units ?? 1) .ceil() .toNumber(), @@ -94,6 +97,4 @@ export const priceToStripePrepaidV2Tiers = ({ .mul(config.billing_units ?? 1) .toString(), })); - - return dividedTiers; }; diff --git a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx index edaf869d8..b2cd1d3ef 100644 --- a/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx +++ b/vite/src/views/products/features/credit-systems/components/AiCreditSchema.tsx @@ -1,7 +1,12 @@ -import { PlusIcon } from "lucide-react"; +import { InfoIcon } from "lucide-react"; import { FormLabel } from "@/components/v2/form/FormLabel"; import { Input } from "@/components/v2/inputs/Input"; import { SearchableSelect } from "@/components/v2/selects/SearchableSelect"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/v2/tooltips/Tooltip"; import { useAiProviders } from "../hooks/useAiProviders"; import type { CreditSystemFormInstance } from "../hooks/useCreditSystemForm"; import { AiCreditSchemaTable } from "./AiCreditSchemaTable"; @@ -43,39 +48,51 @@ export function AiCreditSchema({ form }: AiCreditSchemaProps) { /> -
- {activeProviderKeys.map((providerKey) => { - const provider = providers[providerKey]; - const modelFullIds = providerGroups[providerKey] ?? []; - const providerName = - provider?.name ?? - providerKey.charAt(0).toUpperCase() + providerKey.slice(1); + {activeProviderKeys.length > 0 && ( +
+ {activeProviderKeys.map((providerKey) => { + const provider = providers[providerKey]; + const modelFullIds = providerGroups[providerKey] ?? []; + const providerName = + provider?.name ?? + providerKey.charAt(0).toUpperCase() + providerKey.slice(1); - return ( - - ); - })} -
+ return ( + + ); + })} +
+ )}
e.stopPropagation()} > - Add Provider + + Add Provider Override + + + + + + Add specific markup overrides for certain providers/models. + + + { - if (!releaseDate) return -1; - const timestamp = Date.parse(releaseDate); - return Number.isNaN(timestamp) ? -1 : timestamp; -}; - -function getDefaultModelMarkups( - providers: Record, -): Record { - const result: Record = {}; - const preferredProvider = - providers["openrouter"] ?? Object.values(providers)[0]; - if (!preferredProvider) return result; - - const providerKey = preferredProvider.id; - for (const company of DEFAULT_AI_MODEL_COMPANIES) { - const companyModels = Object.entries(preferredProvider.models).filter( - ([key]) => key.startsWith(company), - ); - - const latestModel = companyModels.reduce< - [string, ModelsDevProvider["models"][string]] | null - >((currentLatest, candidate) => { - if (!currentLatest) return candidate; - const currentRelease = getReleaseDateMs(currentLatest[1].release_date); - const candidateRelease = getReleaseDateMs(candidate[1].release_date); - return candidateRelease > currentRelease ? candidate : currentLatest; - }, null); - - if (!latestModel) continue; - - const [modelKey] = latestModel; - result[joinModelId(providerKey, modelKey)] = {}; - } - return result; -} - interface CreditSystemSchemaProps { form: CreditSystemFormInstance; disableModeSwitch?: boolean; @@ -63,20 +18,16 @@ export function CreditSystemSchema({ form, disableModeSwitch = false, }: CreditSystemSchemaProps) { - const { providers } = useModelsDevPricing(); const type = useStore(form.store, (s) => s.values.type); const mode: CreditSchemaMode = isAiCreditSystem(type) ? "ai" : "classic"; const handleModeChange = (newMode: string) => { if (newMode === "ai") { - const modelMarkups = getDefaultModelMarkups(providers); form.setFieldValue("type", FeatureType.AiCreditSystem); form.setFieldValue("config", { ...form.state.values.config, schema: [] }); - form.setFieldValue( - "model_markups", - Object.keys(modelMarkups).length > 0 ? modelMarkups : {}, - ); + form.setFieldValue("model_markups", {}); + form.setFieldValue("provider_markups", {}); } else { form.setFieldValue("type", FeatureType.CreditSystem); form.setFieldValue("config", { @@ -86,6 +37,7 @@ export function CreditSystemSchema({ ], }); form.setFieldValue("model_markups", {}); + form.setFieldValue("provider_markups", {}); } }; diff --git a/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx b/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx index 3e06d28c3..f046fe991 100644 --- a/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx +++ b/vite/src/views/products/plan/components/new-feature/NewFeatureType.tsx @@ -2,6 +2,7 @@ import { FeatureType as APIFeatureType, type CreateFeature, FeatureUsageType, + isAnyCreditSystem, } from "@autumn/shared"; import { BarcodeIcon, CoinsIcon } from "@phosphor-icons/react"; import { PanelButton } from "@/components/v2/buttons/PanelButton"; @@ -57,7 +58,7 @@ export function NewFeatureType({
{ setFeature({ ...feature,