community[patch]: feat: add milvus partitionKey (#5983)

* feat: add milvus partitionKey * Format --------- Co-authored-by: Jacob Lee <jacoblee93@gmail.com>
langchain-ai · Jul 4, 2024 · ba89ca1 · ba89ca1
1 parent 6b0776e
commit ba89ca1
Showing 1 changed file with 33 additions and 5 deletions.
diff --git a/libs/langchain-community/src/vectorstores/milvus.ts b/libs/langchain-community/src/vectorstores/milvus.ts
@@ -32,6 +32,8 @@ export interface MilvusLibArgs {
   clientConfig?: ClientConfig;
   autoId?: boolean;
   indexCreateOptions?: IndexCreateOptions;
+  partitionKey?: string; // doc: https://milvus.io/docs/use-partition-key.md
+  partitionKeyMaxLength?: number;
 }
 
 export interface IndexCreateOptions {
@@ -74,6 +76,7 @@ const MILVUS_PRIMARY_FIELD_NAME = "langchain_primaryid";
 const MILVUS_VECTOR_FIELD_NAME = "langchain_vector";
 const MILVUS_TEXT_FIELD_NAME = "langchain_text";
 const MILVUS_COLLECTION_NAME_PREFIX = "langchain_col";
+const MILVUS_PARTITION_KEY_MAX_LENGTH = 512;
 
 /**
  * Default parameters for index searching.
@@ -126,6 +129,10 @@ export class Milvus extends VectorStore {
 
   textFieldMaxLength: number;
 
+  partitionKey?: string;
+
+  partitionKeyMaxLength?: number;
+
   fields: string[];
 
   client: MilvusClient;
@@ -146,6 +153,10 @@ export class Milvus extends VectorStore {
 
     this.textFieldMaxLength = args.textFieldMaxLength ?? 0;
 
+    this.partitionKey = args.partitionKey;
+    this.partitionKeyMaxLength =
+      args.partitionKeyMaxLength ?? MILVUS_PARTITION_KEY_MAX_LENGTH;
+
     this.fields = [];
 
     const url = args.url ?? getEnvironmentVariable("MILVUS_URL");
@@ -457,7 +468,13 @@ export class Milvus extends VectorStore {
   ): Promise<void> {
     const fieldList: FieldType[] = [];
 
-    fieldList.push(...createFieldTypeForMetadata(documents, this.primaryField));
+    fieldList.push(
+      ...createFieldTypeForMetadata(
+        documents,
+        this.primaryField,
+        this.partitionKey
+      )
+    );
 
     if (this.autoId) {
       fieldList.push({
@@ -500,6 +517,16 @@ export class Milvus extends VectorStore {
       }
     );
 
+    if (this.partitionKey) {
+      fieldList.push({
+        name: this.partitionKey,
+        description: "Partition key",
+        data_type: DataType.VarChar,
+        max_length: this.partitionKeyMaxLength,
+        is_partition_key: true,
+      });
+    }
+
     fieldList.forEach((field) => {
       if (!field.autoID) {
         this.fields.push(field.name);
@@ -677,7 +704,8 @@ export class Milvus extends VectorStore {
 
 function createFieldTypeForMetadata(
   documents: Document[],
-  primaryFieldName: string
+  primaryFieldName: string,
+  partitionKey?: string
 ): FieldType[] {
   const sampleMetadata = documents[0].metadata;
   let textFieldMaxLength = 0;
@@ -712,10 +740,10 @@ function createFieldTypeForMetadata(
   for (const [key, value] of Object.entries(sampleMetadata)) {
     const type = typeof value;
 
-    if (key === primaryFieldName) {
+    if (key === primaryFieldName || key === partitionKey) {
       /**
-       * skip primary field
-       * because we will create primary field in createCollection
+       * skip primary field and partition key
+       * because we will create primary field and partition key in createCollection
        *  */
     } else if (type === "string") {
       fields.push({