|
|
import express from 'express'; |
|
|
import { SystemMessage, HumanMessage, Runnable, LlamaCppLLM } from './src/index.js'; |
|
|
import bodyParser from 'body-parser'; |
|
|
import path from 'path'; |
|
|
import fs from 'fs'; |
|
|
|
|
|
|
|
|
class EmailClassifierRunnable extends Runnable { |
|
|
constructor(llm) { |
|
|
super(); |
|
|
this.llm = llm; |
|
|
} |
|
|
async _call(input, config) { |
|
|
|
|
|
if (!this.llm) return { category: "Error", confidence: 0, reason: "Model not initialized" }; |
|
|
|
|
|
const messages = this._buildPrompt(input); |
|
|
const response = await this.llm.invoke(messages, config); |
|
|
return this._parseClassification(response.content); |
|
|
} |
|
|
_buildPrompt(email) { |
|
|
return [ |
|
|
new SystemMessage(`You are an email classification assistant. Classify into: Spam, Invoice, Meeting Request, Urgent, Personal, Other. Respond in JSON like {"category": "X", "confidence": 0.9, "reason": "Y"}.`), |
|
|
new HumanMessage(`Classify:\nSubject: ${email.subject}\nBody: ${email.body}`) |
|
|
]; |
|
|
} |
|
|
_parseClassification(response) { |
|
|
try { |
|
|
const jsonMatch = response.match(/\{[\s\S]*\}/); |
|
|
if (!jsonMatch) throw new Error('No JSON found'); |
|
|
return JSON.parse(jsonMatch[0]); |
|
|
} catch (e) { return { category: 'Other', confidence: 0, reason: 'Failed to parse JSON', raw: response }; } |
|
|
} |
|
|
} |
|
|
|
|
|
const app = express(); |
|
|
const PORT = 7860; |
|
|
|
|
|
app.use(bodyParser.json()); |
|
|
|
|
|
let classifier = null; |
|
|
|
|
|
async function initModel() { |
|
|
try { |
|
|
console.log("Loading model..."); |
|
|
|
|
|
if (!fs.existsSync('./models/Qwen3-1.7B-Q8_0.gguf')) { |
|
|
console.error("Model file missing!"); |
|
|
return; |
|
|
} |
|
|
|
|
|
const llm = new LlamaCppLLM({ |
|
|
modelPath: './models/Qwen3-1.7B-Q8_0.gguf', |
|
|
temperature: 0.1, |
|
|
maxTokens: 200 |
|
|
}); |
|
|
|
|
|
|
|
|
await llm.invoke("Hi"); |
|
|
|
|
|
classifier = new EmailClassifierRunnable(llm); |
|
|
console.log("Model loaded successfully!"); |
|
|
} catch (err) { |
|
|
console.error("Failed to load model:", err); |
|
|
} |
|
|
} |
|
|
|
|
|
app.post('/classify', async (req, res) => { |
|
|
if (!classifier) return res.status(503).json({ error: "Model loading or failed" }); |
|
|
try { |
|
|
const { subject, body } = req.body; |
|
|
const result = await classifier.invoke({ subject, body, from: 'api' }); |
|
|
res.json(result); |
|
|
} catch (error) { |
|
|
res.status(500).json({ error: error.message }); |
|
|
} |
|
|
}); |
|
|
|
|
|
app.get('/', (req, res) => res.send('AI Email Classifier Running. POST /classify to use.')); |
|
|
|
|
|
app.listen(PORT, '0.0.0.0', () => { |
|
|
console.log(`Server listening on ${PORT}`); |
|
|
initModel(); |
|
|
}); |
|
|
|