scripts/omniroute.js

/**
 * OmniRoute v2 — The Polyglot Intelligence Engine
 * 
 * Features:
 *   - Tiered Fallback: Big Names (T1) -> Open Weights (T2) -> Free (T3).
 *   - Multi-API Translation: Speaks OpenAI, Anthropic Messages, and Google Gemini.
 *   - Advanced Account Rotation: Round-robin across multiple keys per provider.
 *   - Per-Key Cooldowns: Tracks bans/rate-limits at the individual key level.
 */

const express = require('express');
const fetch = require('node-fetch');
const fs = require('fs');
const path = require('path');

const app = express();
app.use(express.json());

// ─── Load Masterfile ──────────────────────────────────────────────
const configPath = path.join(__dirname, '..', 'config', 'omniroute-masterfile.json');
let masterfile = { providers: {}, model: {}, routing: {} };

function loadConfig() {
  if (fs.existsSync(configPath)) {
    const raw = fs.readFileSync(configPath, 'utf-8');
    const resolved = raw.replace(/\$\{(\w+)\}/g, (_, key) => process.env[key] || '');
    masterfile = JSON.parse(resolved);
    console.log(`📡 OmniRoute: Tiered Masterfile loaded.`);
  }
}
loadConfig();

// ─── Port & Global State ─────────────────────────────────────────
const PORT = process.env.GRAPH_VIEWER_PORT ? 20128 : (process.env.PORT || 20128);
const keyRotation = {}; // { providerName: lastUsedKeyIndex }
const keyCooldowns = {}; // { "provider_key": timestamp }

// ─── Intelligence: The Polyglot Translators ───────────────────────

/**
 * Translates OpenAI-format messages to Anthropic Messages format.
 */
function toAnthropic(messages) {
  const system = messages.find(m => m.role === 'system')?.content || '';
  const contents = messages.filter(m => m.role !== 'system').map(m => ({
    role: m.role === 'assistant' ? 'assistant' : 'user',
    content: m.content
  }));
  return { system, messages: contents };
}

/**
 * Translates OpenAI-format messages to Google Gemini format.
 */
function toGemini(messages) {
  return {
    contents: messages.map(m => ({
      role: m.role === 'assistant' ? 'model' : 'user',
      parts: [{ text: m.content }]
    }))
  };
}

// ─── Logic: Key & Provider Management ───────────────────────────

function getAvailableKeys(providerName) {
  const base = `${providerName.toUpperCase()}_API_KEY`;
  const keys = [];
  if (process.env[base] && process.env[base] !== 'REPLACE_ME') keys.push(process.env[base]);
  for (let i = 2; i <= 10; i++) {
    const k = process.env[`${base}_${i}`];
    if (k && k !== 'REPLACE_ME') keys.push(k);
  }
  return keys;
}

function getKeyCooldown(provider, key) {
  const composite = `${provider}_${key.slice(-5)}`;
  if (keyCooldowns[composite] && Date.now() < keyCooldowns[composite]) return true;
  return false;
}

function setKeyCooldown(provider, key, seconds = 60) {
  const composite = `${provider}_${key.slice(-5)}`;
  keyCooldowns[composite] = Date.now() + (seconds * 1000);
  console.log(`⏸️  Key ${composite} limited for ${seconds}s`);
}

/**
 * Intelligence: Tiered Candidate Finder
 */
function findCandidates(requestedModel) {
  const candidates = [];
  for (const [pName, pConfig] of Object.entries(masterfile.providers)) {
    const keys = getAvailableKeys(pName);
    if (keys.length === 0) continue;

    for (const model of pConfig.models) {
      if (requestedModel && model.id !== requestedModel) continue;
      
      candidates.push({
        provider: pName,
        model: model.id,
        api: pConfig.api,
        baseUrl: pConfig.baseUrl,
        tier: pConfig.tier || 3,
        priority: model.priority || 99,
        keys
      });
    }
  }

  // Tier 1 (Paid) -> Tier 2 (Open Weight) -> Tier 3 (Free)
  candidates.sort((a, b) => a.tier - b.tier || a.priority - b.priority);
  return candidates;
}

// ─── The Main Execution Engine ───────────────────────────────────

app.post('/v1/chat/completions', async (req, res) => {
  const { model: reqModel, messages, ...options } = req.body;
  const chain = reqModel ? [reqModel] : [masterfile.model.primary, masterfile.model.fallback, masterfile.model.review];
  
  for (const targetModel of chain) {
    const candidates = findCandidates(targetModel);

    for (const c of candidates) {
      // Rotation: Start from a fresh key each time
      if (!keyRotation[c.provider]) keyRotation[c.provider] = 0;
      
      for (let i = 0; i < c.keys.length; i++) {
        const keyIdx = (keyRotation[c.provider] + i) % c.keys.length;
        const key = c.keys[keyIdx];

        if (getKeyCooldown(c.provider, key)) continue;

        try {
          console.log(`📡 [Tier ${c.tier}] ${c.provider} -> ${targetModel}`);
          const controller = new AbortController();
          const timeout = setTimeout(() => controller.abort(), 60000);

          let body, url, headers = { 'Authorization': `Bearer ${key}`, 'Content-Type': 'application/json' };

          // Polyglot: Translation Logic
          if (c.api === 'anthropic-messages') {
            const translated = toAnthropic(messages);
            url = `${c.baseUrl}/messages`;
            headers['x-api-key'] = key;
            delete headers['Authorization'];
            body = { model: c.model, ...translated, max_tokens: options.max_tokens || 4096 };
          } else if (c.api === 'google-gemini') {
            const translated = toGemini(messages);
            url = `${c.baseUrl}/${c.model}:generateContent?key=${key}`;
            delete headers['Authorization'];
            body = translated;
          } else {
            // Default: OpenAI Format
            url = `${c.baseUrl}/chat/completions`;
            body = { model: c.model, messages, ...options };
          }

          const response = await fetch(url, {
            method: 'POST',
            headers,
            body: JSON.stringify(body),
            signal: controller.signal
          });
          clearTimeout(timeout);

          if (response.status === 429) {
            setKeyCooldown(c.provider, key, 120);
            continue;
          }

          if (!response.ok) {
            console.warn(`⚠️  ${c.provider} failed (${response.status})`);
            continue;
          }

          const data = await response.json();
          // Update rotation for next time
          keyRotation[c.provider] = (keyIdx + 1) % c.keys.length;
          
          // Return standardized response
          return res.json(data);

        } catch (e) {
          console.error(`❌  ${c.provider} error: ${e.message}`);
          continue;
        }
      }
    }
  }

  res.status(503).json({ error: "All tiers exhausted. Check API keys or cooldowns." });
});

// ─── System Info & Management ────────────────────────────────────

app.get('/v1/models', (req, res) => {
  const list = [];
  for (const [pName, pConfig] of Object.entries(masterfile.providers)) {
    const keys = getAvailableKeys(pName);
    pConfig.models.forEach(m => list.push({
      id: m.id,
      tier: pConfig.tier,
      provider: pName,
      active_keys: keys.length,
      available: !keys.every(k => getKeyCooldown(pName, k))
    }));
  }
  res.json({ data: list.sort((a,b) => a.tier - b.tier) });
});

app.listen(PORT, '0.0.0.0', () => {
  console.log(`🚀 OmniRoute v2: Polyglot Intelligence Engine on port ${PORT}`);
});