estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

test: improve vector store tests (#3855)

Browse Source
This commit is contained in:
Bowen Liang
2024-04-26 19:18:42 +08:00
committed by GitHub
parent 4d66a86579
commit 045827043d
6 changed files with 58 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
api/core/rag/datasource/vdb/milvus/milvus_vector.py
View File

@@ -124,7 +124,7 @@ class MilvusVector(BaseVector):
if ids:
self._client.delete(collection_name=self._collection_name, pks=ids)
def delete_by_ids(self, doc_ids: list[str]) -> None:
def delete_by_ids(self, ids: list[str]) -> None:
alias = uuid4().hex
if self._client_config.secure:
uri = "https://" + str(self._client_config.host) + ":" + str(self._client_config.port)
@@ -136,7 +136,7 @@ class MilvusVector(BaseVector):
if utility.has_collection(self._collection_name, using=alias):
result = self._client.query(collection_name=self._collection_name,
filter=f'metadata["doc_id"] in {doc_ids}',
filter=f'metadata["doc_id"] in {ids}',
output_fields=["id"])
if result:
ids = [item["id"] for item in result]

4
api/core/rag/datasource/vdb/relyt/relyt_vector.py
View File

@@ -199,10 +199,10 @@ class RelytVector(BaseVector):
if ids:
self.delete_by_uuids(ids)
def delete_by_ids(self, doc_ids: list[str]) -> None:
def delete_by_ids(self, ids: list[str]) -> None:
with Session(self.client) as session:
ids_str = ','.join(f"'{doc_id}'" for doc_id in doc_ids)
ids_str = ','.join(f"'{doc_id}'" for doc_id in ids)
select_statement = sql_text(
f"""SELECT id FROM "{self._collection_name}" WHERE metadata->>'doc_id' in ({ids_str}); """
)