voyageai-cli
Version:
CLI for Voyage AI embeddings, reranking, and MongoDB Atlas Vector Search
399 lines (354 loc) • 10.7 kB
JavaScript
/**
* Knowledge Base Manager for RAG Chat
* Handles KB selection, file upload, ingestion, and persistence
*/
class KBManager {
constructor() {
this.currentKB = null;
this.kbs = [];
this.isIngesting = false;
this.ingestionProgress = { current: 0, total: 0, currentFile: '' };
// Local state
this.loadLastKB();
}
/**
* Get the currently selected embedding model from the chat config UI.
* Falls back to 'voyage-4-large' if the element is missing.
*/
getEmbeddingModel() {
const sel = document.getElementById('chatEmbeddingModel');
return sel ? sel.value : 'voyage-4-large';
}
/**
* Load last used KB from localStorage
*/
loadLastKB() {
const lastKBName = localStorage.getItem('__vai_last_kb');
if (lastKBName) {
this.currentKB = lastKBName;
}
}
/**
* Save last used KB to localStorage
*/
saveLastKB() {
if (this.currentKB) {
localStorage.setItem('__vai_last_kb', this.currentKB);
}
}
/**
* Fetch list of all KBs from server
*/
async listKBs() {
try {
const res = await fetch('/api/rag/kbs');
if (!res.ok) throw new Error('Failed to fetch KBs');
const data = await res.json();
this.kbs = data.kbs || [];
return this.kbs;
} catch (err) {
console.error('Error listing KBs:', err);
return [];
}
}
/**
* Select a KB (or null to create new)
*/
async selectKB(kbName) {
try {
const res = await fetch('/api/rag/kb-select', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ kbName })
});
if (!res.ok) throw new Error('Failed to select KB');
const data = await res.json();
this.currentKB = data.selected || null;
this.saveLastKB();
return data;
} catch (err) {
console.error('Error selecting KB:', err);
throw err;
}
}
/**
* Upload files and start ingestion
* Yields progress events
*/
async *ingestFiles(files, kbName = null) {
if (!files || files.length === 0) {
throw new Error('No files selected');
}
// Validate file types and sizes
for (const file of files) {
if (!['text/plain', 'text/markdown', 'application/x-markdown'].includes(file.type)) {
throw new Error(`Invalid file type: ${file.name}. Only .txt and .md supported.`);
}
if (file.size > 10 * 1024 * 1024) { // 10MB limit
throw new Error(`File too large: ${file.name}. Max 10MB.`);
}
}
// Create FormData
const formData = new FormData();
for (const file of files) {
formData.append('files', file);
}
if (kbName) {
formData.append('kbName', kbName);
}
formData.append('embeddingModel', this.getEmbeddingModel());
try {
this.isIngesting = true;
const res = await fetch('/api/rag/ingest', {
method: 'POST',
body: formData
});
if (!res.ok) {
const error = await res.json();
throw new Error(error.error || 'Ingestion failed');
}
// Parse streaming response
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buffer = '';
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split('\n');
buffer = lines.pop(); // Keep incomplete line in buffer
for (const line of lines) {
if (!line.trim()) continue;
try {
const event = JSON.parse(line);
yield event;
// Update progress
if (event.type === 'progress') {
this.ingestionProgress = {
current: event.current || 0,
total: event.total || 0,
currentFile: event.file || ''
};
}
} catch (e) {
console.warn('Failed to parse event:', line, e);
}
}
}
// Final flush of buffer
if (buffer.trim()) {
try {
const event = JSON.parse(buffer);
yield event;
} catch (e) {
console.warn('Failed to parse final event:', buffer, e);
}
}
this.isIngesting = false;
} catch (err) {
this.isIngesting = false;
console.error('Error ingesting files:', err);
throw err;
}
}
/**
* Ingest pasted text into a KB
* Yields progress events via NDJSON stream
*/
async *ingestText(text, kbName, title) {
if (!text || !text.trim()) {
throw new Error('Text content is required');
}
try {
this.isIngesting = true;
const res = await fetch('/api/rag/ingest-text', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ text, kbName, title, embeddingModel: this.getEmbeddingModel() })
});
if (!res.ok) {
const error = await res.json();
throw new Error(error.error || 'Text ingestion failed');
}
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buffer = '';
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split('\n');
buffer = lines.pop();
for (const line of lines) {
if (!line.trim()) continue;
try {
const event = JSON.parse(line);
yield event;
if (event.type === 'progress') {
this.ingestionProgress = {
current: event.current || 0,
total: event.total || 0,
currentFile: event.stage || ''
};
}
} catch (e) {
console.warn('Failed to parse event:', line, e);
}
}
}
if (buffer.trim()) {
try { yield JSON.parse(buffer); } catch (e) { /* ignore */ }
}
this.isIngesting = false;
} catch (err) {
this.isIngesting = false;
console.error('Error ingesting text:', err);
throw err;
}
}
/**
* Ingest content from a URL into a KB
* Yields progress events via NDJSON stream
*/
async *ingestURL(url, kbName) {
if (!url || !/^https?:\/\//i.test(url)) {
throw new Error('A valid URL starting with http:// or https:// is required');
}
try {
this.isIngesting = true;
const res = await fetch('/api/rag/ingest-url', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ url, kbName, embeddingModel: this.getEmbeddingModel() })
});
if (!res.ok) {
const error = await res.json();
throw new Error(error.error || 'URL ingestion failed');
}
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buffer = '';
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split('\n');
buffer = lines.pop();
for (const line of lines) {
if (!line.trim()) continue;
try {
const event = JSON.parse(line);
yield event;
if (event.type === 'progress') {
this.ingestionProgress = {
current: event.current || 0,
total: event.total || 0,
currentFile: event.stage || ''
};
}
} catch (e) {
console.warn('Failed to parse event:', line, e);
}
}
}
if (buffer.trim()) {
try { yield JSON.parse(buffer); } catch (e) { /* ignore */ }
}
this.isIngesting = false;
} catch (err) {
this.isIngesting = false;
console.error('Error ingesting URL:', err);
throw err;
}
}
/**
* Remove a document from KB
*/
async removeDoc(kbName, docId) {
try {
const res = await fetch(`/api/rag/docs/${encodeURIComponent(kbName)}/${encodeURIComponent(docId)}`, {
method: 'DELETE'
});
if (!res.ok) throw new Error('Failed to remove document');
return await res.json();
} catch (err) {
console.error('Error removing doc:', err);
throw err;
}
}
/**
* Get KB details and doc list
*/
async getKBDetails(kbName) {
try {
const res = await fetch(`/api/rag/kb/${encodeURIComponent(kbName)}`);
if (!res.ok) throw new Error('Failed to fetch KB details');
return await res.json();
} catch (err) {
console.error('Error fetching KB details:', err);
throw err;
}
}
/**
* Rename a KB (updates displayName; internal name/collection unchanged)
*/
async renameKB(kbName, displayName) {
const res = await fetch(`/api/rag/kb/${encodeURIComponent(kbName)}`, {
method: 'PATCH',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ newName: displayName })
});
if (!res.ok) {
const err = await res.json();
throw new Error(err.error || 'Failed to rename KB');
}
return await res.json();
}
/**
* List documents in a KB (grouped by fileName)
*/
async listDocs(kbName) {
try {
const res = await fetch(`/api/rag/kb/${encodeURIComponent(kbName)}/docs`);
if (!res.ok) throw new Error('Failed to list documents');
const data = await res.json();
return data.docs || [];
} catch (err) {
console.error('Error listing docs:', err);
return [];
}
}
/**
* Remove all chunks for a file by fileName
*/
async removeDocByName(kbName, fileName) {
try {
const res = await fetch(
`/api/rag/docs/${encodeURIComponent(kbName)}/by-name/${encodeURIComponent(fileName)}`,
{ method: 'DELETE' }
);
if (!res.ok) throw new Error('Failed to remove document');
return await res.json();
} catch (err) {
console.error('Error removing doc:', err);
throw err;
}
}
/**
* Clear a KB (delete all docs)
*/
async clearKB(kbName) {
try {
const res = await fetch(`/api/rag/kb/${encodeURIComponent(kbName)}`, {
method: 'DELETE'
});
if (!res.ok) throw new Error('Failed to clear KB');
return await res.json();
} catch (err) {
console.error('Error clearing KB:', err);
throw err;
}
}
}
// Global instance
window.kbManager = new KBManager();