genkitx-qdrant
Version:
Genkit AI framework plugin for the Qdrant vector database.
244 lines • 8.27 kB
JavaScript
var __defProp = Object.defineProperty;
var __defProps = Object.defineProperties;
var __getOwnPropDescs = Object.getOwnPropertyDescriptors;
var __getOwnPropSymbols = Object.getOwnPropertySymbols;
var __hasOwnProp = Object.prototype.hasOwnProperty;
var __propIsEnum = Object.prototype.propertyIsEnumerable;
var __defNormalProp = (obj, key, value) => key in obj ? __defProp(obj, key, { enumerable: true, configurable: true, writable: true, value }) : obj[key] = value;
var __spreadValues = (a, b) => {
for (var prop in b || (b = {}))
if (__hasOwnProp.call(b, prop))
__defNormalProp(a, prop, b[prop]);
if (__getOwnPropSymbols)
for (var prop of __getOwnPropSymbols(b)) {
if (__propIsEnum.call(b, prop))
__defNormalProp(a, prop, b[prop]);
}
return a;
};
var __spreadProps = (a, b) => __defProps(a, __getOwnPropDescs(b));
var __async = (__this, __arguments, generator) => {
return new Promise((resolve, reject) => {
var fulfilled = (value) => {
try {
step(generator.next(value));
} catch (e) {
reject(e);
}
};
var rejected = (value) => {
try {
step(generator.throw(value));
} catch (e) {
reject(e);
}
};
var step = (x) => x.done ? resolve(x.value) : Promise.resolve(x.value).then(fulfilled, rejected);
step((generator = generator.apply(__this, __arguments)).next());
});
};
import {
CommonRetrieverOptionsSchema,
Document,
indexerRef,
retrieverRef
} from "@genkit-ai/ai/retriever";
import { QdrantClient } from "@qdrant/js-client-rest";
import { z } from "genkit";
import { genkitPlugin } from "genkit/plugin";
import { v5 as uuidv5 } from "uuid";
const FilterType = z.any();
const QdrantRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({
k: z.number().default(10),
filter: FilterType.optional(),
scoreThreshold: z.number().optional()
});
const QdrantIndexerOptionsSchema = z.null().optional();
const CONTENT_PAYLOAD_KEY = "content";
const METADATA_PAYLOAD_KEY = "metadata";
const CONTENT_TYPE_KEY = "_content_type";
const qdrantRetrieverRef = (collectionName, displayName = null) => {
return retrieverRef({
name: `qdrant/${collectionName}`,
info: {
label: displayName != null ? displayName : `Qdrant - ${collectionName}`
},
configSchema: QdrantRetrieverOptionsSchema
});
};
const qdrantIndexerRef = (collectionName, displayName = null) => {
return indexerRef({
name: `qdrant/${collectionName}`,
info: {
label: displayName != null ? displayName : `Qdrant - ${collectionName}`
},
configSchema: QdrantIndexerOptionsSchema
});
};
function qdrant(params) {
return genkitPlugin("qdrant", (ai) => __async(null, null, function* () {
params.forEach((p) => configureQdrantRetriever(ai, p));
params.forEach((p) => configureQdrantIndexer(ai, p));
}));
}
var index_default = qdrant;
function configureQdrantRetriever(ai, params) {
var _a;
const {
embedder,
collectionName,
embedderOptions,
clientParams,
contentPayloadKey,
metadataPayloadKey
} = params;
const client = new QdrantClient(clientParams);
const contentKey = contentPayloadKey != null ? contentPayloadKey : CONTENT_PAYLOAD_KEY;
const metadataKey = metadataPayloadKey != null ? metadataPayloadKey : METADATA_PAYLOAD_KEY;
const dataTypeKey = (_a = params.dataTypePayloadKey) != null ? _a : CONTENT_TYPE_KEY;
return ai.defineRetriever(
{
name: `qdrant/${collectionName}`,
configSchema: QdrantRetrieverOptionsSchema
},
(content, options) => __async(null, null, function* () {
yield ensureCollection(params, false, ai);
const queryEmbeddings = yield ai.embed({
embedder,
content,
options: embedderOptions
});
const results = (yield client.query(collectionName, {
query: queryEmbeddings[0].embedding,
limit: options.k,
filter: options.filter,
score_threshold: options.scoreThreshold,
with_payload: [contentKey, metadataKey, dataTypeKey],
with_vector: false
})).points;
const documents = results.map((result) => {
var _a2, _b, _c, _d, _e, _f;
const content2 = (_b = (_a2 = result.payload) == null ? void 0 : _a2[contentKey]) != null ? _b : "";
const metadata = __spreadProps(__spreadValues({}, (_d = (_c = result.payload) == null ? void 0 : _c[metadataKey]) != null ? _d : {}), {
_similarityScore: result.score
});
const dataType = (_f = (_e = result.payload) == null ? void 0 : _e[dataTypeKey]) != null ? _f : "text";
return Document.fromData(
content2,
dataType,
metadata
).toJSON();
});
return {
documents
};
})
);
}
function configureQdrantIndexer(ai, params) {
var _a;
const {
embedder,
collectionName,
embedderOptions,
clientParams,
contentPayloadKey,
metadataPayloadKey
} = params;
const client = new QdrantClient(clientParams);
const contentKey = contentPayloadKey != null ? contentPayloadKey : CONTENT_PAYLOAD_KEY;
const metadataKey = metadataPayloadKey != null ? metadataPayloadKey : METADATA_PAYLOAD_KEY;
const dataTypeKey = (_a = params.dataTypePayloadKey) != null ? _a : CONTENT_TYPE_KEY;
return ai.defineIndexer(
{
name: `qdrant/${collectionName}`,
configSchema: QdrantIndexerOptionsSchema
},
(docs, options) => __async(null, null, function* () {
yield ensureCollection(params, true, ai);
const embeddings = yield Promise.all(
docs.map(
(doc) => ai.embed({
embedder,
content: doc,
options: embedderOptions
})
)
);
const points = embeddings.map((embeddingArr, i) => {
const doc = docs[i];
const embeddingDocs = doc.getEmbeddingDocuments(embeddingArr);
return embeddingArr.map((docEmbedding, j) => {
const embeddingDoc = embeddingDocs[j] || {};
const id = uuidv5(JSON.stringify(embeddingDoc), uuidv5.URL);
return {
id,
vector: docEmbedding.embedding,
payload: {
[contentKey]: embeddingDoc.data,
[metadataKey]: embeddingDoc.metadata,
[dataTypeKey]: embeddingDoc.dataType
}
};
});
}).reduce((acc, val) => acc.concat(val), []);
yield client.upsert(collectionName, { points });
})
);
}
function createQdrantCollection(params, ai) {
return __async(this, null, function* () {
const { embedder, embedderOptions, clientParams, collectionName } = params;
const client = new QdrantClient(clientParams);
let collectionCreateOptions = params.collectionCreateOptions;
if (!collectionCreateOptions) {
const embeddings = yield ai.embed({
embedder,
content: "SOME_TEXT",
options: embedderOptions
});
const vector = Array.isArray(embeddings) ? embeddings[0].embedding : embeddings.embedding;
collectionCreateOptions = {
vectors: {
size: vector.length,
distance: "Cosine"
}
};
}
return yield client.createCollection(collectionName, collectionCreateOptions);
});
}
function deleteQdrantCollection(params) {
return __async(this, null, function* () {
const client = new QdrantClient(params.clientParams);
return yield client.deleteCollection(params.collectionName);
});
}
function ensureCollection(params, createCollection = true, ai) {
return __async(this, null, function* () {
const { clientParams, collectionName } = params;
const client = new QdrantClient(clientParams);
if ((yield client.collectionExists(collectionName)).exists) {
return;
}
if (createCollection) {
yield createQdrantCollection(params, ai);
} else {
throw new Error(
`Collection ${collectionName} does not exist. Index some documents first.`
);
}
});
}
export {
QdrantIndexerOptionsSchema,
configureQdrantIndexer,
configureQdrantRetriever,
createQdrantCollection,
index_default as default,
deleteQdrantCollection,
qdrant,
qdrantIndexerRef,
qdrantRetrieverRef
};
//# sourceMappingURL=index.mjs.map