genkitx-milvus/lib/index.mjs.map

Genkit AI framework plugin for Milvus vector database.

{"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<EmbedderCustomOptions>;\n  embedderOptions?: z.infer<EmbedderCustomOptions>;\n}\n\n/**\n * Milvus plugin that provides a milvus retriever and indexer.\n */\nexport function milvus<EmbedderCustomOptions extends z.ZodTypeAny>(\n  params: MilvusPluginOptions<EmbedderCustomOptions>[]\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<EmbedderCustomOptions extends z.ZodTypeAny>(\n  ai: Genkit,\n  params: MilvusPluginOptions<EmbedderCustomOptions>\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<string, any> } => !!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<EmbedderCustomOptions extends z.ZodTypeAny>(\n  ai: Genkit,\n  params: MilvusPluginOptions<EmbedderCustomOptions>\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"]}