138 lines
3.5 KiB
JavaScript
138 lines
3.5 KiB
JavaScript
const http = require('http');
|
|
const fs = require('fs');
|
|
const WebSocket = require('ws');
|
|
require('dotenv').config();
|
|
|
|
|
|
// Configuration
|
|
const PORT = 9898;
|
|
const API_KEY = 'your_litellm_secret_key'; // consider process.env.API_KEY; // Replace with your actual API key
|
|
const LLM_MODEL = 'llama3-70b';
|
|
const HTML_FILE = 'index.html';
|
|
|
|
/**
|
|
* Create and configure HTTP server
|
|
* @returns {http.Server} Configured HTTP server
|
|
*/
|
|
function createHttpServer() {
|
|
return http.createServer(handleHttpRequest);
|
|
}
|
|
|
|
/**
|
|
* Handle HTTP requests by serving the HTML file
|
|
* @param {http.IncomingMessage} req - HTTP request object
|
|
* @param {http.ServerResponse} res - HTTP response object
|
|
*/
|
|
function handleHttpRequest(req, res) {
|
|
res.writeHead(200, {'Content-Type': 'text/html'});
|
|
|
|
fs.readFile(HTML_FILE, (err, data) => {
|
|
if (err) {
|
|
console.error('Error reading HTML file:', err);
|
|
res.statusCode = 500;
|
|
res.end('Error getting the file');
|
|
return;
|
|
}
|
|
res.end(data);
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Set up WebSocket server
|
|
* @param {http.Server} server - HTTP server to attach WebSocket to
|
|
* @returns {WebSocket.Server} Configured WebSocket server
|
|
*/
|
|
function setupWebSocketServer(server) {
|
|
const wss = new WebSocket.Server({ server });
|
|
|
|
wss.on('connection', handleWebSocketConnection);
|
|
|
|
return wss;
|
|
}
|
|
|
|
/**
|
|
* Handle new WebSocket connections
|
|
* @param {WebSocket} ws - WebSocket connection
|
|
*/
|
|
function handleWebSocketConnection(ws) {
|
|
console.log('Client connected');
|
|
|
|
ws.on('message', (message) => handleIncomingMessage(ws, message));
|
|
|
|
ws.on('close', () => {
|
|
console.log('Client disconnected');
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Process incoming WebSocket messages
|
|
* @param {WebSocket} ws - WebSocket connection
|
|
* @param {string} message - Received message
|
|
*/
|
|
function handleIncomingMessage(ws, message) {
|
|
console.log(`Received message => ${message}`);
|
|
queryLLM(ws, message.toString());
|
|
}
|
|
|
|
/**
|
|
* Query the language model API
|
|
* @param {WebSocket} ws - WebSocket connection to send the response
|
|
* @param {string} message - User message to process
|
|
* @returns {Promise<string|null>} The LLM response or null on error
|
|
*/
|
|
async function queryLLM(ws, message) {
|
|
try {
|
|
console.log(`\n\n\nWhat is the key: ${API_KEY}`);
|
|
const got = (await import('got')).default;
|
|
const response = await got(process.env.LLM_API_ENDPOINT, {
|
|
method: 'POST',
|
|
headers: {
|
|
'content-type': 'application/json',
|
|
'Authorization': `Bearer ${API_KEY}`,
|
|
},
|
|
json: {
|
|
model: LLM_MODEL,
|
|
messages: [
|
|
{
|
|
"role": "system",
|
|
"content": "Keep response to 4 lines of text."
|
|
},
|
|
{ role: 'user', content: message }
|
|
],
|
|
max_tokens: 1000
|
|
}
|
|
});
|
|
|
|
const data = JSON.parse(response.body);
|
|
|
|
if (data?.choices?.[0]?.message?.content) {
|
|
const content = data.choices[0].message.content;
|
|
ws.send(content);
|
|
return content;
|
|
} else {
|
|
console.error('Invalid or empty response from LLM API');
|
|
ws.send('Sorry, I received an invalid response from the language model.');
|
|
return null;
|
|
}
|
|
} catch (error) {
|
|
console.error('Error querying LLM:', error);
|
|
ws.send('Sorry, there was an error processing your request.');
|
|
return null;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Initialize and start the server
|
|
*/
|
|
function startServer() {
|
|
const server = createHttpServer();
|
|
setupWebSocketServer(server);
|
|
|
|
server.listen(PORT, () => {
|
|
console.log(`Server listening on port ${PORT}`);
|
|
});
|
|
}
|
|
|
|
// Start the server when this file is run
|
|
startServer();
|