{"version":3,"sources":["../src/index.ts"],"sourcesContent":["import { EmbedderArgument } from '@genkit-ai/ai/embedder';\nimport { CommonRetrieverOptionsSchema } from '@genkit-ai/ai/retriever';\nimport { GenkitPlugin, genkitPlugin } from 'genkit/plugin';\nimport { Genkit, z, Document, indexerRef, retrieverRef } from 'genkit';\nimport {\n MilvusClient,\n ClientConfig as MilvusClientParams,\n RowData,\n} from '@zilliz/milvus2-sdk-node';\n\nexport const DocumentSchema = z.object({\n id: z.number(),\n vector: z.array(z.number()),\n metadata: z.record(z.string(), z.any()).optional(),\n});\n\nconst MilvusRetrieverOptionsSchema = CommonRetrieverOptionsSchema.extend({\n k: z.number().max(1000).default(10),\n filter: z.record(z.string(), z.any()).optional(),\n limit: z.number().optional(),\n});\n\nconst MilvusIndexerOptionsSchema = z.null().optional();\n\ninterface MilvusPluginOptions<\n EmbedderCustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> {\n collectionName: string;\n dbName?: string;\n clientParams?: MilvusClientParams;\n embedder: EmbedderArgument;\n embedderOptions?: z.infer;\n}\n\n/**\n * Milvus plugin that provides a milvus retriever and indexer.\n */\nexport function milvus(\n params: MilvusPluginOptions[]\n): GenkitPlugin {\n return genkitPlugin('milvus', async (ai: Genkit) => {\n params.map((i) => milvusRetriever(ai, i));\n params.map((i) => milvusIndexer(ai, i));\n });\n}\n\nexport default milvus;\n\nexport const milvusRetrieverRef = (params: {\n collectionName: string;\n displayName?: string;\n}) => {\n return retrieverRef({\n name: `milvus/${params.collectionName}`,\n info: {\n label: params.displayName ?? `Milvus - ${params.collectionName}`,\n },\n configSchema: MilvusRetrieverOptionsSchema.optional(),\n });\n};\n\nexport const milvusIndexerRef = (params: {\n collectionName: string;\n displayName?: string;\n}) => {\n return indexerRef({\n name: `milvus/${params.collectionName}`,\n info: {\n label: params.displayName ?? `Milvus - ${params.collectionName}`,\n },\n configSchema: MilvusIndexerOptionsSchema.optional(),\n });\n};\n\n/**\n * Configures a Milvus vector store retriever.\n */\nexport function milvusRetriever(\n ai: Genkit,\n params: MilvusPluginOptions\n) {\n const { collectionName, embedder, embedderOptions, dbName } = params;\n const milvusConfig = params.clientParams ?? getDefaultConfig();\n\n return ai.defineRetriever(\n {\n name: `milvus/${collectionName}`,\n configSchema: MilvusRetrieverOptionsSchema,\n },\n async (content, options) => {\n const queryEmbeddings = await ai.embed({\n embedder,\n content,\n options: embedderOptions,\n });\n\n const isCollectionExists = await hasMilvusCollection({\n collectionName,\n dbName,\n clientParams: milvusConfig,\n });\n if (!isCollectionExists) {\n await createMilvusCollection({\n dbName,\n collectionName,\n dimension: queryEmbeddings.length,\n clientParams: milvusConfig,\n });\n }\n\n const response = await searchMilvusData({\n collectionName,\n dbName,\n query: queryEmbeddings.flatMap((e) => e.embedding),\n limit: options.limit,\n filter: JSON.stringify(options.filter),\n topk: options.k,\n clientParams: milvusConfig,\n });\n\n return {\n documents: response.results\n .map((result) => result['$meta'])\n .filter(\n (m): m is { document: string; metadata: Record } => !!m\n )\n .map((m) => {\n return Document.fromText(m.document, m.metadata).toJSON();\n }),\n };\n }\n );\n}\n\n/**\n * Configures a Milvus vector store indexer.\n */\nexport function milvusIndexer(\n ai: Genkit,\n params: MilvusPluginOptions\n) {\n const { collectionName, embedder, embedderOptions, dbName } = params;\n const milvusConfig = params.clientParams ?? getDefaultConfig();\n\n return ai.defineIndexer(\n {\n name: `milvus/${collectionName}`,\n configSchema: MilvusIndexerOptionsSchema.optional(),\n },\n async (docs) => {\n const embeddings = await Promise.all(\n docs.map((doc) =>\n ai.embed({\n embedder,\n content: doc,\n options: embedderOptions,\n })\n )\n );\n\n const entries = embeddings.map((value, i) => {\n const metadata: RowData = {\n ...docs[i].metadata,\n };\n\n return {\n vector: value,\n document: docs[i].text,\n metadata,\n };\n });\n\n const isCollectionExists = await hasMilvusCollection({\n collectionName,\n dbName,\n clientParams: milvusConfig,\n });\n if (!isCollectionExists) {\n await createMilvusCollection({\n dbName,\n collectionName,\n dimension: embeddings[0].length,\n clientParams: milvusConfig,\n });\n }\n\n await insertMilvusData({\n collectionName: collectionName,\n dbName,\n entries,\n clientParams: milvusConfig,\n });\n }\n );\n}\n\nexport async function createMilvusCollection(params: {\n collectionName: string;\n dimension: number;\n dbName?: string;\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dimension, dbName } = params;\n const milvus = getMilvusInstance(clientParams);\n\n return milvus.createCollection({\n collection_name: collectionName,\n db_name: dbName,\n dimension,\n auto_id: true,\n enable_dynamic_field: true,\n });\n}\n\nexport async function describeMilvusCollection(params: {\n collectionName: string;\n dbName?: string;\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dbName } = params;\n const milvus = getMilvusInstance(clientParams);\n\n return await milvus.describeCollection({\n collection_name: collectionName,\n db_name: dbName,\n });\n}\n\nexport async function insertMilvusData(params: {\n collectionName: string;\n dbName?: string;\n entries: RowData[];\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dbName, entries } = params;\n const milvus = getMilvusInstance(clientParams);\n\n return milvus.insert({\n collection_name: collectionName,\n db_name: dbName,\n fields_data: entries,\n });\n}\n\nexport async function searchMilvusData(params: {\n collectionName: string;\n dbName?: string;\n query: number[];\n limit?: number;\n filter?: string;\n topk?: number;\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dbName, query, ...restParams } = params;\n const milvus = getMilvusInstance(clientParams);\n\n return milvus.search({\n collection_name: collectionName,\n db_name: dbName,\n data: query,\n ...restParams,\n });\n}\n\nexport async function deleteMilvusCollection(params: {\n collectionName: string;\n dbName?: string;\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dbName } = params;\n const milvus = getMilvusInstance(clientParams);\n\n try {\n return await milvus.dropCollection({\n collection_name: collectionName,\n db_name: dbName,\n });\n } catch (error) {\n console.error('Failed to delete Milvus collection:', error);\n throw error;\n }\n}\n\nexport async function hasMilvusCollection(params: {\n collectionName: string;\n dbName?: string;\n clientParams?: MilvusClientParams;\n}) {\n const { clientParams, collectionName, dbName } = params;\n const milvus = getMilvusInstance(clientParams);\n\n const response = await milvus.hasCollection({\n collection_name: collectionName,\n db_name: dbName,\n });\n\n return response.value;\n}\n\nfunction getMilvusInstance(clientParams?: MilvusClientParams) {\n const milvusConfig = clientParams ?? getDefaultConfig();\n return new MilvusClient(milvusConfig);\n}\n\nfunction getDefaultConfig(): MilvusClientParams {\n const configOrAddress = process.env.MILVUS_URI ?? 'http://localhost:19530';\n\n return {\n address: configOrAddress,\n token: '', // Default token is an empty string\n username: '', // Default username is an empty string\n password: '', // Default password is an empty string\n };\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AACA,SAAS,oCAAoC;AAC7C,SAAuB,oBAAoB;AAC3C,SAAiB,GAAG,UAAU,YAAY,oBAAoB;AAC9D;AAAA,EACE;AAAA,OAGK;AAEA,MAAM,iBAAiB,EAAE,OAAO;AAAA,EACrC,IAAI,EAAE,OAAO;AAAA,EACb,QAAQ,EAAE,MAAM,EAAE,OAAO,CAAC;AAAA,EAC1B,UAAU,EAAE,OAAO,EAAE,OAAO,GAAG,EAAE,IAAI,CAAC,EAAE,SAAS;AACnD,CAAC;AAED,MAAM,+BAA+B,6BAA6B,OAAO;AAAA,EACvE,GAAG,EAAE,OAAO,EAAE,IAAI,GAAI,EAAE,QAAQ,EAAE;AAAA,EAClC,QAAQ,EAAE,OAAO,EAAE,OAAO,GAAG,EAAE,IAAI,CAAC,EAAE,SAAS;AAAA,EAC/C,OAAO,EAAE,OAAO,EAAE,SAAS;AAC7B,CAAC;AAED,MAAM,6BAA6B,EAAE,KAAK,EAAE,SAAS;AAe9C,SAAS,OACd,QACc;AACd,SAAO,aAAa,UAAU,CAAO,OAAe;AAClD,WAAO,IAAI,CAAC,MAAM,gBAAgB,IAAI,CAAC,CAAC;AACxC,WAAO,IAAI,CAAC,MAAM,cAAc,IAAI,CAAC,CAAC;AAAA,EACxC,EAAC;AACH;AAEA,IAAO,cAAQ;AAER,MAAM,qBAAqB,CAAC,WAG7B;AAnDN;AAoDE,SAAO,aAAa;AAAA,IAClB,MAAM,UAAU,OAAO,cAAc;AAAA,IACrC,MAAM;AAAA,MACJ,QAAO,YAAO,gBAAP,YAAsB,YAAY,OAAO,cAAc;AAAA,IAChE;AAAA,IACA,cAAc,6BAA6B,SAAS;AAAA,EACtD,CAAC;AACH;AAEO,MAAM,mBAAmB,CAAC,WAG3B;AAhEN;AAiEE,SAAO,WAAW;AAAA,IAChB,MAAM,UAAU,OAAO,cAAc;AAAA,IACrC,MAAM;AAAA,MACJ,QAAO,YAAO,gBAAP,YAAsB,YAAY,OAAO,cAAc;AAAA,IAChE;AAAA,IACA,cAAc,2BAA2B,SAAS;AAAA,EACpD,CAAC;AACH;AAKO,SAAS,gBACd,IACA,QACA;AAhFF;AAiFE,QAAM,EAAE,gBAAgB,UAAU,iBAAiB,OAAO,IAAI;AAC9D,QAAM,gBAAe,YAAO,iBAAP,YAAuB,iBAAiB;AAE7D,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,UAAU,cAAc;AAAA,MAC9B,cAAc;AAAA,IAChB;AAAA,IACA,CAAO,SAAS,YAAY;AAC1B,YAAM,kBAAkB,MAAM,GAAG,MAAM;AAAA,QACrC;AAAA,QACA;AAAA,QACA,SAAS;AAAA,MACX,CAAC;AAED,YAAM,qBAAqB,MAAM,oBAAoB;AAAA,QACnD;AAAA,QACA;AAAA,QACA,cAAc;AAAA,MAChB,CAAC;AACD,UAAI,CAAC,oBAAoB;AACvB,cAAM,uBAAuB;AAAA,UAC3B;AAAA,UACA;AAAA,UACA,WAAW,gBAAgB;AAAA,UAC3B,cAAc;AAAA,QAChB,CAAC;AAAA,MACH;AAEA,YAAM,WAAW,MAAM,iBAAiB;AAAA,QACtC;AAAA,QACA;AAAA,QACA,OAAO,gBAAgB,QAAQ,CAAC,MAAM,EAAE,SAAS;AAAA,QACjD,OAAO,QAAQ;AAAA,QACf,QAAQ,KAAK,UAAU,QAAQ,MAAM;AAAA,QACrC,MAAM,QAAQ;AAAA,QACd,cAAc;AAAA,MAChB,CAAC;AAED,aAAO;AAAA,QACL,WAAW,SAAS,QACjB,IAAI,CAAC,WAAW,OAAO,OAAO,CAAC,EAC/B;AAAA,UACC,CAAC,MAAgE,CAAC,CAAC;AAAA,QACrE,EACC,IAAI,CAAC,MAAM;AACV,iBAAO,SAAS,SAAS,EAAE,UAAU,EAAE,QAAQ,EAAE,OAAO;AAAA,QAC1D,CAAC;AAAA,MACL;AAAA,IACF;AAAA,EACF;AACF;AAKO,SAAS,cACd,IACA,QACA;AA5IF;AA6IE,QAAM,EAAE,gBAAgB,UAAU,iBAAiB,OAAO,IAAI;AAC9D,QAAM,gBAAe,YAAO,iBAAP,YAAuB,iBAAiB;AAE7D,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM,UAAU,cAAc;AAAA,MAC9B,cAAc,2BAA2B,SAAS;AAAA,IACpD;AAAA,IACA,CAAO,SAAS;AACd,YAAM,aAAa,MAAM,QAAQ;AAAA,QAC/B,KAAK;AAAA,UAAI,CAAC,QACR,GAAG,MAAM;AAAA,YACP;AAAA,YACA,SAAS;AAAA,YACT,SAAS;AAAA,UACX,CAAC;AAAA,QACH;AAAA,MACF;AAEA,YAAM,UAAU,WAAW,IAAI,CAAC,OAAO,MAAM;AAC3C,cAAM,WAAoB,mBACrB,KAAK,CAAC,EAAE;AAGb,eAAO;AAAA,UACL,QAAQ;AAAA,UACR,UAAU,KAAK,CAAC,EAAE;AAAA,UAClB;AAAA,QACF;AAAA,MACF,CAAC;AAED,YAAM,qBAAqB,MAAM,oBAAoB;AAAA,QACnD;AAAA,QACA;AAAA,QACA,cAAc;AAAA,MAChB,CAAC;AACD,UAAI,CAAC,oBAAoB;AACvB,cAAM,uBAAuB;AAAA,UAC3B;AAAA,UACA;AAAA,UACA,WAAW,WAAW,CAAC,EAAE;AAAA,UACzB,cAAc;AAAA,QAChB,CAAC;AAAA,MACH;AAEA,YAAM,iBAAiB;AAAA,QACrB;AAAA,QACA;AAAA,QACA;AAAA,QACA,cAAc;AAAA,MAChB,CAAC;AAAA,IACH;AAAA,EACF;AACF;AAEA,SAAsB,uBAAuB,QAK1C;AAAA;AACD,UAAM,EAAE,cAAc,gBAAgB,WAAW,OAAO,IAAI;AAC5D,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,WAAOA,QAAO,iBAAiB;AAAA,MAC7B,iBAAiB;AAAA,MACjB,SAAS;AAAA,MACT;AAAA,MACA,SAAS;AAAA,MACT,sBAAsB;AAAA,IACxB,CAAC;AAAA,EACH;AAAA;AAEA,SAAsB,yBAAyB,QAI5C;AAAA;AACD,UAAM,EAAE,cAAc,gBAAgB,OAAO,IAAI;AACjD,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,WAAO,MAAMA,QAAO,mBAAmB;AAAA,MACrC,iBAAiB;AAAA,MACjB,SAAS;AAAA,IACX,CAAC;AAAA,EACH;AAAA;AAEA,SAAsB,iBAAiB,QAKpC;AAAA;AACD,UAAM,EAAE,cAAc,gBAAgB,QAAQ,QAAQ,IAAI;AAC1D,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,WAAOA,QAAO,OAAO;AAAA,MACnB,iBAAiB;AAAA,MACjB,SAAS;AAAA,MACT,aAAa;AAAA,IACf,CAAC;AAAA,EACH;AAAA;AAEA,SAAsB,iBAAiB,QAQpC;AAAA;AACD,UAAuE,aAA/D,gBAAc,gBAAgB,QAAQ,MA7PhD,IA6PyE,IAAf,uBAAe,IAAf,CAAhD,gBAAc,kBAAgB,UAAQ;AAC9C,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,WAAOA,QAAO,OAAO;AAAA,MACnB,iBAAiB;AAAA,MACjB,SAAS;AAAA,MACT,MAAM;AAAA,OACH,WACJ;AAAA,EACH;AAAA;AAEA,SAAsB,uBAAuB,QAI1C;AAAA;AACD,UAAM,EAAE,cAAc,gBAAgB,OAAO,IAAI;AACjD,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,QAAI;AACF,aAAO,MAAMA,QAAO,eAAe;AAAA,QACjC,iBAAiB;AAAA,QACjB,SAAS;AAAA,MACX,CAAC;AAAA,IACH,SAAS,OAAO;AACd,cAAQ,MAAM,uCAAuC,KAAK;AAC1D,YAAM;AAAA,IACR;AAAA,EACF;AAAA;AAEA,SAAsB,oBAAoB,QAIvC;AAAA;AACD,UAAM,EAAE,cAAc,gBAAgB,OAAO,IAAI;AACjD,UAAMA,UAAS,kBAAkB,YAAY;AAE7C,UAAM,WAAW,MAAMA,QAAO,cAAc;AAAA,MAC1C,iBAAiB;AAAA,MACjB,SAAS;AAAA,IACX,CAAC;AAED,WAAO,SAAS;AAAA,EAClB;AAAA;AAEA,SAAS,kBAAkB,cAAmC;AAC5D,QAAM,eAAe,sCAAgB,iBAAiB;AACtD,SAAO,IAAI,aAAa,YAAY;AACtC;AAEA,SAAS,mBAAuC;AAhThD;AAiTE,QAAM,mBAAkB,aAAQ,IAAI,eAAZ,YAA0B;AAElD,SAAO;AAAA,IACL,SAAS;AAAA,IACT,OAAO;AAAA;AAAA,IACP,UAAU;AAAA;AAAA,IACV,UAAU;AAAA;AAAA,EACZ;AACF;","names":["milvus"]}