47 lines
1.4 KiB
JavaScript
47 lines
1.4 KiB
JavaScript
import { ChromaClient } from 'chromadb';
|
|
import { OllamaEmbeddings } from 'langchain/embeddings/ollama';
|
|
import { Document } from 'langchain/document';
|
|
|
|
export async function createVectorStore() {
|
|
const chroma = new ChromaClient({ path: process.env.CHROMA_URL });
|
|
const collection = await chroma.getOrCreateCollection({
|
|
name: process.env.CHROMA_COLLECTION,
|
|
});
|
|
const embeddings = new OllamaEmbeddings({
|
|
model: process.env.OLLAMA_EMBEDDING_MODEL || 'nomic-embed-text',
|
|
});
|
|
return new VectorStore(collection, embeddings);
|
|
}
|
|
|
|
class VectorStore {
|
|
constructor(collection, embeddings) {
|
|
this.collection = collection;
|
|
this.embeddings = embeddings;
|
|
}
|
|
|
|
async addDocuments(docs) {
|
|
const texts = docs.map((d) => d.pageContent);
|
|
const embeddings = await this.embeddings.embedDocuments(texts);
|
|
await this.collection.addDocuments({
|
|
documents: docs,
|
|
embeddings,
|
|
});
|
|
}
|
|
|
|
async similaritySearch(query, k = 4) {
|
|
const embedding = await this.embeddings.embedQuery(query);
|
|
const results = await this.collection.getNearestNeighbors({
|
|
queryEmbeddings: [embedding],
|
|
n: k,
|
|
});
|
|
const ids = results.ids[0];
|
|
const docs = await this.collection.getDocuments({ ids });
|
|
return docs.map(
|
|
(doc) =>
|
|
new Document({
|
|
pageContent: doc.document,
|
|
metadata: doc.metadata,
|
|
})
|
|
);
|
|
}
|
|
} |