files cant't be seen by LLM, but everything else works.
This commit is contained in:
emil
+165
-98
@@ -7,6 +7,9 @@ const path = require('path');
|
||||
|
||||
const app = express();
|
||||
const PORT = process.env.PORT || 3002;
|
||||
const CHROMA_HOST = process.env.CHROMA_HOST || 'chromadb';
|
||||
const CHROMA_PORT = process.env.CHROMA_PORT || '8000';
|
||||
const CHROMA_URL = `http://${CHROMA_HOST}:${CHROMA_PORT}`;
|
||||
|
||||
// Middleware
|
||||
app.use(cors());
|
||||
@@ -26,21 +29,22 @@ const storage = multer.diskStorage({
|
||||
|
||||
const upload = multer({ storage });
|
||||
|
||||
// In-memory storage for knowledge bases metadata
|
||||
// In-memory storage
|
||||
const knowledgeBases = new Map();
|
||||
const knowledgeFiles = new Map();
|
||||
const uploadedFiles = new Map();
|
||||
|
||||
// Token usage tracking
|
||||
const tokenUsageStats = new Map(); // conversationId -> aggregated token stats
|
||||
|
||||
// In-memory storage for token usage events
|
||||
const tokenUsageLog = [];
|
||||
|
||||
// Extract token usage from SSE response
|
||||
// ===== CHROMA CLIENT =====
|
||||
|
||||
function getChromaClient() {
|
||||
return new ChromaClient({ path: CHROMA_URL });
|
||||
}
|
||||
|
||||
// ===== TOKEN USAGE =====
|
||||
|
||||
function extractTokenUsage(fullResponse) {
|
||||
try {
|
||||
// Parse the last SSE data block which contains usage info
|
||||
const lines = fullResponse.split('\n');
|
||||
for (let i = lines.length - 1; i >= 0; i--) {
|
||||
const line = lines[i];
|
||||
@@ -57,7 +61,7 @@ function extractTokenUsage(fullResponse) {
|
||||
};
|
||||
}
|
||||
} catch {
|
||||
// Not valid JSON, skip
|
||||
// skip
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -67,13 +71,6 @@ function extractTokenUsage(fullResponse) {
|
||||
}
|
||||
}
|
||||
|
||||
// Global ChromaDB client (recreated per request when path changes)
|
||||
function getChromaClient(persistDir = './data/chromadb') {
|
||||
return new ChromaClient({
|
||||
path: `file://${path.resolve(persistDir)}`,
|
||||
});
|
||||
}
|
||||
|
||||
// ===== OPENAI PROXY =====
|
||||
|
||||
app.post('/api/v1/chat/completions', async (req, res) => {
|
||||
@@ -107,7 +104,6 @@ app.post('/api/v1/chat/completions', async (req, res) => {
|
||||
});
|
||||
}
|
||||
|
||||
// Stream the response
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-cache');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
@@ -125,7 +121,6 @@ app.post('/api/v1/chat/completions', async (req, res) => {
|
||||
}
|
||||
res.end();
|
||||
|
||||
// Parse token usage from the final SSE message
|
||||
const tokenUsage = extractTokenUsage(fullResponse);
|
||||
if (tokenUsage) {
|
||||
const usageEntry = {
|
||||
@@ -137,13 +132,11 @@ app.post('/api/v1/chat/completions', async (req, res) => {
|
||||
totalTokens: tokenUsage.totalTokens,
|
||||
};
|
||||
tokenUsageLog.push(usageEntry);
|
||||
console.log(`Token usage: ${usageEntry.totalTokens} tokens (${usageEntry.promptTokens} prompt, ${usageEntry.completionTokens} completion) for model ${model}`);
|
||||
console.log(`Token usage: ${usageEntry.totalTokens} tokens for model ${model}`);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Proxy error:', error);
|
||||
res.status(500).json({
|
||||
error: { message: error.message || 'Internal server error' },
|
||||
});
|
||||
res.status(500).json({ error: { message: error.message || 'Internal server error' } });
|
||||
}
|
||||
});
|
||||
|
||||
@@ -156,9 +149,7 @@ app.post('/api/v1/models', async (req, res) => {
|
||||
|
||||
const url = `${apiUrl.replace(/\/$/, '')}/models`;
|
||||
const response = await fetch(url, {
|
||||
headers: {
|
||||
'Authorization': `Bearer ${apiKey}`,
|
||||
},
|
||||
headers: { 'Authorization': `Bearer ${apiKey}` },
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
@@ -224,7 +215,6 @@ function chunkText(text, chunkSize = 500, overlap = 50) {
|
||||
chunks.push(text.slice(start, end));
|
||||
start += chunkSize - overlap;
|
||||
if (start >= text.length) break;
|
||||
// Don't create tiny chunks at the end
|
||||
if (text.length - start < chunkSize * 0.3) {
|
||||
chunks[chunks.length - 1] = text.slice(chunks.length > 0 ? start - (chunkSize - overlap) : 0);
|
||||
break;
|
||||
@@ -253,11 +243,58 @@ async function getEmbeddings(texts, apiUrl, apiKey, model) {
|
||||
return data.data.map((d) => d.embedding);
|
||||
}
|
||||
|
||||
app.get('/api/knowledge', async (req, res) => {
|
||||
const { chromaDir } = req.query;
|
||||
async function processFileToKnowledge(filePath, fileName, baseId, apiUrl, apiKey, embeddingModel) {
|
||||
let text = '';
|
||||
try {
|
||||
const client = getChromaClient(chromaDir || './data/chromadb');
|
||||
// Get all collections to list knowledge bases
|
||||
text = fs.readFileSync(filePath, 'utf-8');
|
||||
} catch (e) {
|
||||
console.warn(`Cannot read ${fileName}:`, e.message);
|
||||
return { success: false, error: e.message };
|
||||
}
|
||||
|
||||
const chunks = chunkText(text, 500, 50);
|
||||
if (chunks.length === 0) return { success: false, error: 'Empty file' };
|
||||
|
||||
const embeddings = await getEmbeddings(chunks, apiUrl, apiKey, embeddingModel || 'text-embedding-3-small');
|
||||
|
||||
const client = getChromaClient();
|
||||
const collection = await client.getCollection({ name: baseId });
|
||||
|
||||
const ids = chunks.map((_, i) => `${baseId}_chunk_${Date.now()}_${i}_${Math.random().toString(36).substring(2, 6)}`);
|
||||
const metadatas = chunks.map((chunk, i) => ({
|
||||
source: fileName,
|
||||
chunkIndex: i,
|
||||
fileId: fileName,
|
||||
}));
|
||||
|
||||
await collection.add({ ids, embeddings, documents: chunks, metadatas });
|
||||
|
||||
const fileInfo = {
|
||||
id: fileName,
|
||||
originalName: fileName,
|
||||
size: fs.statSync(filePath).size,
|
||||
mimeType: 'text/plain',
|
||||
chunkCount: chunks.length,
|
||||
uploadedAt: Date.now(),
|
||||
};
|
||||
|
||||
const files = knowledgeFiles.get(baseId) || [];
|
||||
files.push(fileInfo);
|
||||
knowledgeFiles.set(baseId, files);
|
||||
|
||||
const kb = knowledgeBases.get(baseId);
|
||||
if (kb) {
|
||||
kb.documentCount = (kb.documentCount || 0) + chunks.length;
|
||||
kb.files = files;
|
||||
knowledgeBases.set(baseId, kb);
|
||||
}
|
||||
|
||||
return { success: true, fileInfo, chunks: chunks.length };
|
||||
}
|
||||
|
||||
app.get('/api/knowledge', async (req, res) => {
|
||||
try {
|
||||
const client = getChromaClient();
|
||||
const collections = await client.listCollections();
|
||||
const bases = [];
|
||||
|
||||
@@ -290,16 +327,22 @@ app.post('/api/knowledge', async (req, res) => {
|
||||
const id = name.trim().toLowerCase().replace(/[^a-z0-9]/g, '_') + '_' + Date.now();
|
||||
|
||||
try {
|
||||
const client = getChromaClient(chromaDir || './data/chromadb');
|
||||
const client = getChromaClient();
|
||||
await client.getOrCreateCollection({
|
||||
name: id,
|
||||
metadata: { name: name.trim(), description: description || '', createdAt: Date.now() },
|
||||
metadata: {
|
||||
name: name.trim(),
|
||||
description: description || '',
|
||||
chromaDir: chromaDir || '/data/chromadb',
|
||||
createdAt: Date.now(),
|
||||
},
|
||||
});
|
||||
|
||||
const kb = {
|
||||
id,
|
||||
name: name.trim(),
|
||||
description: description || '',
|
||||
chromaDir: chromaDir || '/data/chromadb',
|
||||
documentCount: 0,
|
||||
files: [],
|
||||
};
|
||||
@@ -315,17 +358,15 @@ app.post('/api/knowledge', async (req, res) => {
|
||||
|
||||
app.delete('/api/knowledge/:id', async (req, res) => {
|
||||
const { id } = req.params;
|
||||
const { chromaDir } = req.query;
|
||||
|
||||
try {
|
||||
const client = getChromaClient(chromaDir || './data/chromadb');
|
||||
const client = getChromaClient();
|
||||
await client.deleteCollection({ name: id });
|
||||
knowledgeBases.delete(id);
|
||||
knowledgeFiles.delete(id);
|
||||
res.json({ success: true });
|
||||
} catch (error) {
|
||||
console.error('Delete knowledge base error:', error);
|
||||
// Even if Chroma throws, clean up our state
|
||||
knowledgeBases.delete(id);
|
||||
knowledgeFiles.delete(id);
|
||||
res.json({ success: true });
|
||||
@@ -343,67 +384,97 @@ app.post('/api/knowledge/:id/files', upload.single('file'), async (req, res) =>
|
||||
}
|
||||
|
||||
try {
|
||||
// Read file content
|
||||
let text = '';
|
||||
if (file.mimetype === 'application/pdf') {
|
||||
text = `[PDF file: ${file.originalname}]`; // Simplified; in production use pdf-parse
|
||||
} else {
|
||||
text = fs.readFileSync(file.path, 'utf-8');
|
||||
}
|
||||
|
||||
// Chunk the text
|
||||
const chunks = chunkText(text, 500, 50);
|
||||
if (chunks.length === 0) {
|
||||
return res.status(400).json({ error: 'File is empty or could not be parsed' });
|
||||
}
|
||||
|
||||
// Get embeddings
|
||||
const embeddings = await getEmbeddings(chunks, apiUrl, apiKey, embeddingModel || 'text-embedding-3-small');
|
||||
|
||||
// Store in ChromaDB
|
||||
const client = getChromaClient();
|
||||
const collection = await client.getCollection({ name: id });
|
||||
|
||||
const ids = chunks.map((_, i) => `${id}_chunk_${Date.now()}_${i}`);
|
||||
const metadatas = chunks.map((chunk, i) => ({
|
||||
source: file.originalname,
|
||||
chunkIndex: i,
|
||||
fileId: file.filename,
|
||||
}));
|
||||
|
||||
await collection.add({
|
||||
ids,
|
||||
embeddings,
|
||||
documents: chunks,
|
||||
metadatas,
|
||||
});
|
||||
|
||||
// Update metadata
|
||||
const fileInfo = {
|
||||
id: file.filename,
|
||||
originalName: file.originalname,
|
||||
size: file.size,
|
||||
chunkCount: chunks.length,
|
||||
};
|
||||
|
||||
const files = knowledgeFiles.get(id) || [];
|
||||
files.push(fileInfo);
|
||||
knowledgeFiles.set(id, files);
|
||||
|
||||
const kb = knowledgeBases.get(id);
|
||||
if (kb) {
|
||||
kb.documentCount = (kb.documentCount || 0) + chunks.length;
|
||||
kb.files = files;
|
||||
knowledgeBases.set(id, kb);
|
||||
}
|
||||
|
||||
res.json({ file: fileInfo, chunks: chunks.length });
|
||||
const result = await processFileToKnowledge(
|
||||
file.path,
|
||||
file.originalname,
|
||||
id,
|
||||
apiUrl,
|
||||
apiKey,
|
||||
embeddingModel
|
||||
);
|
||||
if (!result.success) return res.status(400).json({ error: result.error });
|
||||
res.json({ file: result.fileInfo, chunks: result.chunks });
|
||||
} catch (error) {
|
||||
console.error('Upload to knowledge base error:', error);
|
||||
res.status(500).json({ error: error.message });
|
||||
}
|
||||
});
|
||||
|
||||
// ===== SCAN LOCAL FOLDER =====
|
||||
|
||||
const SKIP_DIRS_SCAN = /node_modules|\.git|\.next|\.vscode|\.idea|dist|build|__pycache__|\.cache|\.venv|venv|env|target|\.turbo/i;
|
||||
|
||||
function shouldSkipPath(p) {
|
||||
const parts = p.split('/');
|
||||
for (const part of parts) {
|
||||
if (part.startsWith('.') || SKIP_DIRS_SCAN.test(part)) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function getFilesRecursively(dirPath, files = []) {
|
||||
if (!fs.existsSync(dirPath) || !fs.statSync(dirPath).isDirectory()) return files;
|
||||
const entries = fs.readdirSync(dirPath);
|
||||
for (const entry of entries) {
|
||||
const fullPath = path.join(dirPath, entry);
|
||||
const relPath = fullPath.replace(dirPath, '');
|
||||
if (shouldSkipPath(relPath)) continue;
|
||||
const stat = fs.statSync(fullPath);
|
||||
if (stat.isDirectory()) {
|
||||
getFilesRecursively(fullPath, files);
|
||||
} else {
|
||||
files.push(fullPath);
|
||||
}
|
||||
}
|
||||
return files;
|
||||
}
|
||||
|
||||
app.post('/api/knowledge/:id/scan', async (req, res) => {
|
||||
const { id } = req.params;
|
||||
const { folderPath, apiUrl, apiKey, embeddingModel } = req.body;
|
||||
|
||||
if (!folderPath || !fs.existsSync(folderPath)) {
|
||||
return res.status(400).json({ error: 'Invalid or non-existent folder path' });
|
||||
}
|
||||
if (!apiUrl || !apiKey) {
|
||||
return res.status(400).json({ error: 'API URL and Key are required for embeddings' });
|
||||
}
|
||||
|
||||
try {
|
||||
const allPaths = getFilesRecursively(folderPath);
|
||||
const results = [];
|
||||
let processed = 0;
|
||||
const maxFiles = 50;
|
||||
|
||||
for (const filePath of allPaths.slice(0, maxFiles)) {
|
||||
try {
|
||||
const fileName = path.relative(folderPath, filePath);
|
||||
const result = await processFileToKnowledge(
|
||||
filePath,
|
||||
fileName,
|
||||
id,
|
||||
apiUrl,
|
||||
apiKey,
|
||||
embeddingModel
|
||||
);
|
||||
if (result.success) {
|
||||
results.push({ file: fileName, chunks: result.chunks });
|
||||
} else {
|
||||
results.push({ file: fileName, error: result.error });
|
||||
}
|
||||
} catch (e) {
|
||||
results.push({ file: path.relative(folderPath, filePath), error: e.message });
|
||||
}
|
||||
processed++;
|
||||
}
|
||||
|
||||
res.json({ processed, total: allPaths.length, results });
|
||||
} catch (error) {
|
||||
console.error('Scan folder error:', error);
|
||||
res.status(500).json({ error: error.message });
|
||||
}
|
||||
});
|
||||
|
||||
app.delete('/api/knowledge/:baseId/files/:fileId', async (req, res) => {
|
||||
const { baseId, fileId } = req.params;
|
||||
|
||||
@@ -411,7 +482,6 @@ app.delete('/api/knowledge/:baseId/files/:fileId', async (req, res) => {
|
||||
const client = getChromaClient();
|
||||
const collection = await client.getCollection({ name: baseId });
|
||||
|
||||
// Find all chunks with this fileId
|
||||
const results = await collection.get({
|
||||
where: { fileId: { $eq: fileId } },
|
||||
});
|
||||
@@ -420,7 +490,6 @@ app.delete('/api/knowledge/:baseId/files/:fileId', async (req, res) => {
|
||||
await collection.delete({ ids: results.ids });
|
||||
}
|
||||
|
||||
// Update metadata
|
||||
const files = (knowledgeFiles.get(baseId) || []).filter((f) => f.id !== fileId);
|
||||
knowledgeFiles.set(baseId, files);
|
||||
|
||||
@@ -478,14 +547,13 @@ app.post('/api/knowledge/:id/query', async (req, res) => {
|
||||
|
||||
// Health check
|
||||
app.get('/api/health', (req, res) => {
|
||||
res.json({ status: 'ok', timestamp: new Date().toISOString() });
|
||||
res.json({ status: 'ok', timestamp: new Date().toISOString(), chromaUrl: CHROMA_URL });
|
||||
});
|
||||
|
||||
// Token usage tracking endpoints
|
||||
// Token usage tracking
|
||||
app.get('/api/token-usage', (req, res) => {
|
||||
try {
|
||||
const { period = 'all', conversationId } = req.query;
|
||||
|
||||
let filtered = [...tokenUsageLog];
|
||||
|
||||
if (conversationId) {
|
||||
@@ -506,7 +574,6 @@ app.get('/api/token-usage', (req, res) => {
|
||||
const completionTokens = filtered.reduce((sum, e) => sum + e.completionTokens, 0);
|
||||
const totalTokens = filtered.reduce((sum, e) => sum + e.totalTokens, 0);
|
||||
|
||||
// Per-model breakdown
|
||||
const modelBreakdown = {};
|
||||
filtered.forEach((entry) => {
|
||||
if (!modelBreakdown[entry.model]) {
|
||||
@@ -521,7 +588,7 @@ app.get('/api/token-usage', (req, res) => {
|
||||
res.json({
|
||||
summary: { promptTokens, completionTokens, totalTokens, requestCount: filtered.length },
|
||||
modelBreakdown,
|
||||
history: filtered.slice(-100), // last 100 entries
|
||||
history: filtered.slice(-100),
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Token usage error:', error);
|
||||
@@ -531,11 +598,11 @@ app.get('/api/token-usage', (req, res) => {
|
||||
|
||||
app.delete('/api/token-usage', (req, res) => {
|
||||
tokenUsageLog.length = 0;
|
||||
tokenUsageStats.clear();
|
||||
res.json({ success: true });
|
||||
});
|
||||
|
||||
app.listen(PORT, () => {
|
||||
console.log(`AIUI Backend running on port ${PORT}`);
|
||||
console.log(`ChromaDB URL: ${CHROMA_URL}`);
|
||||
console.log(`Upload directory: ${uploadDir}`);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user