vibezcheck 0.2.0 → 0.4.0

package/dist/ai-sdk/index.js

@@ -30,11 +30,24 @@ var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: tru
 // src/ai-sdk/index.ts
 var ai_sdk_exports = {};
 __export(ai_sdk_exports, {
+  createVibezModel: () => createVibezModel,
+  createVibezSession: () => createVibezSession,
   meteredModel: () => meteredModel,
+  vibezcheck: () => vibezcheck,
   withBilling: () => withBilling
 });
 module.exports = __toCommonJS(ai_sdk_exports);
 
+// src/types.ts
+var VibezCircuitBreakerError = class extends Error {
+  event;
+  constructor(event) {
+    super(event.message);
+    this.name = "VibezCircuitBreakerError";
+    this.event = event;
+  }
+};
+
 // src/meter/client.ts
 var import_stripe = __toESM(require("stripe"));
 
@@ -576,14 +589,26 @@ function getModelPricing(modelName) {
 
 // src/pricing/calculator.ts
 function calculateCost(params) {
-  const rates = getModelPricing(params.model);
+  let rates;
+  if (params.customRate) {
+    const outRate = params.customRate.out ?? params.customRate.output ?? 0;
+    rates = {
+      inputPer1M: params.customRate.in,
+      outputPer1M: outRate,
+      reasoningPer1M: params.customRate.reasoning ?? outRate,
+      cachedInputPer1M: params.customRate.cached ?? params.customRate.in * 0.15,
+      currency: "USD"
+    };
+  } else {
+    rates = getModelPricing(params.model);
+  }
   const inputTokens = params.inputTokens ?? 0;
   const outputTokens = params.outputTokens ?? 0;
   const reasoningTokens = params.reasoningTokens ?? 0;
   const cachedTokens = params.cachedTokens ?? 0;
   const regularInputTokens = Math.max(0, inputTokens - cachedTokens);
   const regularInputCost = regularInputTokens / 1e6 * rates.inputPer1M;
-  const cachedRate = rates.cachedInputPer1M ?? rates.inputPer1M * 0.5;
+  const cachedRate = rates.cachedInputPer1M ?? rates.inputPer1M * 0.15;
   const cachedInputCost = cachedTokens / 1e6 * cachedRate;
   const inputCostUSD = regularInputCost + cachedInputCost;
   const outputCostUSD = outputTokens / 1e6 * rates.outputPer1M;
@@ -591,20 +616,27 @@ function calculateCost(params) {
   const reasoningCostUSD = reasoningTokens / 1e6 * reasoningRate;
   const standardCacheCost = cachedTokens / 1e6 * rates.inputPer1M;
   const cachedDiscountUSD = Math.max(0, standardCacheCost - cachedInputCost);
-  const totalUSD = inputCostUSD + outputCostUSD;
+  let totalUSD = inputCostUSD + outputCostUSD;
   const markup = params.markupMultiplier ?? 1;
-  const retailUSD = markup !== 1 ? totalUSD * markup : void 0;
+  let retailUSD = void 0;
+  if (markup !== 1 || params.minimumChargeUSD !== void 0) {
+    retailUSD = totalUSD * markup;
+    if (params.minimumChargeUSD !== void 0 && retailUSD < params.minimumChargeUSD) {
+      retailUSD = params.minimumChargeUSD;
+    }
+  }
   return {
     inputCostUSD: Number(inputCostUSD.toFixed(8)),
     outputCostUSD: Number(outputCostUSD.toFixed(8)),
     reasoningCostUSD: reasoningTokens > 0 ? Number(reasoningCostUSD.toFixed(8)) : void 0,
     cachedDiscountUSD: cachedTokens > 0 ? Number(cachedDiscountUSD.toFixed(8)) : void 0,
     totalUSD: Number(totalUSD.toFixed(8)),
-    retailUSD: retailUSD ? Number(retailUSD.toFixed(8)) : void 0,
+    retailUSD: retailUSD !== void 0 ? Number(retailUSD.toFixed(8)) : void 0,
     currency: rates.currency || "USD"
   };
 }
-function calculateUsageCost(model, usage, markupMultiplier) {
+function calculateUsageCost(model, usage, optionsOrMarkup) {
+  const options = typeof optionsOrMarkup === "number" ? { markupMultiplier: optionsOrMarkup } : optionsOrMarkup || {};
   return calculateCost({
     model,
     inputTokens: usage.inputTokens,
@@ -612,7 +644,9 @@ function calculateUsageCost(model, usage, markupMultiplier) {
     reasoningTokens: usage.reasoningTokens,
     cachedTokens: usage.cachedTokens,
     cacheWriteTokens: usage.cacheWriteTokens,
-    markupMultiplier
+    markupMultiplier: options.markupMultiplier,
+    minimumChargeUSD: options.minimumChargeUSD,
+    customRate: options.customRate
   });
 }
 
@@ -863,7 +897,22 @@ function withBilling(model, options = {}) {
   const customerId = typeof options.customer === "string" ? options.customer : options.customer?.id || options.customerId;
   const modelId = model.modelId || "unknown-model";
   const provider = model.provider?.replace(/^@ai-sdk\//, "") || "ai-sdk";
-  const handleUsage = (rawUsage) => {
+  const maxCostPerCall = options.maxCostPerCallUSD !== void 0 ? options.maxCostPerCallUSD : 0.5;
+  const captureOnAbort = options.captureOnAbort !== false;
+  const scheduleFlush = () => {
+    try {
+      const flushPromise = meter.flush();
+      if (typeof globalThis !== "undefined") {
+        const g = globalThis;
+        if (typeof g.after === "function") {
+          g.after(() => flushPromise);
+          return;
+        }
+      }
+    } catch {
+    }
+  };
+  const handleUsage = (rawUsage, extraMeta) => {
     if (!rawUsage) return;
     const inputTokens = rawUsage.promptTokens ?? rawUsage.inputTokens ?? 0;
     const outputTokens = rawUsage.completionTokens ?? rawUsage.outputTokens ?? 0;
@@ -877,7 +926,43 @@ function withBilling(model, options = {}) {
       visibleOutputTokens: Math.max(0, outputTokens - reasoningTokens),
       cachedTokens: cachedTokens > 0 ? cachedTokens : void 0
     };
-    const cost = calculateUsageCost(modelId, usage);
+    const cost = calculateUsageCost(modelId, usage, {
+      markupMultiplier: options.pricing?.margin,
+      minimumChargeUSD: options.pricing?.minimumChargeUSD,
+      customRate: options.rate
+    });
+    if (maxCostPerCall > 0 && maxCostPerCall !== Infinity && cost.totalUSD > maxCostPerCall) {
+      const budgetEvent = {
+        reason: "cost_per_call_exceeded",
+        limit: maxCostPerCall,
+        current: cost.totalUSD,
+        model: modelId,
+        customerId,
+        message: `[vibezcheck] Circuit Breaker tripped: Call cost ($${cost.totalUSD.toFixed(4)}) exceeded budget ceiling ($${maxCostPerCall.toFixed(4)}).`
+      };
+      if (options.onBudgetExceeded) {
+        options.onBudgetExceeded(budgetEvent);
+      }
+      if (options.throwOnBudgetExceeded) {
+        throw new VibezCircuitBreakerError(budgetEvent);
+      }
+    }
+    if (options.maxTokensPerCall && usage.totalTokens > options.maxTokensPerCall) {
+      const budgetEvent = {
+        reason: "max_tokens_exceeded",
+        limit: options.maxTokensPerCall,
+        current: usage.totalTokens,
+        model: modelId,
+        customerId,
+        message: `[vibezcheck] Circuit Breaker tripped: Token count (${usage.totalTokens}) exceeded maxTokensPerCall limit (${options.maxTokensPerCall}).`
+      };
+      if (options.onBudgetExceeded) {
+        options.onBudgetExceeded(budgetEvent);
+      }
+      if (options.throwOnBudgetExceeded) {
+        throw new VibezCircuitBreakerError(budgetEvent);
+      }
+    }
     const event = {
       timestamp: (/* @__PURE__ */ new Date()).toISOString(),
       model: modelId,
@@ -885,7 +970,11 @@ function withBilling(model, options = {}) {
       usage,
       cost,
       customerId,
-      metadata: options.metadata
+      metadata: {
+        ...options.metadata,
+        ...extraMeta,
+        billingMode: options.billing?.mode || "postpaid"
+      }
     };
     meter.recordUsage({
       model: modelId,
@@ -895,24 +984,16 @@ function withBilling(model, options = {}) {
       reasoningTokens: usage.reasoningTokens,
       cachedTokens: usage.cachedTokens,
       customerId,
-      metadata: options.metadata
+      metadata: event.metadata
     });
     if (options.onUsage) {
       options.onUsage(event);
     }
+    scheduleFlush();
   };
-  return new Proxy(model, {
+  const handler = {
     get(target, prop, receiver) {
       const originalValue = Reflect.get(target, prop, receiver);
-      if (prop === "doGenerate" && typeof originalValue === "function") {
-        return async function(...args) {
-          const result = await originalValue.apply(target, args);
-          if (result && result.usage) {
-            handleUsage(result.usage);
-          }
-          return result;
-        };
-      }
       if (prop === "doStream" && typeof originalValue === "function") {
         return async function(...args) {
           const result = await originalValue.apply(target, args);
@@ -920,63 +1001,196 @@ function withBilling(model, options = {}) {
             return result;
           }
           const originalStream = result.stream;
-          if (typeof originalStream.getReader === "function") {
-            const reader = originalStream.getReader();
-            const transformedStream = new ReadableStream({
-              async start(controller) {
-                try {
-                  while (true) {
-                    const { done, value } = await reader.read();
-                    if (done) {
-                      controller.close();
-                      break;
-                    }
-                    if (value && typeof value === "object") {
-                      if (value.type === "finish" && value.usage) {
-                        handleUsage(value.usage);
-                      }
-                    }
-                    controller.enqueue(value);
-                  }
-                } catch (err) {
-                  controller.error(err);
-                }
-              }
-            });
-            return {
-              ...result,
-              stream: transformedStream
-            };
+          let streamCompleted = false;
+          let accumulatedChars = 0;
+          let estimatedInputTokens = 0;
+          try {
+            if (args[0]?.prompt) {
+              const str = typeof args[0].prompt === "string" ? args[0].prompt : JSON.stringify(args[0].prompt);
+              estimatedInputTokens = Math.ceil(str.length / 4);
+            }
+          } catch {
           }
-          if (Symbol.asyncIterator in originalStream) {
-            const wrappedAsyncIterable = {
-              async *[Symbol.asyncIterator]() {
-                for await (const chunk of originalStream) {
-                  if (chunk && typeof chunk === "object") {
-                    if (chunk.type === "finish" && chunk.usage) {
-                      handleUsage(chunk.usage);
-                    }
-                  }
-                  yield chunk;
+          const abortSignal = args[0]?.abortSignal;
+          if (abortSignal && captureOnAbort) {
+            abortSignal.addEventListener(
+              "abort",
+              () => {
+                if (!streamCompleted && accumulatedChars > 0) {
+                  streamCompleted = true;
+                  const estimatedOutputTokens = Math.ceil(accumulatedChars / 3.8);
+                  handleUsage(
+                    {
+                      inputTokens: estimatedInputTokens,
+                      outputTokens: estimatedOutputTokens
+                    },
+                    { aborted: true, partial: true }
+                  );
                 }
+              },
+              { once: true }
+            );
+          }
+          const transformStream = new TransformStream({
+            transform(chunk, controller) {
+              controller.enqueue(chunk);
+              if (chunk.type === "text-delta" && chunk.textDelta) {
+                accumulatedChars += chunk.textDelta.length;
               }
-            };
-            return {
-              ...result,
-              stream: wrappedAsyncIterable
-            };
+              if (chunk.type === "finish" && chunk.usage) {
+                streamCompleted = true;
+                handleUsage(chunk.usage);
+              }
+            },
+            flush() {
+              if (!streamCompleted && captureOnAbort && accumulatedChars > 0) {
+                streamCompleted = true;
+                const estimatedOutputTokens = Math.ceil(accumulatedChars / 3.8);
+                handleUsage(
+                  {
+                    inputTokens: estimatedInputTokens,
+                    outputTokens: estimatedOutputTokens
+                  },
+                  { aborted: true, partial: true }
+                );
+              }
+            }
+          });
+          return {
+            ...result,
+            stream: originalStream.pipeThrough(transformStream)
+          };
+        };
+      }
+      if (prop === "doGenerate" && typeof originalValue === "function") {
+        return async function(...args) {
+          const result = await originalValue.apply(target, args);
+          if (result && result.usage) {
+            handleUsage(result.usage);
           }
           return result;
         };
       }
       return originalValue;
     }
-  });
+  };
+  return new Proxy(model, handler);
 }
 var meteredModel = withBilling;
+
+// src/ai-sdk/declarative.ts
+function createVibezModel(modelOrId, options = {}) {
+  if (typeof modelOrId === "object" && modelOrId !== null) {
+    return withBilling(modelOrId, options);
+  }
+  if (typeof modelOrId === "string") {
+    const rawId = modelOrId;
+    let providerName = "openai";
+    let cleanModelId = rawId;
+    if (rawId.includes("/")) {
+      const parts = rawId.split("/");
+      providerName = parts[0].toLowerCase();
+      cleanModelId = parts.slice(1).join("/");
+    }
+    const apiKey = options.apiKey || process.env.AI_GATEWAY_API_KEY || process.env.OPENAI_API_KEY || process.env.ANTHROPIC_API_KEY || process.env.GOOGLE_GENERATIVE_AI_API_KEY;
+    const baseURL = options.baseURL || process.env.AI_GATEWAY_BASE_URL || (providerName === "openai" ? "https://api.openai.com/v1" : void 0);
+    let baseModelInstance = null;
+    try {
+      if (providerName === "openai" || providerName === "gateway" || !providerName) {
+        const { createOpenAI } = require("@ai-sdk/openai");
+        const openaiProvider = createOpenAI({ apiKey, baseURL });
+        baseModelInstance = openaiProvider(cleanModelId);
+      } else if (providerName === "anthropic") {
+        const { createAnthropic } = require("@ai-sdk/anthropic");
+        const anthropicProvider = createAnthropic({ apiKey, baseURL });
+        baseModelInstance = anthropicProvider(cleanModelId);
+      } else if (providerName === "google") {
+        const { createGoogleGenerativeAI } = require("@ai-sdk/google");
+        const googleProvider = createGoogleGenerativeAI({ apiKey, baseURL });
+        baseModelInstance = googleProvider(cleanModelId);
+      }
+    } catch {
+      baseModelInstance = {
+        specificationVersion: "v2",
+        provider: providerName,
+        modelId: cleanModelId
+      };
+    }
+    if (!baseModelInstance) {
+      baseModelInstance = {
+        specificationVersion: "v2",
+        provider: providerName,
+        modelId: cleanModelId
+      };
+    }
+    return withBilling(baseModelInstance, options);
+  }
+  return modelOrId;
+}
+function createVibezSession(sessionOptions = {}) {
+  const meter = createMeter({
+    apiKey: sessionOptions.stripeApiKey
+  });
+  const customerId = typeof sessionOptions.customer === "string" ? sessionOptions.customer : sessionOptions.customer?.id;
+  return {
+    meter,
+    model: (modelOrId, callOptions = {}) => {
+      return createVibezModel(modelOrId, {
+        customer: sessionOptions.customer,
+        pricing: sessionOptions.pricing,
+        billing: sessionOptions.billing,
+        meter,
+        metadata: {
+          ...sessionOptions.metadata,
+          ...callOptions.metadata
+        },
+        ...callOptions
+      });
+    },
+    trackTool: async (name, { costUSD, metadata }) => {
+      const event = {
+        timestamp: (/* @__PURE__ */ new Date()).toISOString(),
+        model: `tool:${name}`,
+        provider: "tool",
+        usage: {
+          inputTokens: 0,
+          outputTokens: 0,
+          totalTokens: 0
+        },
+        cost: {
+          inputCostUSD: 0,
+          outputCostUSD: costUSD,
+          totalUSD: costUSD,
+          currency: "USD"
+        },
+        customerId,
+        metadata: {
+          ...sessionOptions.metadata,
+          ...metadata,
+          tool: name
+        }
+      };
+      meter.recordUsage({
+        model: `tool:${name}`,
+        provider: "tool",
+        inputTokens: 0,
+        outputTokens: 0,
+        customerId,
+        metadata: event.metadata
+      });
+      await meter.flush();
+    }
+  };
+}
+var vibezcheck = Object.assign(createVibezModel, {
+  session: createVibezSession
+});
 // Annotate the CommonJS export names for ESM import in node:
 0 && (module.exports = {
+  createVibezModel,
+  createVibezSession,
   meteredModel,
+  vibezcheck,
   withBilling
 });
 //# sourceMappingURL=index.js.map