242 lines
8.2 KiB
JavaScript
242 lines
8.2 KiB
JavaScript
|
import 'dotenv/config';
|
||
|
import express from 'express';
|
||
|
import bodyParser from 'body-parser';
|
||
|
import cmd from 'cmd-promise';
|
||
|
import cors from 'cors';
|
||
|
import cheerio from 'cheerio';
|
||
|
import llamaTokenizer from 'llama-tokenizer-js';
|
||
|
import googleIt from 'google-it';
|
||
|
import Groq from 'groq-sdk';
|
||
|
|
||
|
// Constants and initialization
|
||
|
const app = express();
|
||
|
const port = 3000;
|
||
|
const prompt = process.env.PROMPT;
|
||
|
|
||
|
const groq = new Groq({ apiKey: process.env.GROQ });
|
||
|
|
||
|
app.use(cors({
|
||
|
origin: '*',
|
||
|
allowedHeaders: ['Content-Type', 'x-forwarded-for-id', 'x-forwarded-for-name']
|
||
|
}));
|
||
|
app.use(bodyParser.json());
|
||
|
|
||
|
let isProcessing = false;
|
||
|
let conversationHistory = {};
|
||
|
|
||
|
// Helper function to get current timestamp
|
||
|
const getTimestamp = () => {
|
||
|
const now = new Date();
|
||
|
const date = now.toLocaleDateString('en-US');
|
||
|
const time = now.toLocaleTimeString('en-US');
|
||
|
return `${date} [${time}]`;
|
||
|
};
|
||
|
|
||
|
// Middleware to track conversation history by CF-Connecting-IP
|
||
|
app.use((req, res, next) => {
|
||
|
const ip = req.headers['x-forwarded-for-id'] || req.headers['cf-connecting-ip'] || req.headers['x-forwarded-for'] || req.headers['x-real-ip'] || req.ip;
|
||
|
const name = req.headers['x-forwarded-for-name'];
|
||
|
const guild = req.headers['x-forwarded-for-guild'];
|
||
|
|
||
|
req.clientIp = ip; // Store the IP in a request property
|
||
|
|
||
|
if (!conversationHistory[req.clientIp]) {
|
||
|
console.log(`${getTimestamp()} [INFO] Initializing conversation history for: ${req.clientIp}`);
|
||
|
conversationHistory[req.clientIp] = [{
|
||
|
role: 'system',
|
||
|
content: `My name is: ${name || 'Unknown'}, my Discord ID is: ${req.clientIp}.` + (guild ? ` We are chatting inside ${guild} a Discord Server.` : '') + prompt
|
||
|
}];
|
||
|
}
|
||
|
|
||
|
next();
|
||
|
});
|
||
|
|
||
|
function countLlamaTokens(messages) {
|
||
|
let totalTokens = 0;
|
||
|
for (const message of messages) {
|
||
|
if (message.role === 'user' || message.role === 'assistant') {
|
||
|
const encodedTokens = llamaTokenizer.encode(message.content);
|
||
|
totalTokens += encodedTokens.length;
|
||
|
}
|
||
|
}
|
||
|
return totalTokens;
|
||
|
}
|
||
|
|
||
|
function trimConversationHistory(ip, maxLength = 14000, tolerance = 25) {
|
||
|
const messages = conversationHistory[ip];
|
||
|
let totalTokens = countLlamaTokens(messages);
|
||
|
|
||
|
while (totalTokens > maxLength + tolerance && messages.length > 1) {
|
||
|
messages.shift(); // Remove the oldest messages first
|
||
|
totalTokens = countLlamaTokens(messages);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Function to scrape web page
|
||
|
async function scrapeWebPage(url, length = 2000) {
|
||
|
try {
|
||
|
const res = await fetch(url);
|
||
|
const html = await res.text();
|
||
|
const $ = cheerio.load(html);
|
||
|
|
||
|
const pageTitle = $('head title').text().trim();
|
||
|
const pageDescription = $('head meta[name="description"]').attr('content');
|
||
|
let plainTextContent = $('body').text().trim().replace(/[\r\n\t]+/g, ' ');
|
||
|
|
||
|
if (plainTextContent.length > length) {
|
||
|
plainTextContent = plainTextContent.substring(0, length) + '...';
|
||
|
}
|
||
|
|
||
|
return `Title: ${pageTitle}\nDescription: ${pageDescription || 'N/A'}\nContent: ${plainTextContent}\nURL: ${url}`;
|
||
|
} catch (err) {
|
||
|
console.error(`${getTimestamp()} [ERROR] Failed to scrape URL: ${url}`, err);
|
||
|
return null;
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Function to handle IP plugin
|
||
|
async function handleIPPlugin(ipAddr, ip, conversationHistory) {
|
||
|
try {
|
||
|
const url = new URL('https://api.abuseipdb.com/api/v2/check');
|
||
|
url.searchParams.append('ipAddress', ipAddr);
|
||
|
url.searchParams.append('maxAgeInDays', '90');
|
||
|
url.searchParams.append('verbose', '');
|
||
|
|
||
|
const options = {
|
||
|
method: 'GET',
|
||
|
headers: {
|
||
|
'Key': process.env.ABUSE_KEY,
|
||
|
'Accept': 'application/json'
|
||
|
}
|
||
|
};
|
||
|
|
||
|
const response = await fetch(url, options);
|
||
|
const data = await response.json();
|
||
|
|
||
|
let abuseResponse = `IP: ${ipAddr}\n`;
|
||
|
abuseResponse += `Abuse Score: ${data.data.abuseConfidenceScore}\n`;
|
||
|
abuseResponse += `Country: ${data.data.countryCode}\n`;
|
||
|
abuseResponse += `Usage Type: ${data.data.usageType}\n`;
|
||
|
abuseResponse += `ISP: ${data.data.isp}\n`;
|
||
|
abuseResponse += `Domain: ${data.data.domain}\n`;
|
||
|
|
||
|
if (data.data.totalReports) {
|
||
|
abuseResponse += `Total Reports: ${data.data.totalReports}\n`;
|
||
|
abuseResponse += `Last Reported: ${data.data.lastReportedAt}\n`;
|
||
|
}
|
||
|
|
||
|
const lastMessageIndex = conversationHistory[ip].length - 1;
|
||
|
if (lastMessageIndex >= 0) {
|
||
|
conversationHistory[ip][lastMessageIndex].content += "\n" + abuseResponse;
|
||
|
console.log(`${getTimestamp()} [INFO] Processed IP address: ${ipAddr}, response: ${abuseResponse}`);
|
||
|
} else {
|
||
|
console.error(`${getTimestamp()} [ERROR] Conversation history is unexpectedly empty for: ${ip}`);
|
||
|
}
|
||
|
} catch (err) {
|
||
|
console.error(`${getTimestamp()} [ERROR] Failed to process IP address: ${ipAddr}`, err);
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Main chat handler
|
||
|
app.post('/api/v1/chat', async (req, res) => {
|
||
|
const startTime = Date.now();
|
||
|
const ip = req.clientIp;
|
||
|
|
||
|
// if (isProcessing) {
|
||
|
// return res.status(429).json({ message: "System is busy processing another request, try again later" });
|
||
|
// }
|
||
|
|
||
|
isProcessing = true;
|
||
|
|
||
|
try {
|
||
|
const userMessage = req.body.message + `\nDate/Time: ${getTimestamp()}`;
|
||
|
conversationHistory[ip].push({ role: 'user', content: userMessage });
|
||
|
trimConversationHistory(ip);
|
||
|
|
||
|
const pluginTasks = [];
|
||
|
|
||
|
// Check for IPs in user message and process them
|
||
|
const ipRegex = /(\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b)/g;
|
||
|
const ipAddresses = userMessage.match(ipRegex);
|
||
|
if (ipAddresses) {
|
||
|
for (const ipAddr of ipAddresses) {
|
||
|
pluginTasks.push(handleIPPlugin(ipAddr, ip, conversationHistory));
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Check for URLs and scrape them
|
||
|
const urls = userMessage.match(/(https?:\/\/[^\s]+)/g);
|
||
|
if (urls) {
|
||
|
for (const url of urls) {
|
||
|
pluginTasks.push(scrapeWebPage(url).then(content => {
|
||
|
if (content) {
|
||
|
conversationHistory[ip].push({ role: 'assistant', content });
|
||
|
}
|
||
|
}));
|
||
|
}
|
||
|
}
|
||
|
|
||
|
await Promise.all(pluginTasks);
|
||
|
|
||
|
const completion = await groq.chat.completions.create({
|
||
|
messages: conversationHistory[ip],
|
||
|
//model: "llama3-8b-8192"
|
||
|
model: "llama-3.2-3b-preview"
|
||
|
});
|
||
|
|
||
|
const assistantMessage = completion.choices[0].message.content;
|
||
|
conversationHistory[ip].push({ role: 'assistant', content: assistantMessage });
|
||
|
|
||
|
res.json(assistantMessage);
|
||
|
} catch (error) {
|
||
|
console.error(`${getTimestamp()} [ERROR] An error occurred`, error);
|
||
|
res.status(500).json({ message: "An error occurred", error: error.message });
|
||
|
} finally {
|
||
|
isProcessing = false;
|
||
|
const endTime = Date.now();
|
||
|
const processingTime = ((endTime - startTime) / 1000).toFixed(2);
|
||
|
console.log(`${getTimestamp()} [STATS] Processing Time: ${processingTime} seconds`);
|
||
|
}
|
||
|
});
|
||
|
|
||
|
// Endpoint to restart core service
|
||
|
app.post('/api/v1/restart-core', (req, res) => {
|
||
|
console.log(`${getTimestamp()} [INFO] Restarting core service`);
|
||
|
cmd(`docker restart llama-gpu-server`).then(out => {
|
||
|
console.log(`${getTimestamp()} [INFO] Core service restarted`);
|
||
|
res.json(out.stdout);
|
||
|
}).catch(err => {
|
||
|
console.error(`${getTimestamp()} [ERROR] Failed to restart core service`, err);
|
||
|
res.status(500).json({
|
||
|
message: "An error occurred while restarting the core service",
|
||
|
error: err.message
|
||
|
});
|
||
|
});
|
||
|
});
|
||
|
|
||
|
// Endpoint to reset conversation history
|
||
|
app.post('/api/v1/reset-conversation', (req, res) => {
|
||
|
const ip = req.clientIp;
|
||
|
console.log(`${getTimestamp()} [INFO] Resetting conversation history for: ${ip}`);
|
||
|
|
||
|
conversationHistory[ip] = [{
|
||
|
role: 'system',
|
||
|
content: prompt
|
||
|
}];
|
||
|
|
||
|
console.log(`${getTimestamp()} [INFO] Conversation history reset for: ${ip}`);
|
||
|
res.json({ message: "Conversation history reset for: " + ip });
|
||
|
});
|
||
|
|
||
|
// Get conversation history for debugging purposes
|
||
|
app.get('/api/v1/conversation-history', (req, res) => {
|
||
|
const ip = req.clientIp;
|
||
|
console.log(`${getTimestamp()} [INFO] Fetching conversation history for: ${ip}`);
|
||
|
res.json(conversationHistory[ip]);
|
||
|
});
|
||
|
|
||
|
// Start server
|
||
|
app.listen(port, () => {
|
||
|
console.log(`${getTimestamp()} [INFO] Server running at http://localhost:${port}`);
|
||
|
});
|