78 lines
2.9 KiB
JavaScript
78 lines
2.9 KiB
JavaScript
import { ChromaClient } from 'chromadb';
|
|
import { MODEL_ID, embed, embedMany } from './embed.js';
|
|
|
|
// ==============================|| Chroma collections ||============================== //
|
|
//
|
|
// Two collections, deliberately separate:
|
|
//
|
|
// intent_examples — ~20 phrasings per intent. Answers "which question is
|
|
// this?", never "what is the number?".
|
|
// console_docs — chunked markdown. Genuine document Q&A.
|
|
//
|
|
// NO OPERATIONAL DATA IS EMBEDDED. No bookings, riders or customers ever enter
|
|
// the vector store. A vector store is a snapshot and this data changes by the
|
|
// minute; every figure still comes from a live API call through the existing
|
|
// typed functions. That is the whole reason this design is safe to build.
|
|
|
|
export const INTENTS_COLLECTION = 'intent_examples';
|
|
export const DOCS_COLLECTION = 'console_docs';
|
|
|
|
const client = new ChromaClient({ path: process.env.CHROMA_URL || 'http://localhost:8000' });
|
|
|
|
// Chroma would otherwise call its own default embedding function (which
|
|
// downloads a different model). We embed ourselves so queries and documents
|
|
// are guaranteed to come from the same model.
|
|
const noopEmbeddingFunction = { generate: async (texts) => embedMany(texts) };
|
|
|
|
export const getCollection = async (name) =>
|
|
client.getOrCreateCollection({
|
|
name,
|
|
embeddingFunction: noopEmbeddingFunction,
|
|
metadata: {
|
|
// Stamped so a model swap is detectable rather than silently degrading
|
|
// every score. Re-seed after changing it.
|
|
'hnsw:space': 'cosine',
|
|
embeddingModel: MODEL_ID
|
|
}
|
|
});
|
|
|
|
export const resetCollection = async (name) => {
|
|
try {
|
|
await client.deleteCollection({ name });
|
|
} catch {
|
|
// Not present yet — nothing to delete.
|
|
}
|
|
return getCollection(name);
|
|
};
|
|
|
|
// Chroma returns cosine DISTANCE (0 = identical). Similarity is the useful
|
|
// direction for a human-facing confidence, so convert once, here, rather than
|
|
// leaving every call site to remember which way round it is.
|
|
const toSimilarity = (distance) => 1 - Number(distance);
|
|
|
|
export const query = async (name, text, topK = 5) => {
|
|
const collection = await getCollection(name);
|
|
const vector = await embed(text);
|
|
const res = await collection.query({ queryEmbeddings: [vector], nResults: topK });
|
|
|
|
const ids = res.ids?.[0] || [];
|
|
return ids.map((id, i) => ({
|
|
id,
|
|
document: res.documents?.[0]?.[i] || '',
|
|
metadata: res.metadatas?.[0]?.[i] || {},
|
|
score: toSimilarity(res.distances?.[0]?.[i] ?? 1)
|
|
}));
|
|
};
|
|
|
|
export const health = async () => {
|
|
await client.heartbeat();
|
|
const names = (await client.listCollections()).map((c) => c.name ?? c);
|
|
const counts = {};
|
|
for (const name of [INTENTS_COLLECTION, DOCS_COLLECTION]) {
|
|
if (!names.includes(name)) continue;
|
|
// eslint-disable-next-line no-await-in-loop
|
|
counts[name] = await (await getCollection(name)).count();
|
|
}
|
|
return { chroma: 'up', collections: counts };
|
|
};
|