Change text pk by using the parameter primary_field

langchain-ai · Feb 2, 2024 · 6927e81 · 6927e81
1 parent fd26430
commit 6927e81
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/libs/community/langchain_community/vectorstores/milvus.py b/libs/community/langchain_community/vectorstores/milvus.py
@@ -1012,14 +1012,14 @@ def get_ids(self,
         try:
             query_result = self.col.query(
                     expr=expr,
-                    output_fields=["pk"]
+                    output_fields=[self._primary_field]
                 )
         except MilvusException as exc:
             logger.error(
                 "Failed to get ids: %s error: %s", self.collection_name, exc
             )
             raise exc
-        ids = [item.get("pk") for item in query_result]
+        ids = [item.get(self._primary_field) for item in query_result]
         return ids
 
     def upsert(self,