107 lines
3.1 KiB
JavaScript
107 lines
3.1 KiB
JavaScript
const {QdrantClient} = require('@qdrant/js-client-rest');
|
|
const {QDRANT, COLLECTIONS, getEnv} = require('@nexusai/shared');
|
|
|
|
let client;
|
|
|
|
// ********** QDrant Client **********
|
|
function getClient(){
|
|
if(!client){
|
|
client = new QdrantClient({url: getEnv('QDRANT_URL', QDRANT.DEFAULT_URL)});
|
|
}
|
|
return client;
|
|
}
|
|
|
|
// ********** Collections **********
|
|
|
|
async function initCollections() {
|
|
const client = getClient();
|
|
for (const name of Object.values(COLLECTIONS)) {
|
|
const exists = await collectionExists(name);
|
|
if (!exists) {
|
|
await client.createCollection(name, {
|
|
vectors: {
|
|
size: QDRANT.VECTOR_SIZE,
|
|
distance: QDRANT.DISTANCE_METRIC
|
|
}
|
|
});
|
|
console.log(`Created Qdrant collection: ${name}`);
|
|
} else {
|
|
console.log(`Qdrant collection already exists: ${name}`);
|
|
}
|
|
}
|
|
}
|
|
|
|
async function collectionExists(name) {
|
|
try {
|
|
const client = getClient();
|
|
await client.getCollection(name);
|
|
return true;
|
|
} catch (err) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// ********* Upsert Vectors *********
|
|
|
|
//payload is metadata stored alongside the vector in QDrant.
|
|
//SQLite ID stored here to be able to link back to original data
|
|
async function upsertVector(collection, id, vector, payload = {}) {
|
|
const client = getClient();
|
|
await client.upsertVector(collection, {
|
|
wait: true, // Wait for the operation to complete before returning
|
|
points: [{id, vector, payload}]
|
|
});
|
|
}
|
|
|
|
// Upsert an episode vector into the EPISODES collection, with the episode ID as the point ID
|
|
async function upsertEpisode(id, vector, payload={}) {
|
|
return upsertVector(COLLECTIONS.EPISODES, id, vector, payload);
|
|
}
|
|
|
|
async function upsertEntity(id, vector, payload={}) {
|
|
return upsertVector(COLLECTIONS.ENTITIES, id, vector, payload);
|
|
}
|
|
|
|
async function upsertSummary(id, vector, payload={}) {
|
|
return upsertVector(COLLECTIONS.SUMMARIES, id, vector, payload);
|
|
}
|
|
|
|
//********** Search Vectors ********** */
|
|
async function searchCollection(collection, vector, limit = QDRANT.DEFAULT_LIMIT, filter = null){
|
|
const client = getClient();
|
|
const params = {vector,limit, with_payload: true};
|
|
|
|
if (filter) params.filter = filter;
|
|
|
|
return client.search(collection, params);
|
|
}
|
|
async function searchEpisodes(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
|
|
return searchCollection(COLLECTIONS.EPISODES, vector, limit, filter);
|
|
}
|
|
async function searchEntities(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
|
|
return searchCollection(COLLECTIONS.ENTITIES, vector, limit, filter);
|
|
}
|
|
async function searchSummaries(vector, limit = QDRANT.DEFAULT_LIMIT, filter = null) {
|
|
return searchCollection(COLLECTIONS.SUMMARIES, vector, limit, filter);
|
|
}
|
|
|
|
//********** Delete Vectors ********** */
|
|
|
|
async function deleteVector(collection, id) {
|
|
const client = getClient();
|
|
await client.delete(collection, {
|
|
wait: true,
|
|
points: [id]
|
|
});
|
|
}
|
|
|
|
module.exports = {
|
|
initCollections,
|
|
upsertEpisode,
|
|
upsertEntity,
|
|
upsertSummary,
|
|
searchEpisodes,
|
|
searchEntities,
|
|
searchSummaries,
|
|
deleteVector
|
|
}; |