community[patch]: Make some functions work with Milvus (#10695)

**Description**
Make some functions work with Milvus:
1. get_ids: Get primary keys by field in the metadata
2. delete: Delete one or more entities by ids
3. upsert: Update/Insert one or more entities

**Issue**
None
**Dependencies**
None
**Tag maintainer:**
@hwchase17 
**Twitter handle:**
None

---------

Co-authored-by: HoaNQ9 <hoanq.1811@gmail.com>
Co-authored-by: Erick Friis <erick@langchain.dev>
This commit is contained in:
Quang Hoa
2024-02-10 06:21:31 +07:00
committed by GitHub
parent c9999557bf
commit 54c1fb3f25
3 changed files with 188 additions and 38 deletions

View File

@@ -989,3 +989,64 @@ class Milvus(VectorStore):
page_content=data.pop(self._text_field),
metadata=data.pop(self._metadata_field) if self._metadata_field else data,
)
def get_pks(self, expr: str, **kwargs: Any) -> List[int] | None:
"""Get primary keys with expression
Args:
expr: Expression - E.g: "id in [1, 2]", or "title LIKE 'Abc%'"
Returns:
List[int]: List of IDs (Primary Keys)
"""
from pymilvus import MilvusException
if self.col is None:
logger.debug("No existing collection to get pk.")
return None
try:
query_result = self.col.query(
expr=expr, output_fields=[self._primary_field]
)
except MilvusException as exc:
logger.error("Failed to get ids: %s error: %s", self.collection_name, exc)
raise exc
pks = [item.get(self._primary_field) for item in query_result]
return pks
def upsert(
self,
ids: Optional[List[str]] = None,
documents: List[Document] | None = None,
**kwargs: Any,
) -> List[str] | None:
"""Update/Insert documents to the vectorstore.
Args:
ids: IDs to update - Let's call get_pks to get ids with expression \n
documents (List[Document]): Documents to add to the vectorstore.
Returns:
List[str]: IDs of the added texts.
"""
from pymilvus import MilvusException
if documents is None or len(documents) == 0:
logger.debug("No documents to upsert.")
return None
if ids is not None and len(ids):
try:
self.delete(ids=ids)
except MilvusException:
pass
try:
return self.add_documents(documents=documents)
except MilvusException as exc:
logger.error(
"Failed to upsert entities: %s error: %s", self.collection_name, exc
)
raise exc