58 lines
1.6 KiB
JavaScript
58 lines
1.6 KiB
JavaScript
import { ChromaClient } from 'chromadb';
|
|
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
|
|
|
|
const chroma = new ChromaClient({
|
|
host: process.env.CHROMA_HOST || 'localhost',
|
|
port: parseInt(process.env.CHROMA_PORT, 10) || 8000,
|
|
});
|
|
|
|
const embeddings = new OpenAIEmbeddings({
|
|
openAIApiKey: process.env.OPENAI_API_KEY,
|
|
});
|
|
|
|
const COLLECTION_NAME = 'rag_collection';
|
|
|
|
/**
|
|
* Ensure the collection exists in ChromaDB.
|
|
*/
|
|
async function ensureCollection() {
|
|
const collections = await chroma.listCollections();
|
|
if (!collections.includes(COLLECTION_NAME)) {
|
|
await chroma.createCollection({ name: COLLECTION_NAME });
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Add an array of documents to the vector store.
|
|
* @param {string[]} docs
|
|
*/
|
|
export async function addDocuments(docs) {
|
|
await ensureCollection();
|
|
const ids = docs.map((_, idx) => `doc-${Date.now()}-${idx}`);
|
|
const embeddingsResult = await embeddings.embedDocuments(docs);
|
|
await chroma.add({
|
|
collection_name: COLLECTION_NAME,
|
|
ids,
|
|
documents: docs,
|
|
embeddings: embeddingsResult,
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Query the vector store for the most relevant documents.
|
|
* @param {string} queryText
|
|
* @param {number} nResults
|
|
* @returns {Promise<{documents: string[]}>}
|
|
*/
|
|
export async function query(queryText, nResults = 3) {
|
|
await ensureCollection();
|
|
const embedding = await embeddings.embedQuery(queryText);
|
|
const results = await chroma.query({
|
|
collection_name: COLLECTION_NAME,
|
|
query_embeddings: [embedding],
|
|
n_results: nResults,
|
|
});
|
|
// ChromaDB returns an array of objects; extract documents
|
|
const docs = results.documents || [];
|
|
return { documents: docs };
|
|
} |