genkitx-milvus
Version:
Genkit AI framework plugin for Milvus vector database.
292 lines • 9.03 kB
JavaScript
var __defProp = Object.defineProperty;
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 __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 } from "@genkit-ai/ai/retriever";
import { genkitPlugin } from "genkit/plugin";
import { z, Document, indexerRef, retrieverRef } from "genkit";
import {
MilvusClient
} from "@zilliz/milvus2-sdk-node";
const DocumentSchema = z.object({
id: z.number(),
vector: z.array(z.number()),
metadata: z.record(z.string(), z.any()).optional()
});
const MilvusRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({
k: z.number().max(1e3).default(10),
filter: z.record(z.string(), z.any()).optional(),
limit: z.number().optional()
});
const MilvusIndexerOptionsSchema = z.null().optional();
function milvus(params) {
return 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 retrieverRef({
name: `milvus/${params.collectionName}`,
info: {
label: (_a = params.displayName) != null ? _a : `Milvus - ${params.collectionName}`
},
configSchema: MilvusRetrieverOptionsSchema.optional()
});
};
const milvusIndexerRef = (params) => {
var _a;
return 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 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 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
};
}
export {
DocumentSchema,
createMilvusCollection,
src_default as default,
deleteMilvusCollection,
describeMilvusCollection,
hasMilvusCollection,
insertMilvusData,
milvus,
milvusIndexer,
milvusIndexerRef,
milvusRetriever,
milvusRetrieverRef,
searchMilvusData
};
//# sourceMappingURL=index.mjs.map