genkitx-milvus
Version:
Genkit AI framework plugin for Milvus vector database.
323 lines • 10.7 kB
JavaScript
;
var __defProp = Object.defineProperty;
var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
var __getOwnPropNames = Object.getOwnPropertyNames;
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 __objRest = (source, exclude) => {
var target = {};
for (var prop in source)
if (__hasOwnProp.call(source, prop) && exclude.indexOf(prop) < 0)
target[prop] = source[prop];
if (source != null && __getOwnPropSymbols)
for (var prop of __getOwnPropSymbols(source)) {
if (exclude.indexOf(prop) < 0 && __propIsEnum.call(source, prop))
target[prop] = source[prop];
}
return target;
};
var __export = (target, all) => {
for (var name in all)
__defProp(target, name, { get: all[name], enumerable: true });
};
var __copyProps = (to, from, except, desc) => {
if (from && typeof from === "object" || typeof from === "function") {
for (let key of __getOwnPropNames(from))
if (!__hasOwnProp.call(to, key) && key !== except)
__defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
}
return to;
};
var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
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());
});
};
var src_exports = {};
__export(src_exports, {
DocumentSchema: () => DocumentSchema,
createMilvusCollection: () => createMilvusCollection,
default: () => src_default,
deleteMilvusCollection: () => deleteMilvusCollection,
describeMilvusCollection: () => describeMilvusCollection,
hasMilvusCollection: () => hasMilvusCollection,
insertMilvusData: () => insertMilvusData,
milvus: () => milvus,
milvusIndexer: () => milvusIndexer,
milvusIndexerRef: () => milvusIndexerRef,
milvusRetriever: () => milvusRetriever,
milvusRetrieverRef: () => milvusRetrieverRef,
searchMilvusData: () => searchMilvusData
});
module.exports = __toCommonJS(src_exports);
var import_retriever = require("@genkit-ai/ai/retriever");
var import_plugin = require("genkit/plugin");
var import_genkit = require("genkit");
var import_milvus2_sdk_node = require("@zilliz/milvus2-sdk-node");
const DocumentSchema = import_genkit.z.object({
id: import_genkit.z.number(),
vector: import_genkit.z.array(import_genkit.z.number()),
metadata: import_genkit.z.record(import_genkit.z.string(), import_genkit.z.any()).optional()
});
const MilvusRetrieverOptionsSchema = import_retriever.CommonRetrieverOptionsSchema.extend({
k: import_genkit.z.number().max(1e3).default(10),
filter: import_genkit.z.record(import_genkit.z.string(), import_genkit.z.any()).optional(),
limit: import_genkit.z.number().optional()
});
const MilvusIndexerOptionsSchema = import_genkit.z.null().optional();
function milvus(params) {
return (0, import_plugin.genkitPlugin)("milvus", (ai) => __async(this, null, function* () {
params.map((i) => milvusRetriever(ai, i));
params.map((i) => milvusIndexer(ai, i));
}));
}
var src_default = milvus;
const milvusRetrieverRef = (params) => {
var _a;
return (0, import_genkit.retrieverRef)({
name: `milvus/${params.collectionName}`,
info: {
label: (_a = params.displayName) != null ? _a : `Milvus - ${params.collectionName}`
},
configSchema: MilvusRetrieverOptionsSchema.optional()
});
};
const milvusIndexerRef = (params) => {
var _a;
return (0, import_genkit.indexerRef)({
name: `milvus/${params.collectionName}`,
info: {
label: (_a = params.displayName) != null ? _a : `Milvus - ${params.collectionName}`
},
configSchema: MilvusIndexerOptionsSchema.optional()
});
};
function milvusRetriever(ai, params) {
var _a;
const { collectionName, embedder, embedderOptions, dbName } = params;
const milvusConfig = (_a = params.clientParams) != null ? _a : getDefaultConfig();
return ai.defineRetriever(
{
name: `milvus/${collectionName}`,
configSchema: MilvusRetrieverOptionsSchema
},
(content, options) => __async(this, null, function* () {
const queryEmbeddings = yield ai.embed({
embedder,
content,
options: embedderOptions
});
const isCollectionExists = yield hasMilvusCollection({
collectionName,
dbName,
clientParams: milvusConfig
});
if (!isCollectionExists) {
yield createMilvusCollection({
dbName,
collectionName,
dimension: queryEmbeddings.length,
clientParams: milvusConfig
});
}
const response = yield searchMilvusData({
collectionName,
dbName,
query: queryEmbeddings.flatMap((e) => e.embedding),
limit: options.limit,
filter: JSON.stringify(options.filter),
topk: options.k,
clientParams: milvusConfig
});
return {
documents: response.results.map((result) => result["$meta"]).filter(
(m) => !!m
).map((m) => {
return import_genkit.Document.fromText(m.document, m.metadata).toJSON();
})
};
})
);
}
function milvusIndexer(ai, params) {
var _a;
const { collectionName, embedder, embedderOptions, dbName } = params;
const milvusConfig = (_a = params.clientParams) != null ? _a : getDefaultConfig();
return ai.defineIndexer(
{
name: `milvus/${collectionName}`,
configSchema: MilvusIndexerOptionsSchema.optional()
},
(docs) => __async(this, null, function* () {
const embeddings = yield Promise.all(
docs.map(
(doc) => ai.embed({
embedder,
content: doc,
options: embedderOptions
})
)
);
const entries = embeddings.map((value, i) => {
const metadata = __spreadValues({}, docs[i].metadata);
return {
vector: value,
document: docs[i].text,
metadata
};
});
const isCollectionExists = yield hasMilvusCollection({
collectionName,
dbName,
clientParams: milvusConfig
});
if (!isCollectionExists) {
yield createMilvusCollection({
dbName,
collectionName,
dimension: embeddings[0].length,
clientParams: milvusConfig
});
}
yield insertMilvusData({
collectionName,
dbName,
entries,
clientParams: milvusConfig
});
})
);
}
function createMilvusCollection(params) {
return __async(this, null, function* () {
const { clientParams, collectionName, dimension, dbName } = params;
const milvus2 = getMilvusInstance(clientParams);
return milvus2.createCollection({
collection_name: collectionName,
db_name: dbName,
dimension,
auto_id: true,
enable_dynamic_field: true
});
});
}
function describeMilvusCollection(params) {
return __async(this, null, function* () {
const { clientParams, collectionName, dbName } = params;
const milvus2 = getMilvusInstance(clientParams);
return yield milvus2.describeCollection({
collection_name: collectionName,
db_name: dbName
});
});
}
function insertMilvusData(params) {
return __async(this, null, function* () {
const { clientParams, collectionName, dbName, entries } = params;
const milvus2 = getMilvusInstance(clientParams);
return milvus2.insert({
collection_name: collectionName,
db_name: dbName,
fields_data: entries
});
});
}
function searchMilvusData(params) {
return __async(this, null, function* () {
const _a = params, { clientParams, collectionName, dbName, query } = _a, restParams = __objRest(_a, ["clientParams", "collectionName", "dbName", "query"]);
const milvus2 = getMilvusInstance(clientParams);
return milvus2.search(__spreadValues({
collection_name: collectionName,
db_name: dbName,
data: query
}, restParams));
});
}
function deleteMilvusCollection(params) {
return __async(this, null, function* () {
const { clientParams, collectionName, dbName } = params;
const milvus2 = getMilvusInstance(clientParams);
try {
return yield milvus2.dropCollection({
collection_name: collectionName,
db_name: dbName
});
} catch (error) {
console.error("Failed to delete Milvus collection:", error);
throw error;
}
});
}
function hasMilvusCollection(params) {
return __async(this, null, function* () {
const { clientParams, collectionName, dbName } = params;
const milvus2 = getMilvusInstance(clientParams);
const response = yield milvus2.hasCollection({
collection_name: collectionName,
db_name: dbName
});
return response.value;
});
}
function getMilvusInstance(clientParams) {
const milvusConfig = clientParams != null ? clientParams : getDefaultConfig();
return new import_milvus2_sdk_node.MilvusClient(milvusConfig);
}
function getDefaultConfig() {
var _a;
const configOrAddress = (_a = process.env.MILVUS_URI) != null ? _a : "http://localhost:19530";
return {
address: configOrAddress,
token: "",
// Default token is an empty string
username: "",
// Default username is an empty string
password: ""
// Default password is an empty string
};
}
// Annotate the CommonJS export names for ESM import in node:
0 && (module.exports = {
DocumentSchema,
createMilvusCollection,
deleteMilvusCollection,
describeMilvusCollection,
hasMilvusCollection,
insertMilvusData,
milvus,
milvusIndexer,
milvusIndexerRef,
milvusRetriever,
milvusRetrieverRef,
searchMilvusData
});
//# sourceMappingURL=index.js.map