@theihtisham/budget-llm 1.0.0

package/src/task-classifier.ts

@@ -0,0 +1,134 @@
+import type { TaskType, ChatMessage } from './types';
+
+// Keyword patterns for each task type, ordered by specificity
+interface TaskPattern {
+  type: TaskType;
+  keywords: string[];
+  roleWeights: Partial<Record<ChatMessage['role'], number>>;
+}
+
+const TASK_PATTERNS: TaskPattern[] = [
+  {
+    type: 'code',
+    keywords: [
+      'function', 'class', 'method', 'variable', 'import', 'export',
+      'return', 'interface', 'type ', 'const ', 'let ', 'var ',
+      'async', 'await', 'promise', 'callback', 'error handling',
+      'algorithm', 'data structure', 'api endpoint', 'sql query',
+      'database', 'debug', 'fix bug', 'refactor', 'code review',
+      'typescript', 'javascript', 'python', 'rust', 'golang',
+      'implementation', 'compile', 'deploy', 'unit test',
+      'regex', 'html', 'css', 'component', 'render',
+      'write a function', 'write a class', 'create a module',
+      'implement', 'optimize', 'performance',
+    ],
+    roleWeights: { system: 1.5, user: 1.0 },
+  },
+  {
+    type: 'math',
+    keywords: [
+      'calculate', 'equation', 'formula', 'derivative', 'integral',
+      'polynomial', 'theorem', 'proof', 'linear algebra', 'calculus',
+      'probability', 'statistics', 'matrix', 'vector', 'solve for x',
+      'mathematical', 'compute', 'arithmetic', 'geometric',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'reasoning',
+    keywords: [
+      'analyze', 'evaluate', 'compare', 'contrast', 'pros and cons',
+      'should i', 'which is better', 'explain why', 'logical',
+      'argument', 'evidence', 'implication', 'consequence',
+      'what would happen', 'predict', 'deduce', 'infer',
+      'step by step reasoning', 'think through',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'creative',
+    keywords: [
+      'write a story', 'poem', 'creative', 'imagine', 'fictional',
+      'narrative', 'dialogue', 'character', 'plot', 'brainstorm',
+      'idea', 'invent', 'design a', 'slogan', 'tagline',
+      'marketing copy', 'blog post', 'article about',
+      'song lyrics', 'joke', 'creative writing',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'translation',
+    keywords: [
+      'translate', 'translation', 'in french', 'in spanish',
+      'in german', 'in japanese', 'in chinese', 'in korean',
+      'convert to', 'say this in',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'summarization',
+    keywords: [
+      'summarize', 'summary', 'tldr', 'brief overview',
+      'key points', 'main ideas', 'condense', 'shorten',
+      'bullet points', 'executive summary',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'analysis',
+    keywords: [
+      'analyze', 'analysis', 'break down', 'examine',
+      'data analysis', 'trends', 'insights', 'report',
+      'metrics', 'dashboard', 'visualization',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+  {
+    type: 'chat',
+    keywords: [
+      'hello', 'hi ', 'hey', 'how are you', 'thanks', 'thank you',
+      'help me', 'can you', 'what is', 'who is', 'when did',
+      'where is', 'tell me about',
+    ],
+    roleWeights: { user: 1.0 },
+  },
+];
+
+export function classifyTask(messages: ChatMessage[]): TaskType {
+  const text = messages.map((m) => m.content.toLowerCase()).join(' ');
+
+  let bestType: TaskType = 'chat';
+  let bestScore = 0;
+
+  for (const pattern of TASK_PATTERNS) {
+    let score = 0;
+    for (const keyword of pattern.keywords) {
+      if (text.includes(keyword)) {
+        score += 1;
+      }
+    }
+
+    // Weight by role
+    for (const msg of messages) {
+      const weight = pattern.roleWeights[msg.role] ?? 0.5;
+      for (const keyword of pattern.keywords) {
+        if (msg.content.toLowerCase().includes(keyword)) {
+          score += weight * 0.5;
+        }
+      }
+    }
+
+    if (score > bestScore) {
+      bestScore = score;
+      bestType = pattern.type;
+    }
+  }
+
+  return bestType;
+}
+
+export function estimateTokens(messages: ChatMessage[]): number {
+  // Rough estimate: ~4 chars per token
+  const totalChars = messages.reduce((sum, m) => sum + m.content.length + m.role.length, 0);
+  return Math.ceil(totalChars / 4);
+}

package/src/types/sql.js.d.ts

@@ -0,0 +1,27 @@
+declare module 'sql.js' {
+  interface SqlJsStatic {
+    Database: new (data?: ArrayLike<number | Buffer>) => Database;
+  }
+
+  interface Database {
+    run(sql: string, params?: (string | number | null | undefined)[]): Database;
+    exec(sql: string, params?: (string | number | null | undefined)[]): Array<{ columns: string[]; values: unknown[][] }>;
+    prepare(sql: string): Statement;
+    export(): Uint8Array;
+    close(): void;
+    getRowsModified(): number;
+  }
+
+  interface Statement {
+    bind(params?: (string | number | null | undefined)[]): boolean;
+    step(): boolean;
+    getAsObject(): Record<string, unknown>;
+    get(column: number): unknown;
+    getColumnNames(): string[];
+    free(): boolean;
+    reset(): boolean;
+  }
+
+  export default function initSqlJs(config?: { locateFile?: (file: string) => string }): Promise<SqlJsStatic>;
+  export { Database, Statement };
+}

package/src/types.ts

@@ -0,0 +1,258 @@
+// ---- Provider & Model Types ----
+
+export type ProviderId = 'openai' | 'anthropic' | 'google' | 'deepseek';
+
+export interface ProviderConfig {
+  id: ProviderId;
+  name: string;
+  baseUrl: string;
+  apiKey: string;
+  enabled: boolean;
+  priority: number; // lower = higher priority in fallback
+  timeoutMs: number;
+}
+
+export interface ModelInfo {
+  id: string;
+  provider: ProviderId;
+  displayName: string;
+  inputPricePer1M: number;  // USD per 1M input tokens
+  outputPricePer1M: number; // USD per 1M output tokens
+  contextWindow: number;
+  maxOutputTokens: number;
+  capabilities: ModelCapability[];
+  qualityScore: number;     // 1-10 overall quality
+  speedScore: number;       // 1-10 latency speed
+  costScore: number;        // 1-10 cost efficiency (10 = cheapest)
+}
+
+export type ModelCapability =
+  | 'code'
+  | 'creative'
+  | 'reasoning'
+  | 'chat'
+  | 'summarization'
+  | 'translation'
+  | 'analysis'
+  | 'math';
+
+export type TaskType =
+  | 'code'
+  | 'creative'
+  | 'reasoning'
+  | 'chat'
+  | 'summarization'
+  | 'translation'
+  | 'analysis'
+  | 'math'
+  | 'auto';
+
+// ---- Request/Response Types ----
+
+export interface ChatMessage {
+  role: 'system' | 'user' | 'assistant';
+  content: string;
+}
+
+export interface ProxyRequest {
+  model?: string;
+  messages: ChatMessage[];
+  temperature?: number;
+  max_tokens?: number;
+  top_p?: number;
+  stream?: boolean;
+  /** BudgetLLM extension: force a specific task type hint */
+  task_type?: TaskType;
+  /** BudgetLLM extension: request ID for tracking */
+  request_id?: string;
+  /** BudgetLLM extension: per-request budget cap in USD */
+  budget_cap?: number;
+}
+
+export interface ProxyResponse {
+  id: string;
+  object: 'chat.completion';
+  created: number;
+  model: string;
+  choices: Array<{
+    index: number;
+    message: ChatMessage;
+    finish_reason: string;
+  }>;
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+  /** BudgetLLM extension: cost breakdown */
+  cost: CostBreakdown;
+}
+
+export interface CostBreakdown {
+  inputCost: number;
+  outputCost: number;
+  totalCost: number;
+  currency: string;
+  model: string;
+  provider: ProviderId;
+  savingsVsGpt4: number; // how much saved vs sending to GPT-4
+}
+
+// ---- Budget Types ----
+
+export interface BudgetConfig {
+  dailyBudget: number;
+  monthlyBudget: number;
+  perRequestCap: number;
+}
+
+export interface BudgetStatus {
+  daily: { spent: number; limit: number; remaining: number; percentUsed: number };
+  monthly: { spent: number; limit: number; remaining: number; percentUsed: number };
+}
+
+// ---- Cost Tracking Types ----
+
+export interface CostRecord {
+  id: string;
+  requestId: string;
+  timestamp: string;
+  provider: ProviderId;
+  model: string;
+  taskType: TaskType;
+  inputTokens: number;
+  outputTokens: number;
+  inputCost: number;
+  outputCost: number;
+  totalCost: number;
+  latencyMs: number;
+  cached: boolean;
+}
+
+export interface CostSummary {
+  totalSpent: number;
+  totalRequests: number;
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  cacheHitRate: number;
+  averageLatencyMs: number;
+  byProvider: Record<ProviderId, ProviderSummary>;
+  byModel: Record<string, ModelSummary>;
+  savingsVsGpt4: number;
+}
+
+export interface ProviderSummary {
+  totalSpent: number;
+  requestCount: number;
+  inputTokens: number;
+  outputTokens: number;
+}
+
+export interface ModelSummary {
+  totalSpent: number;
+  requestCount: number;
+  inputTokens: number;
+  outputTokens: number;
+  averageLatencyMs: number;
+}
+
+// ---- Cache Types ----
+
+export interface CacheEntry {
+  promptHash: string;
+  model: string;
+  provider: ProviderId;
+  response: string;
+  inputTokens: number;
+  outputTokens: number;
+  cost: number;
+  createdAt: number;
+  expiresAt: number;
+  hitCount: number;
+}
+
+// ---- Rate Limiting Types ----
+
+export interface TokenBucket {
+  tokens: number;
+  maxTokens: number;
+  refillRate: number; // tokens per second
+  lastRefill: number;
+}
+
+export interface RateLimitStatus {
+  allowed: boolean;
+  remaining: number;
+  resetMs: number;
+}
+
+// ---- Routing Types ----
+
+export interface RoutingDecision {
+  provider: ProviderId;
+  model: string;
+  reason: string;
+  estimatedCost: number;
+  estimatedSavings: number;
+}
+
+// ---- Dashboard Types ----
+
+export interface DashboardData {
+  overview: {
+    totalSpent: number;
+    totalRequests: number;
+    totalSaved: number;
+    cacheHitRate: number;
+  };
+  budget: BudgetStatus;
+  recentRequests: Array<{
+    time: string;
+    model: string;
+    provider: string;
+    taskType: string;
+    cost: number;
+    tokens: number;
+    cached: boolean;
+  }>;
+  costByDay: Array<{ date: string; cost: number; requests: number }>;
+  modelDistribution: Array<{ model: string; count: number; cost: number }>;
+}
+
+// ---- Error Types ----
+
+export class BudgetExceededError extends Error {
+  constructor(
+    public readonly type: 'daily' | 'monthly' | 'per_request',
+    public readonly spent: number,
+    public readonly limit: number
+  ) {
+    super(`Budget exceeded (${type}): $${spent.toFixed(4)} / $${limit.toFixed(2)}`);
+    this.name = 'BudgetExceededError';
+  }
+}
+
+export class RateLimitError extends Error {
+  constructor(public readonly resetMs: number) {
+    super(`Rate limit exceeded. Retry after ${Math.ceil(resetMs / 1000)}s`);
+    this.name = 'RateLimitError';
+  }
+}
+
+export class ProviderError extends Error {
+  constructor(
+    public readonly provider: ProviderId,
+    public readonly statusCode: number,
+    message: string
+  ) {
+    super(`Provider ${provider} error (${statusCode}): ${message}`);
+    this.name = 'ProviderError';
+  }
+}
+
+export class NoProviderAvailableError extends Error {
+  constructor() {
+    super('No LLM provider available — check API keys and provider health');
+    this.name = 'NoProviderAvailableError';
+  }
+}

package/src/utils/encryption.ts

@@ -0,0 +1,36 @@
+import crypto from 'crypto';
+import { env } from '../config';
+
+const ALGORITHM = 'aes-256-gcm';
+const IV_LENGTH = 16;
+const TAG_LENGTH = 16;
+
+function getKey(): Buffer {
+  const key = env.ENCRYPTION_KEY;
+  // Ensure exactly 32 bytes for AES-256
+  return crypto.createHash('sha256').update(key).digest();
+}
+
+export function encrypt(plaintext: string): string {
+  const iv = crypto.randomBytes(IV_LENGTH);
+  const key = getKey();
+  const cipher = crypto.createCipheriv(ALGORITHM, key, iv);
+  const encrypted = Buffer.concat([cipher.update(plaintext, 'utf8'), cipher.final()]);
+  const tag = cipher.getAuthTag();
+  return Buffer.concat([iv, tag, encrypted]).toString('base64');
+}
+
+export function decrypt(ciphertext: string): string {
+  const data = Buffer.from(ciphertext, 'base64');
+  const iv = data.subarray(0, IV_LENGTH);
+  const tag = data.subarray(IV_LENGTH, IV_LENGTH + TAG_LENGTH);
+  const encrypted = data.subarray(IV_LENGTH + TAG_LENGTH);
+  const key = getKey();
+  const decipher = crypto.createDecipheriv(ALGORITHM, key, iv);
+  decipher.setAuthTag(tag);
+  return decipher.update(encrypted) + decipher.final('utf8');
+}
+
+export function hashPrompt(messages: string): string {
+  return crypto.createHash('sha256').update(messages).digest('hex');
+}

package/tests/config.test.ts

@@ -0,0 +1,85 @@
+import { describe, it, expect } from 'vitest';
+import './setup';
+import { env, getProviders, MODEL_CATALOG, getDefaultBudget } from '../src/config';
+
+describe('Config', () => {
+  it('should parse environment variables', () => {
+    expect(env.PORT).toBe(3211);
+    expect(env.NODE_ENV).toBe('test');
+    expect(env.ENCRYPTION_KEY).toBeTruthy();
+  });
+
+  it('should return provider configurations', () => {
+    const providers = getProviders();
+    expect(providers).toHaveLength(4);
+    expect(providers.map(p => p.id)).toEqual(
+      expect.arrayContaining(['openai', 'anthropic', 'google', 'deepseek'])
+    );
+  });
+
+  it('should mark providers as enabled when API key is set', () => {
+    const providers = getProviders();
+    for (const provider of providers) {
+      expect(provider.enabled).toBe(true);
+    }
+  });
+
+  it('should have valid model catalog entries', () => {
+    for (const model of MODEL_CATALOG) {
+      expect(model.id).toBeTruthy();
+      expect(model.provider).toBeTruthy();
+      expect(model.displayName).toBeTruthy();
+      expect(model.inputPricePer1M).toBeGreaterThan(0);
+      expect(model.outputPricePer1M).toBeGreaterThan(0);
+      expect(model.contextWindow).toBeGreaterThan(0);
+      expect(model.maxOutputTokens).toBeGreaterThan(0);
+      expect(model.capabilities.length).toBeGreaterThan(0);
+      expect(model.qualityScore).toBeGreaterThanOrEqual(1);
+      expect(model.qualityScore).toBeLessThanOrEqual(10);
+      expect(model.speedScore).toBeGreaterThanOrEqual(1);
+      expect(model.speedScore).toBeLessThanOrEqual(10);
+      expect(model.costScore).toBeGreaterThanOrEqual(1);
+      expect(model.costScore).toBeLessThanOrEqual(10);
+    }
+  });
+
+  it('should include models from all providers', () => {
+    const providers = new Set(MODEL_CATALOG.map(m => m.provider));
+    expect(providers.has('openai')).toBe(true);
+    expect(providers.has('anthropic')).toBe(true);
+    expect(providers.has('google')).toBe(true);
+    expect(providers.has('deepseek')).toBe(true);
+  });
+
+  it('should return default budget', () => {
+    const budget = getDefaultBudget();
+    expect(budget.dailyBudget).toBeGreaterThan(0);
+    expect(budget.monthlyBudget).toBeGreaterThan(0);
+    expect(budget.perRequestCap).toBeGreaterThan(0);
+  });
+
+  it('should have provider config with required fields', () => {
+    const providers = getProviders();
+    for (const provider of providers) {
+      expect(provider.id).toBeTruthy();
+      expect(provider.name).toBeTruthy();
+      expect(provider.baseUrl).toBeTruthy();
+      expect(provider.baseUrl).toMatch(/^https:\/\//);
+      expect(typeof provider.priority).toBe('number');
+      expect(typeof provider.timeoutMs).toBe('number');
+    }
+  });
+
+  it('should have unique model IDs', () => {
+    const ids = MODEL_CATALOG.map(m => m.id);
+    const uniqueIds = new Set(ids);
+    expect(ids.length).toBe(uniqueIds.size);
+  });
+
+  it('should have unique provider priorities', () => {
+    const providers = getProviders();
+    const priorities = providers.map(p => p.priority);
+    const uniquePriorities = new Set(priorities);
+    expect(priorities.length).toBe(uniquePriorities.size);
+  });
+});