@theihtisham/budget-llm 1.0.0

package/src/server.ts

@@ -0,0 +1,754 @@
+import express from 'express';
+import cors from 'cors';
+import helmet from 'helmet';
+import compression from 'compression';
+import { env, log } from './config';
+import { handleProxyRequest } from './proxy';
+import {
+  getCostSummary,
+  getDashboardData,
+  getBudgetStatus,
+  getBudgetConfig,
+  setBudgetConfig,
+  clearCache,
+  initDb,
+  closeDb,
+} from './database';
+import { checkRateLimit } from './rate-limiter';
+import {
+  BudgetExceededError,
+  RateLimitError,
+  NoProviderAvailableError,
+} from './types';
+import type { ProxyRequest, BudgetConfig } from './types';
+
+const app = express();
+
+// ---- Middleware ----
+
+app.use(helmet({ contentSecurityPolicy: false }));
+app.use(compression());
+app.use(cors());
+app.use(express.json({ limit: '10mb' }));
+
+// Strip prompt content from logs for security
+app.use((req, _res, next) => {
+  // Only log metadata, never the prompt body
+  log.info(`${req.method} ${req.path} from ${req.ip}`);
+  next();
+});
+
+// ---- Health Check ----
+
+app.get('/health', (_req, res) => {
+  res.json({ status: 'ok', version: '1.0.0', timestamp: new Date().toISOString() });
+});
+
+// ---- OpenAI-Compatible Chat Completions Endpoint ----
+
+app.post('/v1/chat/completions', async (req, res) => {
+  try {
+    const body = req.body as ProxyRequest;
+
+    // Validate required fields
+    if (!body.messages || !Array.isArray(body.messages)) {
+      res.status(400).json({
+        error: {
+          message: 'messages is required and must be an array',
+          type: 'invalid_request_error',
+        },
+      });
+      return;
+    }
+
+    const clientIp = req.ip || req.socket.remoteAddress || 'unknown';
+    const response = await handleProxyRequest(body, clientIp);
+
+    res.json(response);
+  } catch (err) {
+    handleProxyError(err, res);
+  }
+});
+
+// Also support without /v1 prefix for convenience
+app.post('/chat/completions', async (req, res) => {
+  try {
+    const body = req.body as ProxyRequest;
+
+    if (!body.messages || !Array.isArray(body.messages)) {
+      res.status(400).json({
+        error: {
+          message: 'messages is required and must be an array',
+          type: 'invalid_request_error',
+        },
+      });
+      return;
+    }
+
+    const clientIp = req.ip || req.socket.remoteAddress || 'unknown';
+    const response = await handleProxyRequest(body, clientIp);
+
+    res.json(response);
+  } catch (err) {
+    handleProxyError(err, res);
+  }
+});
+
+// ---- Models Endpoint (OpenAI-compatible) ----
+
+app.get('/v1/models', (_req, res) => {
+  const { MODEL_CATALOG } = require('./config');
+  const models = MODEL_CATALOG.map((m: { id: string; displayName: string; provider: string; capabilities: string[]; contextWindow: number }) => ({
+    id: m.id,
+    object: 'model',
+    created: 1700000000,
+    owned_by: m.provider,
+    display_name: m.displayName,
+    capabilities: m.capabilities,
+    context_window: m.contextWindow,
+  }));
+
+  res.json({
+    object: 'list',
+    data: models,
+  });
+});
+
+app.get('/models', (_req, res) => {
+  const { MODEL_CATALOG } = require('./config');
+  const models = MODEL_CATALOG.map((m: { id: string; displayName: string; provider: string; capabilities: string[]; contextWindow: number }) => ({
+    id: m.id,
+    object: 'model',
+    created: 1700000000,
+    owned_by: m.provider,
+    display_name: m.displayName,
+    capabilities: m.capabilities,
+    context_window: m.contextWindow,
+  }));
+
+  res.json({
+    object: 'list',
+    data: models,
+  });
+});
+
+// ---- Dashboard API ----
+
+app.get('/api/dashboard', (_req, res) => {
+  try {
+    const data = getDashboardData();
+    res.json(data);
+  } catch (err) {
+    res.status(500).json({ error: 'Failed to load dashboard data' });
+  }
+});
+
+app.get('/api/costs', (req, res) => {
+  try {
+    const days = parseInt(req.query.days as string) || 30;
+    const summary = getCostSummary(Math.min(days, 365));
+    res.json(summary);
+  } catch (err) {
+    res.status(500).json({ error: 'Failed to load cost data' });
+  }
+});
+
+app.get('/api/budget', (_req, res) => {
+  try {
+    const config = getBudgetConfig();
+    const status = getBudgetStatus(config);
+    res.json({ config, status });
+  } catch (err) {
+    res.status(500).json({ error: 'Failed to load budget data' });
+  }
+});
+
+app.put('/api/budget', (req, res) => {
+  try {
+    const updates = req.body as Partial<BudgetConfig>;
+
+    // Validate inputs
+    if (updates.dailyBudget !== undefined && (typeof updates.dailyBudget !== 'number' || updates.dailyBudget <= 0)) {
+      res.status(400).json({ error: 'dailyBudget must be a positive number' });
+      return;
+    }
+    if (updates.monthlyBudget !== undefined && (typeof updates.monthlyBudget !== 'number' || updates.monthlyBudget <= 0)) {
+      res.status(400).json({ error: 'monthlyBudget must be a positive number' });
+      return;
+    }
+    if (updates.perRequestCap !== undefined && (typeof updates.perRequestCap !== 'number' || updates.perRequestCap <= 0)) {
+      res.status(400).json({ error: 'perRequestCap must be a positive number' });
+      return;
+    }
+
+    setBudgetConfig(updates);
+    const config = getBudgetConfig();
+    const status = getBudgetStatus(config);
+    res.json({ config, status });
+  } catch (err) {
+    res.status(500).json({ error: 'Failed to update budget' });
+  }
+});
+
+app.get('/api/rate-limit', (req, res) => {
+  const ip = req.ip || req.socket.remoteAddress || 'unknown';
+  const status = checkRateLimit(ip);
+  res.json({ ip, ...status });
+});
+
+app.delete('/api/cache', (_req, res) => {
+  try {
+    const cleared = clearCache();
+    res.json({ cleared, message: `Cleared ${cleared} cache entries` });
+  } catch (err) {
+    res.status(500).json({ error: 'Failed to clear cache' });
+  }
+});
+
+// ---- Dashboard HTML ----
+
+app.get('/dashboard', (_req, res) => {
+  res.setHeader('Content-Type', 'text/html');
+  res.send(getDashboardHtml());
+});
+
+app.get('/', (_req, res) => {
+  res.setHeader('Content-Type', 'text/html');
+  res.send(getLandingHtml());
+});
+
+// ---- Error Handler ----
+
+function handleProxyError(err: unknown, res: express.Response): void {
+  if (err instanceof BudgetExceededError) {
+    res.status(429).json({
+      error: {
+        message: err.message,
+        type: 'budget_exceeded',
+        budget_type: err.type,
+        spent: err.spent,
+        limit: err.limit,
+      },
+    });
+    return;
+  }
+
+  if (err instanceof RateLimitError) {
+    res.status(429).json({
+      error: {
+        message: err.message,
+        type: 'rate_limit_exceeded',
+        retry_after_ms: err.resetMs,
+      },
+    });
+    return;
+  }
+
+  if (err instanceof NoProviderAvailableError) {
+    res.status(503).json({
+      error: {
+        message: err.message,
+        type: 'no_provider_available',
+      },
+    });
+    return;
+  }
+
+  if (err instanceof Error) {
+    // Don't leak internal error details in production
+    const message =
+      env.NODE_ENV === 'production'
+        ? 'Internal server error'
+        : err.message;
+    res.status(500).json({
+      error: {
+        message,
+        type: 'internal_error',
+      },
+    });
+    return;
+  }
+
+  res.status(500).json({
+    error: { message: 'Unknown error', type: 'internal_error' },
+  });
+}
+
+// ---- Dashboard HTML ----
+
+function getLandingHtml(): string {
+  return `<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>BudgetLLM - Smart LLM Cost Optimizer</title>
+  <style>
+    * { margin: 0; padding: 0; box-sizing: border-box; }
+    body {
+      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', system-ui, sans-serif;
+      background: linear-gradient(135deg, #0f172a 0%, #1e293b 100%);
+      color: #e2e8f0;
+      min-height: 100vh;
+      display: flex;
+      align-items: center;
+      justify-content: center;
+    }
+    .container {
+      max-width: 700px;
+      text-align: center;
+      padding: 2rem;
+    }
+    h1 {
+      font-size: 3rem;
+      font-weight: 800;
+      background: linear-gradient(135deg, #38bdf8, #818cf8, #c084fc);
+      -webkit-background-clip: text;
+      -webkit-text-fill-color: transparent;
+      margin-bottom: 1rem;
+    }
+    .tagline {
+      font-size: 1.2rem;
+      color: #94a3b8;
+      margin-bottom: 2.5rem;
+      line-height: 1.6;
+    }
+    .endpoint {
+      background: #1e293b;
+      border: 1px solid #334155;
+      border-radius: 12px;
+      padding: 1.5rem;
+      margin-bottom: 1.5rem;
+      text-align: left;
+    }
+    .endpoint h3 { color: #38bdf8; margin-bottom: 0.5rem; font-size: 0.9rem; text-transform: uppercase; letter-spacing: 0.05em; }
+    .endpoint code {
+      display: block;
+      background: #0f172a;
+      padding: 1rem;
+      border-radius: 8px;
+      font-family: 'Fira Code', 'Cascadia Code', monospace;
+      font-size: 0.85rem;
+      color: #a5f3fc;
+      overflow-x: auto;
+      white-space: pre-wrap;
+    }
+    .links {
+      display: flex;
+      gap: 1rem;
+      justify-content: center;
+      margin-top: 2rem;
+    }
+    .links a {
+      background: linear-gradient(135deg, #6366f1, #8b5cf6);
+      color: white;
+      padding: 0.75rem 2rem;
+      border-radius: 10px;
+      text-decoration: none;
+      font-weight: 600;
+      transition: transform 0.2s, box-shadow 0.2s;
+    }
+    .links a:hover {
+      transform: translateY(-2px);
+      box-shadow: 0 8px 25px rgba(99, 102, 241, 0.4);
+    }
+    .links a.secondary {
+      background: #334155;
+    }
+    .links a.secondary:hover {
+      box-shadow: 0 8px 25px rgba(51, 65, 85, 0.6);
+    }
+  </style>
+</head>
+<body>
+  <div class="container">
+    <h1>BudgetLLM</h1>
+    <p class="tagline">Cut your AI costs by 60% &mdash; one API endpoint that<br>automatically picks the cheapest model for every request.</p>
+    <div class="endpoint">
+      <h3>OpenAI-Compatible Endpoint</h3>
+      <code>POST http://localhost:${env.PORT}/v1/chat/completions
+
+{
+  "messages": [
+    {"role": "user", "content": "Hello!"}
+  ]
+}</code>
+    </div>
+    <div class="endpoint">
+      <h3>With Task Hint</h3>
+      <code>{
+  "messages": [
+    {"role": "user", "content": "Write a function to sort an array"}
+  ],
+  "task_type": "code"
+}</code>
+    </div>
+    <div class="links">
+      <a href="/dashboard">Dashboard</a>
+      <a href="/v1/models" class="secondary">Models</a>
+      <a href="/api/costs" class="secondary">Costs API</a>
+    </div>
+  </div>
+</body>
+</html>`;
+}
+
+function getDashboardHtml(): string {
+  return `<!DOCTYPE html>
+<html lang="en">
+<head>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>BudgetLLM Dashboard</title>
+  <style>
+    * { margin: 0; padding: 0; box-sizing: border-box; }
+    body {
+      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', system-ui, sans-serif;
+      background: #0f172a;
+      color: #e2e8f0;
+      min-height: 100vh;
+    }
+    .header {
+      background: linear-gradient(135deg, #1e293b, #0f172a);
+      border-bottom: 1px solid #334155;
+      padding: 1.5rem 2rem;
+      display: flex;
+      justify-content: space-between;
+      align-items: center;
+    }
+    .header h1 {
+      font-size: 1.5rem;
+      background: linear-gradient(135deg, #38bdf8, #818cf8);
+      -webkit-background-clip: text;
+      -webkit-text-fill-color: transparent;
+    }
+    .header .links { display: flex; gap: 1rem; }
+    .header a { color: #94a3b8; text-decoration: none; font-size: 0.9rem; }
+    .header a:hover { color: #38bdf8; }
+    .grid {
+      display: grid;
+      grid-template-columns: repeat(auto-fit, minmax(240px, 1fr));
+      gap: 1rem;
+      padding: 1.5rem 2rem;
+    }
+    .card {
+      background: #1e293b;
+      border: 1px solid #334155;
+      border-radius: 12px;
+      padding: 1.25rem;
+    }
+    .card h3 {
+      font-size: 0.8rem;
+      text-transform: uppercase;
+      letter-spacing: 0.05em;
+      color: #64748b;
+      margin-bottom: 0.5rem;
+    }
+    .card .value {
+      font-size: 1.8rem;
+      font-weight: 700;
+    }
+    .card .value.green { color: #4ade80; }
+    .card .value.blue { color: #38bdf8; }
+    .card .value.purple { color: #a78bfa; }
+    .card .value.yellow { color: #facc15; }
+    .budget-bar {
+      width: 100%;
+      height: 8px;
+      background: #334155;
+      border-radius: 4px;
+      margin-top: 0.5rem;
+      overflow: hidden;
+    }
+    .budget-bar .fill {
+      height: 100%;
+      border-radius: 4px;
+      transition: width 0.5s;
+    }
+    .fill.green { background: #4ade80; }
+    .fill.yellow { background: #facc15; }
+    .fill.red { background: #f87171; }
+    .section {
+      padding: 0 2rem 2rem;
+    }
+    .section h2 {
+      font-size: 1.1rem;
+      margin-bottom: 1rem;
+      color: #94a3b8;
+    }
+    table {
+      width: 100%;
+      border-collapse: collapse;
+      background: #1e293b;
+      border-radius: 12px;
+      overflow: hidden;
+      border: 1px solid #334155;
+    }
+    th {
+      text-align: left;
+      padding: 0.75rem 1rem;
+      font-size: 0.8rem;
+      text-transform: uppercase;
+      letter-spacing: 0.05em;
+      color: #64748b;
+      border-bottom: 1px solid #334155;
+    }
+    td {
+      padding: 0.75rem 1rem;
+      font-size: 0.9rem;
+      border-bottom: 1px solid #1e293b;
+    }
+    tr:last-child td { border-bottom: none; }
+    .badge {
+      display: inline-block;
+      padding: 2px 8px;
+      border-radius: 6px;
+      font-size: 0.75rem;
+      font-weight: 600;
+    }
+    .badge.cached { background: #065f46; color: #6ee7b7; }
+    .badge.live { background: #1e3a5f; color: #7dd3fc; }
+    .loading {
+      text-align: center;
+      padding: 3rem;
+      color: #64748b;
+    }
+    .chart-container {
+      background: #1e293b;
+      border: 1px solid #334155;
+      border-radius: 12px;
+      padding: 1.5rem;
+      margin-bottom: 1.5rem;
+    }
+    .bar-chart { display: flex; align-items: flex-end; gap: 4px; height: 120px; }
+    .bar {
+      flex: 1;
+      background: linear-gradient(to top, #6366f1, #818cf8);
+      border-radius: 4px 4px 0 0;
+      min-width: 8px;
+      position: relative;
+      transition: height 0.5s;
+    }
+    .bar:hover { filter: brightness(1.2); }
+    .bar-label {
+      position: absolute;
+      bottom: -20px;
+      left: 50%;
+      transform: translateX(-50%);
+      font-size: 0.65rem;
+      color: #64748b;
+      white-space: nowrap;
+    }
+    .refresh-btn {
+      background: #334155;
+      border: none;
+      color: #e2e8f0;
+      padding: 0.5rem 1rem;
+      border-radius: 8px;
+      cursor: pointer;
+      font-size: 0.85rem;
+    }
+    .refresh-btn:hover { background: #475569; }
+  </style>
+</head>
+<body>
+  <div class="header">
+    <h1>BudgetLLM Dashboard</h1>
+    <div>
+      <button class="refresh-btn" onclick="loadData()">Refresh</button>
+      <a href="/" style="color:#94a3b8;text-decoration:none;margin-left:1rem;font-size:0.9rem">Home</a>
+    </div>
+  </div>
+
+  <div class="grid" id="stats">
+    <div class="card">
+      <h3>Total Spent</h3>
+      <div class="value blue" id="total-spent">$0.00</div>
+    </div>
+    <div class="card">
+      <h3>Total Saved vs GPT-4</h3>
+      <div class="value green" id="total-saved">$0.00</div>
+    </div>
+    <div class="card">
+      <h3>Total Requests</h3>
+      <div class="value purple" id="total-requests">0</div>
+    </div>
+    <div class="card">
+      <h3>Cache Hit Rate</h3>
+      <div class="value yellow" id="cache-rate">0%</div>
+    </div>
+  </div>
+
+  <div class="grid">
+    <div class="card">
+      <h3>Daily Budget</h3>
+      <div class="value blue" id="daily-spent">$0.00 / $10.00</div>
+      <div class="budget-bar"><div class="fill green" id="daily-bar" style="width:0%"></div></div>
+    </div>
+    <div class="card">
+      <h3>Monthly Budget</h3>
+      <div class="value blue" id="monthly-spent">$0.00 / $200.00</div>
+      <div class="budget-bar"><div class="fill green" id="monthly-bar" style="width:0%"></div></div>
+    </div>
+  </div>
+
+  <div class="section">
+    <h2>Cost by Day (Last 14 Days)</h2>
+    <div class="chart-container">
+      <div class="bar-chart" id="daily-chart">
+        <div class="loading">Loading...</div>
+      </div>
+    </div>
+  </div>
+
+  <div class="section">
+    <h2>Recent Requests</h2>
+    <table>
+      <thead>
+        <tr>
+          <th>Time</th>
+          <th>Model</th>
+          <th>Provider</th>
+          <th>Task</th>
+          <th>Cost</th>
+          <th>Tokens</th>
+          <th>Status</th>
+        </tr>
+      </thead>
+      <tbody id="requests-table">
+        <tr><td colspan="7" class="loading">Loading...</td></tr>
+      </tbody>
+    </table>
+  </div>
+
+  <div class="section">
+    <h2>Model Distribution</h2>
+    <table>
+      <thead>
+        <tr>
+          <th>Model</th>
+          <th>Requests</th>
+          <th>Cost</th>
+        </tr>
+      </thead>
+      <tbody id="model-table">
+        <tr><td colspan="3" class="loading">Loading...</td></tr>
+      </tbody>
+    </table>
+  </div>
+
+  <script>
+    async function loadData() {
+      try {
+        const resp = await fetch('/api/dashboard');
+        const data = await resp.json();
+
+        document.getElementById('total-spent').textContent = '$' + data.overview.totalSpent.toFixed(4);
+        document.getElementById('total-saved').textContent = '$' + data.overview.totalSaved.toFixed(2);
+        document.getElementById('total-requests').textContent = data.overview.totalRequests.toLocaleString();
+        document.getElementById('cache-rate').textContent = (data.overview.cacheHitRate * 100).toFixed(1) + '%';
+
+        // Budget
+        document.getElementById('daily-spent').textContent =
+          '$' + data.budget.daily.spent.toFixed(2) + ' / $' + data.budget.daily.limit.toFixed(2);
+        document.getElementById('monthly-spent').textContent =
+          '$' + data.budget.monthly.spent.toFixed(2) + ' / $' + data.budget.monthly.limit.toFixed(2);
+
+        const dailyPct = Math.min(100, data.budget.daily.percentUsed);
+        const monthlyPct = Math.min(100, data.budget.monthly.percentUsed);
+
+        const dailyBar = document.getElementById('daily-bar');
+        dailyBar.style.width = dailyPct + '%';
+        dailyBar.className = 'fill ' + (dailyPct > 90 ? 'red' : dailyPct > 70 ? 'yellow' : 'green');
+
+        const monthlyBar = document.getElementById('monthly-bar');
+        monthlyBar.style.width = monthlyPct + '%';
+        monthlyBar.className = 'fill ' + (monthlyPct > 90 ? 'red' : monthlyPct > 70 ? 'yellow' : 'green');
+
+        // Daily chart
+        const chartEl = document.getElementById('daily-chart');
+        const days = data.costByDay.slice(0, 14).reverse();
+        if (days.length > 0) {
+          const maxCost = Math.max(...days.map(d => d.cost), 0.01);
+          chartEl.innerHTML = days.map(d => {
+            const height = Math.max(2, (d.cost / maxCost) * 100);
+            return '<div class="bar" style="height:' + height + '%" title="$' + d.cost.toFixed(4) + ' (' + d.requests + ' reqs)"><span class="bar-label">' + d.date.slice(5) + '</span></div>';
+          }).join('');
+        } else {
+          chartEl.innerHTML = '<div style="color:#64748b;text-align:center;width:100%">No data yet</div>';
+        }
+
+        // Recent requests
+        const reqBody = document.getElementById('requests-table');
+        if (data.recentRequests.length > 0) {
+          reqBody.innerHTML = data.recentRequests.slice(0, 20).map(r =>
+            '<tr>' +
+            '<td>' + new Date(r.time).toLocaleString() + '</td>' +
+            '<td>' + r.model + '</td>' +
+            '<td>' + r.provider + '</td>' +
+            '<td>' + r.taskType + '</td>' +
+            '<td>$' + r.cost.toFixed(4) + '</td>' +
+            '<td>' + r.tokens.toLocaleString() + '</td>' +
+            '<td><span class="badge ' + (r.cached ? 'cached">Cached' : 'live">Live') + '</span></td>' +
+            '</tr>'
+          ).join('');
+        } else {
+          reqBody.innerHTML = '<tr><td colspan="7" class="loading">No requests yet. Send a request to POST /v1/chat/completions</td></tr>';
+        }
+
+        // Model distribution
+        const modelBody = document.getElementById('model-table');
+        if (data.modelDistribution.length > 0) {
+          modelBody.innerHTML = data.modelDistribution.map(m =>
+            '<tr><td>' + m.model + '</td><td>' + m.count.toLocaleString() + '</td><td>$' + m.cost.toFixed(4) + '</td></tr>'
+          ).join('');
+        } else {
+          modelBody.innerHTML = '<tr><td colspan="3" class="loading">No data yet</td></tr>';
+        }
+      } catch (err) {
+        console.error('Dashboard load error:', err);
+      }
+    }
+
+    loadData();
+    setInterval(loadData, 30000);
+  </script>
+</body>
+</html>`;
+}
+
+// ---- Start Server ----
+
+export async function startServer(): Promise<express.Application> {
+  // Initialize database (async for sql.js)
+  await initDb();
+
+  app.listen(env.PORT, () => {
+    log.info(`BudgetLLM proxy running on http://localhost:${env.PORT}`);
+    log.info(`Dashboard: http://localhost:${env.PORT}/dashboard`);
+    log.info(`Endpoint:  POST http://localhost:${env.PORT}/v1/chat/completions`);
+  });
+
+  // Graceful shutdown
+  process.on('SIGTERM', () => {
+    log.info('SIGTERM received, shutting down');
+    closeDb();
+    process.exit(0);
+  });
+
+  process.on('SIGINT', () => {
+    log.info('SIGINT received, shutting down');
+    closeDb();
+    process.exit(0);
+  });
+
+  return app;
+}
+
+// Start if run directly
+if (require.main === module) {
+  startServer().catch((err) => {
+    console.error('Failed to start server:', err);
+    process.exit(1);
+  });
+}