diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..883465d --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +node_modules/ +apiKeys.db +*.conf +.env diff --git a/APIMyLlama.js b/APIMyLlama.js index 612b9c9..54c4f08 100644 --- a/APIMyLlama.js +++ b/APIMyLlama.js @@ -1,69 +1,50 @@ -const express = require('express'); -const { initializeDatabase, closeDatabase, getDb } = require('./db'); -const { startServer, askForPort, askForOllamaURL, startCLI } = require('./utils'); -const { setupRoutes } = require('./api'); -const fs = require('fs'); - -const app = express(); -app.use(express.json()); - -console.log('APIMyLlama V2 is being started. Thanks for choosing Gimer Studios.'); - -// Middleware for logging requests -app.use((req, res, next) => { - console.log(`Received a ${req.method} request at ${req.url}`); - next(); -}); - -// Initialize the database -const db = initializeDatabase(); - -// Setup API routes -setupRoutes(app, db); - -// Close the database connection when the application is closed -process.on('SIGINT', () => { - closeDatabase(); -}); - -// Start the server -setTimeout(() => { - if (fs.existsSync('port.conf')) { - fs.readFile('port.conf', 'utf8', (err, data) => { - if (err) { - console.error('Error reading port number from file:', err.message); - askForPort(app, startServer, askForOllamaURL, startCLI, db); - } else { - const port = parseInt(data.trim()); - if (isNaN(port)) { - console.error('Invalid port number in port.conf'); - askForPort(app, startServer, askForOllamaURL, startCLI, db); - } else { - if (fs.existsSync('ollamaURL.conf')) { - fs.readFile('ollamaURL.conf', 'utf8', (err, data) => { - if (err) { - console.error('Error reading Ollama url from file:', err.message); - askForOllamaURL(app, startServer, startCLI, port, db); - } else { - const ollamaURL = data.trim(); - if (typeof ollamaURL !== 'string' || ollamaURL === '') { - console.error('Invalid Ollama url in ollamaURL.conf'); - askForOllamaURL(app, startServer, startCLI, port, db); - } else { - startServer(port, app); - startCLI(db); - } - } - }); - } else { - askForOllamaURL(app, startServer, startCLI, port, db); - } - } - } - }); - } else { - askForPort(app, startServer, askForOllamaURL, startCLI, db); - } -}, 1000); - -module.exports = app; +const express = require('express'); +const db = require('./db'); +const { startServer, resolveConfig, startCLI } = require('./utils'); +const { setupRoutes } = require('./api'); + +const app = express(); + +async function main() { + console.log('APIMyLlama V2 is being started. Thanks for choosing Gimer Studios.'); + + app.use(express.json({ limit: '10mb' })); + + app.use((req, res, next) => { + console.log(`Received a ${req.method} request at ${req.url}`); + next(); + }); + + app.use((req, res, next) => { + res.header('Access-Control-Allow-Origin', '*'); + res.header('Access-Control-Allow-Headers', 'Origin, X-Requested-With, Content-Type, Accept, Authorization'); + res.header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS, DELETE'); + if (req.method === 'OPTIONS') return res.sendStatus(200); + next(); + }); + + await db.initialize(); + + setupRoutes(app); + + const port = await resolveConfig('port number', 'PORT', 'port.conf', '3000'); + await resolveConfig('Ollama server URL', 'OLLAMA_URL', 'ollamaURL.conf', 'http://localhost:11434'); + + startServer(parseInt(port), app); + startCLI(); +} + +process.on('SIGINT', async () => { + console.log('\nShutting down...'); + try { + await db.close(); + } catch {} + process.exit(0); +}); + +main().catch(err => { + console.error('Failed to start:', err); + process.exit(1); +}); + +module.exports = app; diff --git a/README.md b/README.md index 58049c2..0415263 100644 --- a/README.md +++ b/README.md @@ -434,6 +434,13 @@ prompt: Text prompt to generate a response. model: Machine learning model to use for text generation. stream: Boolean indicating whether to stream the response. ``` +## curl Use Case +``` +curl -X POST http://localhost:/generate \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{"prompt":"","model":"","stream":true}' +``` # Support If there are any issues please make a Github Issue Report. To get quicker support join our discord server. -[Discord Server](https://discord.gg/r6XazGtKg7) If there are any feature requests you may request them in the discord server. PLEASE NOTE this project is still in EARLY BETA. diff --git a/api.js b/api.js index d35d31f..469b6d6 100644 --- a/api.js +++ b/api.js @@ -1,130 +1,179 @@ -const { getOllamaURL, sendWebhookNotification } = require('./utils'); const axios = require('axios'); +const db = require('./db'); +const { getOllamaURL, sendWebhookNotification } = require('./utils'); const rateLimits = new Map(); -function setupRoutes(app, db) { - app.use((req, res, next) => rateLimitMiddleware(req, res, next, db)); - app.get('/health', (req, res) => healthCheck(req, res, db)); - app.post('/generate', (req, res) => generateResponse(req, res, db)); +function extractApiKey(req) { + const authHeader = req.headers.authorization; + if (authHeader && authHeader.startsWith('Bearer ')) { + return authHeader.slice(7); + } + return req.body?.apikey || req.query?.apikey || null; } -function rateLimitMiddleware(req, res, next, db) { - const { apikey } = req.body; - if (!apikey) return next(); +async function verifyApiKey(apikey) { + if (!apikey) return null; + try { + return await db.get('SELECT * FROM apiKeys WHERE key = ?', [apikey]); + } catch { + return null; + } +} - db.get('SELECT tokens, last_used, rate_limit, active FROM apiKeys WHERE key = ?', [apikey], (err, row) => { - if (err) { - console.error('Error checking API key for rate limit:', err.message); - return res.status(500).json({ error: 'Internal server error' }); +function setupRoutes(app) { + const healthHandler = async (req, res) => { + const apikey = extractApiKey(req); + if (!apikey) { + return res.status(401).json({ error: 'API key is required (use Authorization: Bearer header or ?apikey= param)' }); } - if (row) { - if (row.active === 0) { - return res.status(403).json({ error: 'API key is deactivated' }); - } - - const currentTime = Date.now(); - const minute = 60000; - const rateLimit = row.rate_limit; - - if (!rateLimits.has(apikey)) { - rateLimits.set(apikey, { tokens: row.tokens, lastUsed: new Date(row.last_used).getTime() }); - } - - const rateLimitInfo = rateLimits.get(apikey); - const timeElapsed = currentTime - rateLimitInfo.lastUsed; - - if (timeElapsed >= minute) { - rateLimitInfo.tokens = rateLimit; - } - - if (rateLimitInfo.tokens > 0) { - rateLimitInfo.tokens -= 1; - rateLimitInfo.lastUsed = currentTime; - rateLimits.set(apikey, rateLimitInfo); - - db.run('UPDATE apiKeys SET tokens = ?, last_used = ? WHERE key = ?', [rateLimitInfo.tokens, new Date(rateLimitInfo.lastUsed).toISOString(), apikey], (err) => { - if (err) { - console.error('Error updating tokens and last_used:', err.message); - return res.status(500).json({ error: 'Internal server error' }); - } - next(); - }); - } else { - return res.status(429).json({ error: 'Rate limit exceeded. Try again later.' }); - } - } else { + const keyInfo = await verifyApiKey(apikey); + if (!keyInfo) { return res.status(403).json({ error: 'Invalid API key' }); } - }); -} -function healthCheck(req, res, db) { - const apikey = req.query.apikey; + let ollamaHealthy = false; + try { + const url = await getOllamaURL(); + await axios.get(`${url}/api/tags`, { timeout: 5000 }); + ollamaHealthy = true; + } catch { + ollamaHealthy = false; + } - if (!apikey) { - return res.status(400).json({ error: 'API key is required' }); - } + res.json({ + status: ollamaHealthy ? 'healthy' : 'degraded', + ollama: ollamaHealthy ? 'reachable' : 'unreachable', + timestamp: new Date().toISOString() + }); + }; + + const generateHandler = async (req, res) => { + const apikey = extractApiKey(req); + if (!apikey) { + return res.status(401).json({ error: 'API key is required (use Authorization: Bearer header)' }); + } - db.get('SELECT key FROM apiKeys WHERE key = ?', [apikey], (err, row) => { - if (err) { - console.error('Error checking API key:', err.message); - return res.status(500).json({ error: 'Internal server error' }); + const keyInfo = await verifyApiKey(apikey); + if (!keyInfo) { + return res.status(403).json({ error: 'Invalid API key' }); } - if (!row) { - console.log('Invalid API key:', apikey); - return res.status(403).json({ error: 'Invalid API Key' }); + if (keyInfo.active === 0) { + return res.status(403).json({ error: 'API key is deactivated' }); } - res.json({ status: 'API is healthy', timestamp: new Date() }); - }); + const rateLimitError = checkRateLimit(apikey, keyInfo); + if (rateLimitError) { + return res.status(429).json({ error: rateLimitError }); + } + + await handleGenerate(req, res, apikey); + }; + + app.get('/v1/health', healthHandler); + app.post('/v1/generate', generateHandler); + app.get('/health', healthHandler); + app.post('/generate', generateHandler); } -async function generateResponse(req, res, db) { - const { apikey, prompt, model, stream, images, raw } = req.body; +function checkRateLimit(apikey, keyInfo) { + const currentTime = Date.now(); + const minute = 60000; + const rateLimit = keyInfo.rate_limit; + + if (!rateLimits.has(apikey)) { + const lastUsed = new Date(keyInfo.last_used).getTime(); + const timeElapsed = currentTime - lastUsed; + const tokens = timeElapsed >= minute ? rateLimit : Math.min(keyInfo.tokens, rateLimit); + rateLimits.set(apikey, { tokens, lastUsed }); + } + + const rateLimitInfo = rateLimits.get(apikey); + const timeElapsed = currentTime - rateLimitInfo.lastUsed; + + if (timeElapsed >= minute) { + rateLimitInfo.tokens = rateLimit; + rateLimitInfo.lastUsed = currentTime; + } + + if (rateLimitInfo.tokens <= 0) { + return 'Rate limit exceeded. Try again later.'; + } + + rateLimitInfo.tokens -= 1; + rateLimitInfo.lastUsed = currentTime; + + db.run('UPDATE apiKeys SET tokens = ?, last_used = ? WHERE key = ?', [ + rateLimitInfo.tokens, + new Date(rateLimitInfo.lastUsed).toISOString(), + apikey + ]).catch(err => console.error('Error updating tokens:', err.message)); - console.log('Request body:', req.body); + return null; +} - if (!apikey) { - return res.status(400).json({ error: 'API key is required' }); +async function handleGenerate(req, res, apikey) { + const { prompt, model, stream, images, raw } = req.body; + + if (!prompt || !model) { + return res.status(400).json({ error: 'Both prompt and model are required' }); } - db.get('SELECT key FROM apiKeys WHERE key = ?', [apikey], async (err, row) => { - if (err) { - console.error('Error checking API key:', err.message); - return res.status(500).json({ error: 'Internal server error' }); + try { + const ollamaURL = await getOllamaURL(); + const OLLAMA_API_URL = `${ollamaURL}/api/generate`; + + if (stream) { + const ollamaResponse = await axios({ + method: 'post', + url: OLLAMA_API_URL, + data: { model, prompt, stream: true, images, raw }, + responseType: 'stream', + timeout: 300000 + }); + + res.setHeader('Content-Type', 'application/x-ndjson'); + ollamaResponse.data.pipe(res); + + ollamaResponse.data.on('end', () => { + logUsage(apikey); + sendWebhook(apikey, prompt, model, stream, images, raw); + }); + } else { + const ollamaResponse = await axios.post(OLLAMA_API_URL, { model, prompt, stream: false, images, raw }, { + timeout: 300000 + }); + + logUsage(apikey); + sendWebhook(apikey, prompt, model, stream, images, raw); + + res.json(ollamaResponse.data); } - if (!row) { - console.log('Invalid API key:', apikey); - return res.status(403).json({ error: 'Invalid API Key' }); + } catch (error) { + if (error.response) { + console.error('Ollama API error:', error.response.status, error.response.data); + res.status(error.response.status).json({ error: 'Ollama API error', detail: error.response.data }); + } else if (error.code === 'ECONNREFUSED') { + console.error('Ollama server is not reachable:', error.message); + res.status(503).json({ error: 'Ollama server is not reachable' }); + } else { + console.error('Error making request to Ollama API:', error.message); + res.status(500).json({ error: 'Error making request to Ollama API' }); } + } +} - try { - const ollamaURL = await getOllamaURL(); - const OLLAMA_API_URL = `${ollamaURL}/api/generate`; - - axios.post(OLLAMA_API_URL, { model, prompt, stream, images, raw }) - .then(response => { - db.run('INSERT INTO apiUsage (key) VALUES (?)', [apikey], (err) => { - if (err) console.error('Error logging API usage:', err.message); - }); - - sendWebhookNotification(db, { apikey, prompt, model, stream, images, raw, timestamp: new Date() }); - - res.json(response.data); - }) - .catch(error => { - console.error('Error making request to Ollama API:', error.message); - res.status(500).json({ error: 'Error making request to Ollama API' }); - }); - } catch (error) { - console.error(error); - res.status(500).json({ error: 'Error retrieving Ollama server port' }); - } +function logUsage(apikey) { + db.run('INSERT INTO apiUsage (key) VALUES (?)', [apikey]) + .catch(err => console.error('Error logging API usage:', err.message)); +} + +function sendWebhook(apikey, prompt, model, stream, images, raw) { + sendWebhookNotification({ + apikey, prompt, model, stream, images, raw, + timestamp: new Date().toISOString() }); } -module.exports = { - setupRoutes -}; +module.exports = { setupRoutes }; diff --git a/db.js b/db.js index 9449859..7dd9abc 100644 --- a/db.js +++ b/db.js @@ -1,101 +1,107 @@ const sqlite3 = require('sqlite3').verbose(); -let db; +class Database { + constructor() { + this.db = null; + } -function initializeDatabase() { - db = new sqlite3.Database('./apiKeys.db', sqlite3.OPEN_READWRITE | sqlite3.OPEN_CREATE, (err) => { - if (err) { - console.error('Error connecting to the database:', err.message); - } else { - console.log('Connected to the apiKeys.db database.'); - createTables(); - } - }); - return db; -} + initialize() { + return new Promise((resolve, reject) => { + const dbPath = process.env.API_KEYS_DB_PATH || './apiKeys.db'; + this.db = new sqlite3.Database(dbPath, sqlite3.OPEN_READWRITE | sqlite3.OPEN_CREATE, (err) => { + if (err) { + console.error('Error connecting to the database:', err.message); + reject(err); + } else { + console.log('Connected to the apiKeys.db database.'); + this.createTables().then(resolve).catch(reject); + } + }); + }); + } - // Function to create the tables even if they do not exist in the database -function createTables() { - db.run(`CREATE TABLE IF NOT EXISTS apiKeys ( - key TEXT PRIMARY KEY, - created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - last_used TIMESTAMP DEFAULT CURRENT_TIMESTAMP, - tokens INTEGER DEFAULT 10, - rate_limit INTEGER DEFAULT 10, - active INTEGER DEFAULT 1, - description TEXT - )`, (err) => { - if (err) { - console.error('Error creating apiKeys table:', err.message); - } - }); + async createTables() { + await this.run(`CREATE TABLE IF NOT EXISTS apiKeys ( + key TEXT PRIMARY KEY, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + last_used TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + tokens INTEGER DEFAULT 10, + rate_limit INTEGER DEFAULT 10, + active INTEGER DEFAULT 1, + description TEXT + )`); - db.run(`CREATE TABLE IF NOT EXISTS apiUsage ( - key TEXT, - timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP - )`, (err) => { - if (err) { - console.error('Error creating apiUsage table:', err.message); - } - }); + await this.run(`CREATE TABLE IF NOT EXISTS apiUsage ( + key TEXT, + timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP + )`); - db.run(`CREATE TABLE IF NOT EXISTS webhooks ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - url TEXT NOT NULL - )`, (err) => { - if (err) { - console.error('Error creating webhooks table:', err.message); - } - }); + await this.run(`CREATE TABLE IF NOT EXISTS webhooks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + url TEXT NOT NULL + )`); - ensureColumns(); -} + await this.ensureColumns(); + } -function ensureColumns() { - db.all("PRAGMA table_info(apiKeys)", (err, rows) => { - if (err) { - console.error('Error checking table info:', err.message); - } else { - const columns = rows.map(row => row.name); - if (!columns.includes('active')) { - db.run("ALTER TABLE apiKeys ADD COLUMN active INTEGER DEFAULT 1", (err) => { - if (err) { - console.error('Error adding active column:', err.message); - } else { - console.log("Added 'active' column to 'apiKeys' table."); - } - }); - } - if (!columns.includes('description')) { - db.run("ALTER TABLE apiKeys ADD COLUMN description TEXT", (err) => { - if (err) { - console.error('Error adding description column:', err.message); - } else { - console.log("Added 'description' column to 'apiKeys' table."); - } - }); - } - } - }); -} + async ensureColumns() { + const rows = await this.all("PRAGMA table_info(apiKeys)"); + const columns = rows.map(row => row.name); -function closeDatabase() { - db.close((err) => { - if (err) { - console.error('Error closing the database connection:', err.message); - } else { - console.log('Closed the database connection.'); + if (!columns.includes('active')) { + await this.run("ALTER TABLE apiKeys ADD COLUMN active INTEGER DEFAULT 1"); + console.log("Added 'active' column to 'apiKeys' table."); } - process.exit(0); - }); -} + if (!columns.includes('description')) { + await this.run("ALTER TABLE apiKeys ADD COLUMN description TEXT"); + console.log("Added 'description' column to 'apiKeys' table."); + } + } -function getDb() { - return db; + get(sql, params = []) { + return new Promise((resolve, reject) => { + this.db.get(sql, params, (err, row) => { + if (err) reject(err); + else resolve(row); + }); + }); + } + + run(sql, params = []) { + return new Promise((resolve, reject) => { + this.db.run(sql, params, function (err) { + if (err) reject(err); + else resolve(this); + }); + }); + } + + all(sql, params = []) { + return new Promise((resolve, reject) => { + this.db.all(sql, params, (err, rows) => { + if (err) reject(err); + else resolve(rows); + }); + }); + } + + close() { + return new Promise((resolve, reject) => { + if (!this.db) { + resolve(); + return; + } + this.db.close((err) => { + if (err) { + console.error('Error closing the database connection:', err.message); + reject(err); + } else { + console.log('Closed the database connection.'); + resolve(); + } + }); + }); + } } -module.exports = { - initializeDatabase, - closeDatabase, - getDb -}; \ No newline at end of file +module.exports = new Database(); diff --git a/ollamaURL.conf b/ollamaURL.conf deleted file mode 100644 index 1916ae8..0000000 --- a/ollamaURL.conf +++ /dev/null @@ -1 +0,0 @@ -http://127.0.0.1:11434 \ No newline at end of file diff --git a/package.json b/package.json index ce812f8..f8ecf8b 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,14 @@ { + "private": true, + "engines": { + "node": ">=18.0.0" + }, + "scripts": { + "start": "node APIMyLlama.js" + }, "dependencies": { - "axios": ">=1.13.5", - "express": ">=5.2.1", + "axios": "^1.13.5", + "express": "^5.2.1", "sqlite3": "^6.0.1" } } diff --git a/utils.js b/utils.js index 9f45334..ba56a6f 100644 --- a/utils.js +++ b/utils.js @@ -2,416 +2,362 @@ const fs = require('fs'); const readline = require('readline'); const crypto = require('crypto'); const axios = require('axios'); -const { getDb } = require('./db'); +const db = require('./db'); let server; let currentPort; let expressApp; +const VALID_URL_PATTERN = /^https?:\/\/[^\s$.?#].[^\s]*$/i; + function startServer(port, app) { currentPort = port; - expressApp = app; // Store the app object + expressApp = app; server = expressApp.listen(currentPort, () => console.log(`Server running on port ${currentPort}`)); } -function askForPort(app, startServerCallback, askForOllamaURLCallback, startCLICallback, db) { - const rl = readline.createInterface({ - input: process.stdin, - output: process.stdout - }); - - rl.question('Enter the port number for the API server: ', (port) => { - fs.writeFile('port.conf', port, (err) => { - if (err) { - console.error('Error saving port number:', err.message); - } else { - console.log(`Port number saved to port.conf: ${port}`); - currentPort = parseInt(port); - askForOllamaURLCallback(app, startServerCallback, startCLICallback, currentPort, db); - } - }); - rl.close(); - }); -} - -function askForOllamaURL(app, startServerCallback, startCLICallback, port, db) { - const rl = readline.createInterface({ - input: process.stdin, - output: process.stdout - }); +async function resolveConfig(name, envVar, confFile, defaultValue) { + if (process.env[envVar]) { + console.log(`${name} set from environment variable ${envVar}`); + return process.env[envVar]; + } - rl.question('Enter the URL for the Ollama server (URL that your Ollama server is running on. By default it is "http://localhost:11434" so if you didnt change anything it should be that.): ', (ollamaURL) => { - fs.writeFile('ollamaURL.conf', ollamaURL, (err) => { - if (err) { - console.error('Error saving Ollama url:', err.message); - } else { - console.log(`Ollama url saved to ollamaURL.conf: ${ollamaURL}`); - startServerCallback(port, app); - startCLICallback(db); - } + try { + const data = await fs.promises.readFile(confFile, 'utf8'); + const value = data.trim(); + if (value) { + console.log(`${name} loaded from ${confFile}: ${value}`); + return value; + } + } catch {} + + return new Promise((resolve) => { + const rl = readline.createInterface({ input: process.stdin, output: process.stdout }); + rl.question(`Enter the ${name} (default: ${defaultValue}): `, (answer) => { + rl.close(); + const value = answer.trim() || defaultValue; + fs.promises.writeFile(confFile, value, 'utf8') + .then(() => console.log(`${name} saved to ${confFile}: ${value}`)) + .catch(err => console.error(`Error saving ${confFile}:`, err.message)); + resolve(value); }); - rl.close(); }); } -function startCLI(db) { +async function startCLI() { const rl = readline.createInterface({ input: process.stdin, output: process.stdout }); - rl.on('line', (input) => { + rl.on('line', async (input) => { const [command, argument, ...rest] = input.trim().split(' '); const description = rest.join(' '); - switch (command) { - case 'generatekey': - generateKey(db); - break; - case 'generatekeys': - generateKeys(db, argument); - break; - case 'listkey': - listKeys(db); - break; - case 'removekey': - removeKey(db, argument); - break; - case 'addkey': - addKey(db, argument); - break; - case 'changeport': - changePort(argument); - break; - case 'changeollamaurl': - changeOllamaURL(argument); - break; - case 'ratelimit': - setRateLimit(db, argument, rest[0]); - break; - case 'addwebhook': - addWebhook(db, argument); - break; - case 'deletewebhook': - deleteWebhook(db, argument); - break; - case 'listwebhooks': - listWebhooks(db); - break; - case 'activatekey': - activateKey(db, argument); - break; - case 'deactivatekey': - deactivateKey(db, argument); - break; - case 'addkeydescription': - addKeyDescription(db, argument, description); - break; - case 'listkeydescription': - listKeyDescription(db, argument); - break; - case 'regeneratekey': - regenerateKey(db, argument); - break; - case 'activateallkeys': - activateAllKeys(db); - break; - case 'deactivateallkeys': - deactivateAllKeys(db); - break; - case 'getkeyinfo': - getKeyInfo(db, argument); - break; - case 'listinactivekeys': - listInactiveKeys(db); - break; - case 'listactivekeys': - listActiveKeys(db); - break; - case 'exit': - rl.close(); - process.exit(0); - break; - default: - console.log('Unknown command'); + try { + switch (command) { + case 'generatekey': + await generateKey(); + break; + case 'generatekeys': + await generateKeys(argument); + break; + case 'listkey': + await listKeys(); + break; + case 'removekey': + await removeKey(argument); + break; + case 'addkey': + await addKey(argument); + break; + case 'changeport': + await changePort(argument); + break; + case 'changeollamaurl': + await changeOllamaURL(argument); + break; + case 'ratelimit': + await setRateLimit(argument, rest[0]); + break; + case 'addwebhook': + await addWebhook(argument); + break; + case 'deletewebhook': + await deleteWebhook(argument); + break; + case 'listwebhooks': + await listWebhooks(); + break; + case 'activatekey': + await activateKey(argument); + break; + case 'deactivatekey': + await deactivateKey(argument); + break; + case 'addkeydescription': + await addKeyDescription(argument, description); + break; + case 'listkeydescription': + await listKeyDescription(argument); + break; + case 'regeneratekey': + await regenerateKey(argument); + break; + case 'activateallkeys': + await activateAllKeys(); + break; + case 'deactivateallkeys': + await deactivateAllKeys(); + break; + case 'getkeyinfo': + await getKeyInfo(argument); + break; + case 'listinactivekeys': + await listInactiveKeys(); + break; + case 'listactivekeys': + await listActiveKeys(); + break; + case 'help': + console.log(` +Available commands: + generatekey Generate a single API key + generatekeys Generate multiple API keys + listkey List all API keys + listactivekeys List active API keys + listinactivekeys List inactive API keys + removekey Remove an API key + addkey Add your own API key + activatekey Activate an API key + deactivatekey Deactivate an API key + activateallkeys Activate all API keys + deactivateallkeys Deactivate all API keys + addkeydescription Add description to a key + listkeydescription Show description for a key + regeneratekey Regenerate (replace) an API key + getkeyinfo Show full info for a key + ratelimit Set rate limit for a key + changeport Change the server port + changeollamaurl Change the Ollama URL + addwebhook Add a webhook + deletewebhook Delete a webhook by ID + listwebhooks List all webhooks + help Show this help message + exit Shut down the server +`); + break; + case 'exit': + console.log('Shutting down...'); + await db.close(); + rl.close(); + process.exit(0); + break; + default: + console.log('Unknown command. Type "help" to see available commands.'); + } + } catch (err) { + console.error('Command error:', err.message); } }); } -function generateKey(db) { +async function generateKey() { const apiKey = crypto.randomBytes(20).toString('hex'); - db.run('INSERT INTO apiKeys(key, rate_limit) VALUES(?, 10)', [apiKey], (err) => { - if (err) { - console.error('Error generating API key:', err.message); - } else { - console.log(`API key generated: ${apiKey}`); - } - }); + await db.run('INSERT INTO apiKeys(key, rate_limit) VALUES(?, 10)', [apiKey]); + console.log(`API key generated: ${apiKey}`); } -function generateKeys(db, count) { +async function generateKeys(count) { if (!count || isNaN(count)) { console.log('Invalid number of keys'); return; } const numberOfKeys = parseInt(count); for (let i = 0; i < numberOfKeys; i++) { - generateKey(db); + await generateKey(); } } -function listKeys(db) { - db.all('SELECT key, active, description FROM apiKeys', [], (err, rows) => { - if (err) { - console.error('Error listing API keys:', err.message); - } else { - console.log('API keys:', rows); - } - }); +async function listKeys() { + const rows = await db.all('SELECT key, active, description FROM apiKeys'); + console.log('API keys:', rows); } -function removeKey(db, key) { - db.run('DELETE FROM apiKeys WHERE key = ?', [key], (err) => { - if (err) { - console.error('Error removing API key:', err.message); - } else { - console.log('API key removed'); - } - }); +async function removeKey(key) { + if (!key) { + console.log('API key is required'); + return; + } + await db.run('DELETE FROM apiKeys WHERE key = ?', [key]); + console.log('API key removed'); } -function addKey(db, key) { +async function addKey(key) { + if (!key) { + console.log('API key is required'); + return; + } console.log('Warning: Adding your own keys may be unsafe. It is recommended to generate keys using the generatekey command.'); - db.run('INSERT INTO apiKeys(key, rate_limit) VALUES(?, 10)', [key], (err) => { - if (err) { - console.error('Error adding API key:', err.message); - } else { - console.log(`API key added: ${key}`); - } - }); + await db.run('INSERT INTO apiKeys(key, rate_limit) VALUES(?, 10)', [key]); + console.log(`API key added: ${key}`); } -function changePort(newPort) { +async function changePort(newPort) { if (!newPort || isNaN(newPort)) { console.log('Invalid port number'); return; } const port = parseInt(newPort); if (server) { - server.close((err) => { - if (err) { - console.error('Error closing the server:', err.message); - } else { - console.log(`Server closed on port ${currentPort}`); - updatePortAndRestart(port); - } + await new Promise((resolve, reject) => { + server.close((err) => { + if (err) { + console.error('Error closing the server:', err.message); + reject(err); + } else { + console.log(`Server closed on port ${currentPort}`); + resolve(); + } + }); }); + } + await fs.promises.writeFile('port.conf', port.toString(), 'utf8'); + console.log(`Port number saved to port.conf: ${port}`); + if (expressApp) { + startServer(port, expressApp); } else { - updatePortAndRestart(port); + console.error('Express app not available. Unable to restart server.'); } } -function updatePortAndRestart(port) { - fs.writeFile('port.conf', port.toString(), (err) => { - if (err) { - console.error('Error saving port number:', err.message); - } else { - console.log(`Port number saved to port.conf: ${port}`); - if (expressApp) { - startServer(port, expressApp); - } else { - console.error('Express app not available. Unable to restart server.'); - } - } - }); -} - -function changeOllamaURL(newURL) { - const urlPattern = /^(http|https):\/\/[^\s$.?#].[^\s]*$/gm; - if (!newURL || !urlPattern.test(newURL)) { - console.log('Invalid Ollama URL'); +async function changeOllamaURL(newURL) { + if (!newURL || !VALID_URL_PATTERN.test(newURL)) { + console.log('Invalid Ollama URL. Must be a valid http/https URL.'); return; } - const URL = newURL; - fs.writeFile('ollamaURL.conf', URL, (err) => { - if (err) { - console.error('Error saving Ollama URL:', err.message); - } else { - console.log(`Ollama URL saved to ollamaURL.conf: ${URL}`); - } - }); + await fs.promises.writeFile('ollamaURL.conf', newURL, 'utf8'); + console.log(`Ollama URL saved to ollamaURL.conf: ${newURL}`); } -function setRateLimit(db, key, limit) { +async function setRateLimit(key, limit) { if (!key || !limit || isNaN(limit)) { console.log('Invalid API key or rate limit number'); return; } const rateLimit = parseInt(limit); - db.run('UPDATE apiKeys SET rate_limit = ? WHERE key = ?', [rateLimit, key], (err) => { - if (err) { - console.error('Error setting rate limit:', err.message); - } else { - console.log(`Rate limit set to ${rateLimit} requests per minute for API key: ${key}`); - } - }); + await db.run('UPDATE apiKeys SET rate_limit = ? WHERE key = ?', [rateLimit, key]); + console.log(`Rate limit set to ${rateLimit} requests per minute for API key: ${key}`); } -function addWebhook(db, url) { +async function addWebhook(url) { if (!url) { console.log('Webhook URL is required'); return; } - db.run('INSERT INTO webhooks (url) VALUES (?)', [url], (err) => { - if (err) { - console.error('Error adding webhook:', err.message); - } else { - console.log(`Webhook added: ${url}`); - } - }); + if (!VALID_URL_PATTERN.test(url)) { + console.log('Invalid webhook URL. Must be a valid http/https URL.'); + return; + } + await db.run('INSERT INTO webhooks (url) VALUES (?)', [url]); + console.log(`Webhook added: ${url}`); } -function deleteWebhook(db, id) { +async function deleteWebhook(id) { if (!id) { console.log('Webhook ID is required'); return; } - db.run('DELETE FROM webhooks WHERE id = ?', [id], (err) => { - if (err) { - console.error('Error deleting webhook:', err.message); - } else { - console.log('Webhook deleted'); - } - }); + await db.run('DELETE FROM webhooks WHERE id = ?', [id]); + console.log('Webhook deleted'); } -function listWebhooks(db) { - db.all('SELECT id, url FROM webhooks', [], (err, rows) => { - if (err) { - console.error('Error listing webhooks:', err.message); - } else { - console.log('Webhooks:', rows); - } - }); +async function listWebhooks() { + const rows = await db.all('SELECT id, url FROM webhooks'); + console.log('Webhooks:', rows); } -function activateKey(db, key) { - db.run('UPDATE apiKeys SET active = 1 WHERE key = ?', [key], (err) => { - if (err) { - console.error('Error activating API key:', err.message); - } else { - console.log(`API key ${key} activated`); - } - }); +async function activateKey(key) { + if (!key) { + console.log('API key is required'); + return; + } + await db.run('UPDATE apiKeys SET active = 1 WHERE key = ?', [key]); + console.log(`API key ${key} activated`); } -function deactivateKey(db, key) { - db.run('UPDATE apiKeys SET active = 0 WHERE key = ?', [key], (err) => { - if (err) { - console.error('Error deactivating API key:', err.message); - } else { - console.log(`API key ${key} deactivated`); - } - }); +async function deactivateKey(key) { + if (!key) { + console.log('API key is required'); + return; + } + await db.run('UPDATE apiKeys SET active = 0 WHERE key = ?', [key]); + console.log(`API key ${key} deactivated`); } -function addKeyDescription(db, key, description) { +async function addKeyDescription(key, description) { if (!key || !description) { console.log('Invalid API key or description'); return; } - db.run('UPDATE apiKeys SET description = ? WHERE key = ?', [description, key], (err) => { - if (err) { - console.error('Error adding description:', err.message); - } else { - console.log(`Description added to API key ${key}`); - } - }); + await db.run('UPDATE apiKeys SET description = ? WHERE key = ?', [description, key]); + console.log(`Description added to API key ${key}`); } -function listKeyDescription(db, key) { +async function listKeyDescription(key) { if (!key) { console.log('Invalid API key'); return; } - db.get('SELECT description FROM apiKeys WHERE key = ?', [key], (err, row) => { - if (err) { - console.error('Error retrieving description:', err.message); - } else { - if (row) { - console.log(`Description for API key ${key}: ${row.description}`); - } else { - console.log(`No description found for API key ${key}`); - } - } - }); + const row = await db.get('SELECT description FROM apiKeys WHERE key = ?', [key]); + if (row) { + console.log(`Description for API key ${key}: ${row.description}`); + } else { + console.log(`No description found for API key ${key}`); + } } -function regenerateKey(db, oldKey) { +async function regenerateKey(oldKey) { if (!oldKey) { console.log('Invalid API key'); return; } const newApiKey = crypto.randomBytes(20).toString('hex'); - db.run('UPDATE apiKeys SET key = ? WHERE key = ?', [newApiKey, oldKey], (err) => { - if (err) { - console.error('Error regenerating API key:', err.message); - } else { - console.log(`API key regenerated. New API key: ${newApiKey}`); - } - }); + await db.run('UPDATE apiKeys SET key = ? WHERE key = ?', [newApiKey, oldKey]); + console.log(`API key regenerated. New API key: ${newApiKey}`); } -function activateAllKeys(db) { - db.run('UPDATE apiKeys SET active = 1', (err) => { - if (err) { - console.error('Error activating all API keys:', err.message); - } else { - console.log('All API keys activated'); - } - }); +async function activateAllKeys() { + await db.run('UPDATE apiKeys SET active = 1'); + console.log('All API keys activated'); } -function deactivateAllKeys(db) { - db.run('UPDATE apiKeys SET active = 0', (err) => { - if (err) { - console.error('Error deactivating all API keys:', err.message); - } else { - console.log('All API keys deactivated'); - } - }); +async function deactivateAllKeys() { + await db.run('UPDATE apiKeys SET active = 0'); + console.log('All API keys deactivated'); } -function getKeyInfo(db, key) { - db.get('SELECT * FROM apiKeys WHERE key = ?', [key], (err, row) => { - if (err) { - console.error('Error retrieving API key info:', err.message); - } else if (row) { - console.log('API key info:', row); - } else { - console.log('No API key found with the given key.'); - } - }); +async function getKeyInfo(key) { + if (!key) { + console.log('API key is required'); + return; + } + const row = await db.get('SELECT * FROM apiKeys WHERE key = ?', [key]); + if (row) { + console.log('API key info:', row); + } else { + console.log('No API key found with the given key.'); + } } -function listInactiveKeys(db) { - db.all('SELECT key FROM apiKeys WHERE active = 0', [], (err, rows) => { - if (err) { - console.error('Error listing inactive API keys:', err.message); - } else { - console.log('Inactive API keys:', rows); - } - }); +async function listInactiveKeys() { + const rows = await db.all('SELECT key FROM apiKeys WHERE active = 0'); + console.log('Inactive API keys:', rows); } -function listActiveKeys(db) { - db.all('SELECT key FROM apiKeys WHERE active = 1', [], (err, rows) => { - if (err) { - console.error('Error listing active API keys:', err.message); - } else { - console.log('Active API keys:', rows); - } - }); +async function listActiveKeys() { + const rows = await db.all('SELECT key FROM apiKeys WHERE active = 1'); + console.log('Active API keys:', rows); } function getOllamaURL() { @@ -419,50 +365,37 @@ function getOllamaURL() { if (fs.existsSync('ollamaURL.conf')) { fs.readFile('ollamaURL.conf', 'utf8', (err, data) => { if (err) { - reject('Error reading Ollama url from file:', err.message); + reject(new Error('Error reading Ollama url from file: ' + err.message)); } else { const ollamaURL = data.trim(); if (typeof ollamaURL !== 'string' || ollamaURL === '') { - reject('Invalid Ollama url in ollamaURL.conf'); + reject(new Error('Invalid Ollama url in ollamaURL.conf')); } else { resolve(ollamaURL); } } }); } else { - reject('Ollama url configuration file not found'); + reject(new Error('Ollama url configuration file not found')); } }); } -function sendWebhookNotification(db, payload) { - db.all('SELECT url FROM webhooks', [], (err, rows) => { - if (err) { - console.error('Error retrieving webhooks:', err.message); - } else { - rows.forEach(row => { - const webhookPayload = { - content: JSON.stringify(payload, null, 2) - }; - - axios.post(row.url, webhookPayload) - .then(response => { - console.log('Webhook notification sent successfully:', response.data); - }) - .catch(error => { - console.error('Error sending webhook notification:', error.message); - }); - }); +function sendWebhookNotification(payload) { + db.all('SELECT url FROM webhooks').then(rows => { + for (const row of rows) { + axios.post(row.url, { content: JSON.stringify(payload, null, 2) }) + .catch(err => console.error('Error sending webhook notification:', err.message)); } + }).catch(err => { + console.error('Error retrieving webhooks:', err.message); }); } module.exports = { startServer, - askForPort, - askForOllamaURL, + resolveConfig, startCLI, getOllamaURL, - sendWebhookNotification, - updatePortAndRestart + sendWebhookNotification };