@od-oneapp/ai-platform 0.1.0

package/src/observability/cost/tracking/budget-guard.ts

@@ -0,0 +1,110 @@
+/**
+ * Budget guard middleware for cost control
+ */
+
+import { logWarn } from '@repo/shared/logs';
+
+import { BudgetExceededError } from '../types';
+
+import type { WrapGenerateArgs, WrapStreamArgs } from '../../../sdk/middleware/model-middleware';
+import type { SDKLanguageModelV3Middleware } from '../../../shared';
+import type { BudgetConfig, UsageSummary } from '../types';
+
+
+/**
+ * Function to get current usage - implement with your storage
+ */
+export type GetCurrentUsage = () => Promise<UsageSummary>;
+
+/**
+ * Create budget guard middleware that enforces spending limits
+ *
+ * @example
+ * ```ts
+ * const budgetGuard = createBudgetGuard(
+ *   {
+ *     maxDailyCost: 100,
+ *     maxMonthlyCost: 2000,
+ *     warningThreshold: 0.8,
+ *   },
+ *   async () => {
+ *     const [daily, monthly] = await Promise.all([
+ *       db.query.dailyUsage(),
+ *       db.query.monthlyUsage(),
+ *     ]);
+ *     return { daily, monthly, requestCount: 0 };
+ *   }
+ * );
+ *
+ * const model = wrapLanguageModel({
+ *   model: createLanguageModel(Model.SONNET),
+ *   middleware: budgetGuard,
+ * });
+ * ```
+ */
+export function createBudgetGuard(
+  config: BudgetConfig,
+  getCurrentUsage: GetCurrentUsage,
+): SDKLanguageModelV3Middleware {
+  const { maxDailyCost, maxMonthlyCost, warningThreshold = 0.8 } = config;
+
+  return {
+    specificationVersion: 'v3',
+    wrapGenerate: async ({ doGenerate }: WrapGenerateArgs) => {
+      const usage = await getCurrentUsage();
+
+      // Check daily limit
+      if (usage.daily >= maxDailyCost) {
+        throw new BudgetExceededError(
+          `Daily budget exceeded: $${usage.daily.toFixed(2)} / $${maxDailyCost}`,
+          'daily',
+          usage.daily,
+          maxDailyCost,
+        );
+      }
+
+      // Check monthly limit
+      if (usage.monthly >= maxMonthlyCost) {
+        throw new BudgetExceededError(
+          `Monthly budget exceeded: $${usage.monthly.toFixed(2)} / $${maxMonthlyCost}`,
+          'monthly',
+          usage.monthly,
+          maxMonthlyCost,
+        );
+      }
+
+      // Log warning if approaching limits
+      if (usage.daily >= maxDailyCost * warningThreshold) {
+        logWarn(
+          `[BudgetGuard] Daily budget warning: ${((usage.daily / maxDailyCost) * 100).toFixed(1)}% used`,
+        );
+      }
+      if (usage.monthly >= maxMonthlyCost * warningThreshold) {
+        logWarn(
+          `[BudgetGuard] Monthly budget warning: ${((usage.monthly / maxMonthlyCost) * 100).toFixed(1)}% used`,
+        );
+      }
+
+      return doGenerate();
+    },
+
+    wrapStream: async ({ doStream }: WrapStreamArgs) => {
+      const usage = await getCurrentUsage();
+
+      if (usage.daily >= maxDailyCost) {
+        throw new BudgetExceededError(`Daily budget exceeded`, 'daily', usage.daily, maxDailyCost);
+      }
+
+      if (usage.monthly >= maxMonthlyCost) {
+        throw new BudgetExceededError(
+          `Monthly budget exceeded`,
+          'monthly',
+          usage.monthly,
+          maxMonthlyCost,
+        );
+      }
+
+      return doStream();
+    },
+  };
+}

package/src/observability/cost/tracking/usage-tracker.ts

@@ -0,0 +1,120 @@
+/**
+ * Usage tracking middleware for AI SDK
+ */
+
+import { logError } from '@repo/shared/logs';
+
+import { Model, getModelMetadata } from '../../../models/registry';
+
+import type { WrapGenerateArgs, WrapStreamArgs } from '../../../sdk/middleware/model-middleware';
+import type { SDKLanguageModelV3Middleware } from '../../../shared';
+import type { UsageEvent, UsageCallback, CostMetadata } from '../types';
+
+/**
+ * Calculate cost based on token usage and model pricing
+ */
+const getUsageTotal = (value: unknown): number => {
+  if (typeof value === 'number') return value;
+  if (value && typeof value === 'object' && 'total' in value) {
+    const { total } = value as { total?: unknown };
+    return typeof total === 'number' ? total : 0;
+  }
+  return 0;
+};
+
+function calculateCost(
+  model: Model,
+  usage?: { inputTokens?: unknown; outputTokens?: unknown },
+): number {
+  if (!usage) return 0;
+
+  const metadata = getModelMetadata(model);
+  const inputTokens = getUsageTotal(usage.inputTokens);
+  const outputTokens = getUsageTotal(usage.outputTokens);
+  const inputCost = (inputTokens / 1000) * metadata.costPer1kInput;
+  const outputCost = (outputTokens / 1000) * metadata.costPer1kOutput;
+
+  return inputCost + outputCost;
+}
+
+/**
+ * Extract cost metadata from request parameters
+ */
+function extractMetadata(params: unknown): CostMetadata {
+  const p = params as Record<string, unknown>;
+  const headers = p.headers as Record<string, string> | undefined;
+
+  return {
+    userId: headers?.['x-user-id'] ?? undefined,
+    teamId: headers?.['x-team-id'] ?? undefined,
+    feature: headers?.['x-feature'] ?? undefined,
+    requestId: headers?.['x-request-id'] ?? undefined,
+  };
+}
+
+/**
+ * Create usage tracking middleware
+ *
+ * @example
+ * ```ts
+ * const tracker = createUsageTracker(async (event) => {
+ *   await db.insert(usageTable).values(event);
+ *   console.log(`Cost: $${event.cost.toFixed(4)}`);
+ * });
+ *
+ * const model = wrapLanguageModel({
+ *   model: createLanguageModel(Model.SONNET),
+ *   middleware: tracker,
+ * });
+ * ```
+ */
+export function createUsageTracker(callback: UsageCallback): SDKLanguageModelV3Middleware {
+  return {
+    specificationVersion: 'v3',
+    wrapGenerate: async ({ doGenerate, params }: WrapGenerateArgs) => {
+      const result = await doGenerate();
+
+      // Try to determine model from params
+      const rawParams = params as Record<string, unknown>;
+      const modelId =
+        (typeof rawParams.model === 'string' ? rawParams.model : undefined) ??
+        (typeof rawParams.modelId === 'string' ? rawParams.modelId : undefined);
+      const model = Object.values(Model).find(entry => entry === modelId) ?? Model.SONNET;
+
+      const event: UsageEvent = {
+        model,
+        inputTokens: getUsageTotal(result.usage?.inputTokens),
+        outputTokens: getUsageTotal(result.usage?.outputTokens),
+        cost: calculateCost(model, result.usage),
+        timestamp: new Date(),
+        metadata: extractMetadata(params),
+      };
+
+      // Fire callback async (don't block response)
+      void (async () => {
+        try {
+          await callback(event);
+        } catch (err) {
+          logError('[UsageTracker] Callback error', {
+            error: err instanceof Error ? err : new Error(String(err)),
+          });
+        }
+      })();
+
+      return result;
+    },
+
+    wrapStream: async ({ doStream }: WrapStreamArgs) => {
+      const { stream, ...rest } = await doStream();
+
+      // For streaming, we track usage from the final chunk
+      // The AI SDK accumulates usage across chunks
+      return {
+        stream,
+        ...rest,
+      };
+    },
+  };
+}
+
+export { calculateCost, extractMetadata };

package/src/observability/cost/types.ts

@@ -0,0 +1,85 @@
+/**
+ * Cost tracking and attribution types
+ */
+
+import type { Model } from '../../models/registry';
+
+/**
+ * Token usage from a single AI request
+ */
+export interface UsageEvent {
+  /** Semantic model name */
+  model: Model;
+  /** Input/prompt tokens */
+  inputTokens: number;
+  /** Output/completion tokens */
+  outputTokens: number;
+  /** Calculated cost in USD */
+  cost: number;
+  /** When the request occurred */
+  timestamp: Date;
+  /** Attribution metadata */
+  metadata: CostMetadata;
+}
+
+/**
+ * Cost attribution metadata
+ */
+export interface CostMetadata {
+  /** User ID for per-user tracking */
+  userId?: string;
+  /** Team ID for team-level aggregation */
+  teamId?: string;
+  /** Feature/product area for attribution */
+  feature?: string;
+  /** Request ID for correlation */
+  requestId?: string;
+  /** Additional custom tags */
+  tags?: Record<string, string>;
+}
+
+/**
+ * Budget configuration for cost controls
+ */
+export interface BudgetConfig {
+  /** Maximum daily spend in USD */
+  maxDailyCost: number;
+  /** Maximum monthly spend in USD */
+  maxMonthlyCost: number;
+  /** Warning threshold (0-1), e.g., 0.8 = 80% */
+  warningThreshold?: number;
+  /** Per-request cost limit */
+  maxRequestCost?: number;
+}
+
+/**
+ * Current usage summary
+ */
+export interface UsageSummary {
+  /** Daily spend so far */
+  daily: number;
+  /** Monthly spend so far */
+  monthly: number;
+  /** Request count */
+  requestCount: number;
+}
+
+/**
+ * Usage callback for tracking
+ */
+export type UsageCallback = (event: UsageEvent) => Promise<void>;
+
+/**
+ * Budget exceeded error
+ */
+export class BudgetExceededError extends Error {
+  constructor(
+    message: string,
+    public readonly budgetType: 'daily' | 'monthly' | 'request',
+    public readonly current: number,
+    public readonly limit: number,
+  ) {
+    super(message);
+    this.name = 'BudgetExceededError';
+  }
+}

package/src/observability/index.ts

@@ -0,0 +1,17 @@
+/**
+ * Observability Module
+ *
+ * Monitoring, tracing, and cost management:
+ * - Telemetry (OpenTelemetry integration)
+ * - Distributed tracing
+ * - Usage analytics
+ * - Cost tracking and attribution
+ *
+ * @module observability
+ */
+
+// Namespaced exports to avoid symbol collisions across modules
+export * as telemetry from "./telemetry";
+export * as tracing from "./tracing";
+export * as analytics from "./analytics";
+export * as cost from "./cost";