@theihtisham/budget-llm 1.0.0

package/tests/server.test.ts

@@ -0,0 +1,253 @@
+import { describe, it, expect, beforeAll } from 'vitest';
+import './setup';
+import express from 'express';
+
+describe('Server API Endpoints', () => {
+  let app: express.Express;
+
+  beforeAll(async () => {
+    // Create a test app with the same routes
+    app = express();
+    app.use(express.json());
+
+    // Import after env is set
+    const database = await import('../src/database');
+    const { MODEL_CATALOG } = await import('../src/config');
+    const { checkRateLimit } = await import('../src/rate-limiter');
+
+    // Initialize DB first
+    await database.initDb();
+
+    // Health
+    app.get('/health', (_req, res) => {
+      res.json({ status: 'ok', version: '1.0.0', timestamp: new Date().toISOString() });
+    });
+
+    // Models
+    app.get('/v1/models', (_req, res) => {
+      res.json({
+        object: 'list',
+        data: MODEL_CATALOG.map((m: { id: string; displayName: string; provider: string; capabilities: string[]; contextWindow: number }) => ({
+          id: m.id,
+          object: 'model',
+          owned_by: m.provider,
+          display_name: m.displayName,
+        })),
+      });
+    });
+
+    // Dashboard
+    app.get('/api/dashboard', (_req, res) => {
+      res.json(database.getDashboardData());
+    });
+
+    // Costs
+    app.get('/api/costs', (req, res) => {
+      const days = parseInt(req.query.days as string) || 30;
+      res.json(database.getCostSummary(Math.min(days, 365)));
+    });
+
+    // Budget
+    app.get('/api/budget', (_req, res) => {
+      const config = database.getBudgetConfig();
+      const status = database.getBudgetStatus(config);
+      res.json({ config, status });
+    });
+
+    app.put('/api/budget', (req, res) => {
+      const { dailyBudget, monthlyBudget, perRequestCap } = req.body;
+      const updates: Record<string, number> = {};
+      if (dailyBudget !== undefined) updates.dailyBudget = dailyBudget;
+      if (monthlyBudget !== undefined) updates.monthlyBudget = monthlyBudget;
+      if (perRequestCap !== undefined) updates.perRequestCap = perRequestCap;
+      database.setBudgetConfig(updates);
+      const config = database.getBudgetConfig();
+      const status = database.getBudgetStatus(config);
+      res.json({ config, status });
+    });
+
+    // Cache clear
+    app.delete('/api/cache', (_req, res) => {
+      const cleared = database.clearCache();
+      res.json({ cleared });
+    });
+
+    // Rate limit
+    app.get('/api/rate-limit', (req, res) => {
+      const ip = req.ip || '127.0.0.1';
+      res.json({ ip, ...checkRateLimit(ip) });
+    });
+
+    // Chat completions with validation
+    app.post('/v1/chat/completions', (req, res) => {
+      if (!req.body.messages || !Array.isArray(req.body.messages)) {
+        return res.status(400).json({
+          error: { message: 'messages is required and must be an array', type: 'invalid_request_error' },
+        });
+      }
+      // In test mode, we don't actually call providers
+      res.json({
+        id: 'test-response',
+        object: 'chat.completion',
+        created: Math.floor(Date.now() / 1000),
+        model: req.body.model || 'gpt-4o-mini',
+        choices: [{
+          index: 0,
+          message: { role: 'assistant', content: 'Test response' },
+          finish_reason: 'stop',
+        }],
+        usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
+        cost: { inputCost: 0, outputCost: 0, totalCost: 0, currency: 'USD', model: 'test', provider: 'openai', savingsVsGpt4: 0 },
+      });
+    });
+  });
+
+  describe('GET /health', () => {
+    it('should return health status', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/health');
+      expect(res.status).toBe(200);
+      expect(res.body.status).toBe('ok');
+      expect(res.body.version).toBe('1.0.0');
+      expect(res.body.timestamp).toBeTruthy();
+    });
+  });
+
+  describe('GET /v1/models', () => {
+    it('should return list of models', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/v1/models');
+      expect(res.status).toBe(200);
+      expect(res.body.object).toBe('list');
+      expect(Array.isArray(res.body.data)).toBe(true);
+      expect(res.body.data.length).toBeGreaterThan(0);
+    });
+
+    it('should include model details', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/v1/models');
+      const model = res.body.data[0];
+      expect(model.id).toBeTruthy();
+      expect(model.object).toBe('model');
+      expect(model.owned_by).toBeTruthy();
+    });
+  });
+
+  describe('POST /v1/chat/completions', () => {
+    it('should reject requests without messages', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app)
+        .post('/v1/chat/completions')
+        .send({});
+      expect(res.status).toBe(400);
+      expect(res.body.error.message).toContain('messages');
+    });
+
+    it('should reject requests with non-array messages', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app)
+        .post('/v1/chat/completions')
+        .send({ messages: 'not an array' });
+      expect(res.status).toBe(400);
+    });
+
+    it('should accept valid requests', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app)
+        .post('/v1/chat/completions')
+        .send({
+          messages: [{ role: 'user', content: 'Hello' }],
+        });
+      expect(res.status).toBe(200);
+      expect(res.body.object).toBe('chat.completion');
+      expect(res.body.choices[0].message.content).toBeTruthy();
+    });
+
+    it('should include OpenAI-compatible response fields', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app)
+        .post('/v1/chat/completions')
+        .send({
+          messages: [{ role: 'user', content: 'Hello' }],
+        });
+      expect(res.body.id).toBeTruthy();
+      expect(res.body.created).toBeTruthy();
+      expect(res.body.model).toBeTruthy();
+      expect(res.body.usage).toBeTruthy();
+      expect(res.body.usage.prompt_tokens).toBeGreaterThanOrEqual(0);
+      expect(res.body.usage.completion_tokens).toBeGreaterThanOrEqual(0);
+      expect(res.body.usage.total_tokens).toBeGreaterThanOrEqual(0);
+    });
+  });
+
+  describe('GET /api/dashboard', () => {
+    it('should return dashboard data', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/api/dashboard');
+      expect(res.status).toBe(200);
+      expect(res.body.overview).toBeTruthy();
+      expect(res.body.budget).toBeTruthy();
+      expect(Array.isArray(res.body.recentRequests)).toBe(true);
+      expect(Array.isArray(res.body.costByDay)).toBe(true);
+      expect(Array.isArray(res.body.modelDistribution)).toBe(true);
+    });
+  });
+
+  describe('GET /api/costs', () => {
+    it('should return cost summary', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/api/costs');
+      expect(res.status).toBe(200);
+      expect(res.body.totalSpent).toBeGreaterThanOrEqual(0);
+      expect(res.body.totalRequests).toBeGreaterThanOrEqual(0);
+    });
+
+    it('should accept days parameter', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/api/costs?days=7');
+      expect(res.status).toBe(200);
+    });
+  });
+
+  describe('GET /api/budget', () => {
+    it('should return budget config and status', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/api/budget');
+      expect(res.status).toBe(200);
+      expect(res.body.config).toBeTruthy();
+      expect(res.body.status).toBeTruthy();
+      expect(res.body.status.daily).toBeTruthy();
+      expect(res.body.status.monthly).toBeTruthy();
+    });
+  });
+
+  describe('PUT /api/budget', () => {
+    it('should update budget configuration', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app)
+        .put('/api/budget')
+        .send({ dailyBudget: 15 });
+      expect(res.status).toBe(200);
+      expect(res.body.config.dailyBudget).toBe(15);
+    });
+  });
+
+  describe('DELETE /api/cache', () => {
+    it('should clear cache', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).delete('/api/cache');
+      expect(res.status).toBe(200);
+      expect(typeof res.body.cleared).toBe('number');
+    });
+  });
+
+  describe('GET /api/rate-limit', () => {
+    it('should return rate limit status', async () => {
+      const { default: supertest } = await import('supertest');
+      const res = await supertest(app).get('/api/rate-limit');
+      expect(res.status).toBe(200);
+      expect(typeof res.body.allowed).toBe('boolean');
+      expect(typeof res.body.remaining).toBe('number');
+    });
+  });
+});

package/tests/setup.ts

@@ -0,0 +1,15 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+
+// Set test environment variables before importing any modules
+process.env.NODE_ENV = 'test';
+process.env.PORT = '3211';
+process.env.ENCRYPTION_KEY = 'test_key_for_testing_exactly_32_chars!!';
+process.env.OPENAI_API_KEY = 'test-openai-key';
+process.env.ANTHROPIC_API_KEY = 'test-anthropic-key';
+process.env.GOOGLE_API_KEY = 'test-google-key';
+process.env.DEEPSEEK_API_KEY = 'test-deepseek-key';
+process.env.DEFAULT_DAILY_BUDGET = '10';
+process.env.DEFAULT_MONTHLY_BUDGET = '200';
+process.env.DEFAULT_PER_REQUEST_CAP = '1';
+process.env.CACHE_TTL = '3600';
+process.env.RATE_LIMIT_MAX_REQUESTS = '100';

package/tests/task-classifier.test.ts

@@ -0,0 +1,117 @@
+import { describe, it, expect } from 'vitest';
+import './setup';
+import { classifyTask, estimateTokens } from '../src/task-classifier';
+import type { ChatMessage } from '../src/types';
+
+describe('classifyTask', () => {
+  it('should classify code tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Write a function that sorts an array using quicksort' },
+    ];
+    expect(classifyTask(messages)).toBe('code');
+  });
+
+  it('should classify code tasks with system instructions', () => {
+    const messages: ChatMessage[] = [
+      { role: 'system', content: 'You are a helpful coding assistant' },
+      { role: 'user', content: 'Implement a binary search algorithm in TypeScript' },
+    ];
+    expect(classifyTask(messages)).toBe('code');
+  });
+
+  it('should classify creative tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Write a poem about the ocean at sunset' },
+    ];
+    expect(classifyTask(messages)).toBe('creative');
+  });
+
+  it('should classify translation tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Translate this sentence to French' },
+    ];
+    expect(classifyTask(messages)).toBe('translation');
+  });
+
+  it('should classify summarization tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Summarize this article in a few bullet points' },
+    ];
+    expect(classifyTask(messages)).toBe('summarization');
+  });
+
+  it('should classify reasoning tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Analyze the pros and cons of microservices vs monolithic architecture' },
+    ];
+    expect(classifyTask(messages)).toBe('reasoning');
+  });
+
+  it('should classify math tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Calculate the derivative of x^3 + 2x^2 - 5x + 3' },
+    ];
+    expect(classifyTask(messages)).toBe('math');
+  });
+
+  it('should default to chat for simple greetings', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Hello, how are you today?' },
+    ];
+    expect(classifyTask(messages)).toBe('chat');
+  });
+
+  it('should classify analysis tasks', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Analyze the trends in this data and provide insights' },
+    ];
+    expect(classifyTask(messages)).toBe('analysis');
+  });
+
+  it('should handle multiple messages', () => {
+    const messages: ChatMessage[] = [
+      { role: 'system', content: 'You help with code' },
+      { role: 'user', content: 'Can you fix the bug in my Python script?' },
+      { role: 'assistant', content: 'Sure, show me the code' },
+      { role: 'user', content: 'Here is the code, the debug says error on line 5' },
+    ];
+    // Multiple code-related keywords should make this 'code'
+    const result = classifyTask(messages);
+    expect(['code', 'chat']).toContain(result);
+  });
+
+  it('should return a valid TaskType for empty messages', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'xyzabc123' },
+    ];
+    const result = classifyTask(messages);
+    expect(result).toBeTruthy();
+  });
+});
+
+describe('estimateTokens', () => {
+  it('should estimate tokens based on character count', () => {
+    const messages: ChatMessage[] = [
+      { role: 'user', content: 'Hello world' }, // 11 chars
+    ];
+    const tokens = estimateTokens(messages);
+    // ~4 chars per token => ~3 tokens for 11 chars
+    expect(tokens).toBeGreaterThan(0);
+    expect(tokens).toBeLessThan(100);
+  });
+
+  it('should increase with longer messages', () => {
+    const short = estimateTokens([{ role: 'user', content: 'Hi' }]);
+    const long = estimateTokens([{ role: 'user', content: 'This is a much longer message with many more words and characters in it' }]);
+    expect(long).toBeGreaterThan(short);
+  });
+
+  it('should account for multiple messages', () => {
+    const one = estimateTokens([{ role: 'user', content: 'Hello' }]);
+    const two = estimateTokens([
+      { role: 'user', content: 'Hello' },
+      { role: 'assistant', content: 'World' },
+    ]);
+    expect(two).toBeGreaterThan(one);
+  });
+});

package/tsconfig.json

@@ -0,0 +1,25 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "commonjs",
+    "lib": ["ES2022"],
+    "outDir": "./dist",
+    "rootDir": "./src",
+    "strict": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true,
+    "forceConsistentCasingInFileNames": true,
+    "resolveJsonModule": true,
+    "declaration": true,
+    "declarationMap": true,
+    "sourceMap": true,
+    "noUnusedLocals": true,
+    "noUnusedParameters": true,
+    "noImplicitReturns": true,
+    "noFallthroughCasesInSwitch": true,
+    "exactOptionalPropertyTypes": false,
+    "noUncheckedIndexedAccess": true
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "dist", "tests"]
+}

package/vitest.config.ts

@@ -0,0 +1,15 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    globals: true,
+    environment: 'node',
+    include: ['tests/**/*.test.ts'],
+    coverage: {
+      provider: 'v8',
+      include: ['src/**/*.ts'],
+      exclude: ['src/server.ts'],
+    },
+    testTimeout: 10000,
+  },
+});