40 lines
1.1 KiB
JavaScript
40 lines
1.1 KiB
JavaScript
import { ChromaClient } from "chromadb";
|
|
import { OpenAIEmbeddings } from "langchain/embeddings/openai";
|
|
import { OpenAI } from "langchain/llms/openai";
|
|
|
|
const client = new ChromaClient({
|
|
path: process.env.CHROMA_DB_PATH || "./chromadb",
|
|
});
|
|
|
|
const embeddings = new OpenAIEmbeddings({
|
|
openAIApiKey: process.env.OPENAI_API_KEY,
|
|
});
|
|
|
|
export async function addDocument(collectionName, text, metadata = {}) {
|
|
const collection = await client.getOrCreateCollection({
|
|
name: collectionName,
|
|
});
|
|
const embedding = await embeddings.embedQuery(text);
|
|
await collection.add({
|
|
documents: [text],
|
|
embeddings: [embedding],
|
|
metadatas: [metadata],
|
|
});
|
|
}
|
|
|
|
export async function getSimilarDocuments(collectionName, query, k = 5) {
|
|
const collection = await client.getOrCreateCollection({
|
|
name: collectionName,
|
|
});
|
|
const embedding = await embeddings.embedQuery(query);
|
|
const results = await collection.query({
|
|
queryEmbeddings: [embedding],
|
|
nResults: k,
|
|
});
|
|
return results.ids[0].map((id, idx) => ({
|
|
id,
|
|
score: results.scores[0][idx],
|
|
document: results.documents[0][idx],
|
|
metadata: results.metadatas[0][idx],
|
|
}));
|
|
} |