server.js

const express = require('express');
const http = require('http');
const proxy = require('express-http-proxy');
const socketIo = require('socket.io');
const moment = require('moment-timezone');

const app = express();
const server = http.createServer(app);
const io = socketIo(server);

const targetUrl = 'https://api.openai.com';
const openaiKey = process.env.OPENAI_KEY;
const port = 7860;
const baseUrl = getExternalUrl(process.env.SPACE_ID);

let requestDetails = [];
let apiKeyDetails = {}; // Track API key usage and first use timestamp

// Allowed models array
const allowedModels = [
    'gpt-3.5-turbo-0125',
    'gpt-3.5-turbo',
    'gpt-3.5-turbo-1106',
    'gpt-3.5-turbo-instruct',
    'gpt-3.5-turbo-16k',
    'gpt-3.5-turbo-0613',
    'gpt-3.5-turbo-16k-0613'
];

io.on('connection', (socket) => {
    console.log('A user connected to the websocket for live logs.');
});

function logAndEmit(message, includeTotal = false) {
    if (includeTotal) {
        const totalRequests = getTotalRequests();
        message += ` Total requests so far: ${totalRequests}`;
    }
    console.log(message);
    io.emit('log', message);
}

function getTotalRequests() {
    return Object.values(apiKeyDetails).reduce((acc, details) => acc + details.count, 0);
}

function authenticateApiKey(req, res, next) {
    const receivedApiKey = req.headers['authorization'];
    if (!receivedApiKey) {
        return res.status(401).send('Unauthorized: No API key provided');
    }

    const validApiKeys = (process.env.SECRET_API_KEYS || '').split(',');
    const apiKeyMatch = validApiKeys.find(key => `Bearer ${key.trim()}` === receivedApiKey);

    if (!apiKeyMatch) {
        return res.status(401).send('Unauthorized: API key is invalid');
    }

    const displayKey = apiKeyMatch.trim().substring(0, 7);

    if (!apiKeyDetails[apiKeyMatch]) {
        apiKeyDetails[apiKeyMatch] = {
            count: 0,
            firstUse: moment()
        };
    }

    if (moment().diff(apiKeyDetails[apiKeyMatch].firstUse, 'weeks') >= 1) {
        return res.status(429).send(`API Key ${displayKey}... is blocked after one week of use.`);
    }

    apiKeyDetails[apiKeyMatch].count++;

    const logMessage = `API Key ${displayKey}... used ${apiKeyDetails[apiKeyMatch].count} times so far.`;
    logAndEmit(logMessage, true);

    req.apiKey = apiKeyMatch;
    next();
}

function validateModel(req, res, next) {
    const pathComponents = req.path.split('/');
    const modelUsed = pathComponents.filter(component => allowedModels.includes(component))[0];

    if (!modelUsed) {
        return res.status(403).send("Access to this model is not allowed.");
    }

    next();
}

app.use('/api', authenticateApiKey, validateModel, proxy(targetUrl, {
    proxyReqPathResolver: function(req) {
        return require('url').parse(req.url).path;
    },
    proxyReqOptDecorator: function(proxyReqOpts, srcReq) {
        proxyReqOpts.headers['Authorization'] = 'Bearer ' + openaiKey;
        return proxyReqOpts;
    }
}));

// Additional server endpoints remain unchanged

server.listen(port, () => {
    const message = `Reverse proxy server and WebSocket running on port ${port}`;
    logAndEmit(message, true); // true to include total requests in the log
});