From bc5a214c200bb14e1053b74b2a4039bac4f6b7e5 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Mon, 7 Jul 2025 22:29:19 -0400 Subject: [PATCH] added tests for all methods Signed-off-by: Francisco Javier Arceo --- .../remote/vector_io/milvus/milvus.py | 45 +++++----- .../test_vector_io_openai_vector_stores.py | 82 +++++++++++++++++++ 2 files changed, 105 insertions(+), 22 deletions(-) diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index 486221601..d97582a13 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -379,6 +379,29 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP logger.error(f"Error loading openai vector store file {file_id} for store {store_id}: {e}") return {} + async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: + """Update vector store file metadata in Milvus database.""" + try: + if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"): + return + + file_data = [ + { + "store_file_id": f"{store_id}_{file_id}", + "store_id": store_id, + "file_id": file_id, + "file_info": json.dumps(file_info), + } + ] + await asyncio.to_thread( + self.client.upsert, + collection_name="openai_vector_store_files", + data=file_data, + ) + except Exception as e: + logger.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}") + raise + async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]: """Load vector store file contents from Milvus database.""" try: @@ -407,28 +430,6 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP logger.error(f"Error loading openai vector store file contents for {file_id} in store {store_id}: {e}") return [] - async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: - """Update vector store file metadata in Milvus database.""" - try: - if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"): - return - - file_data = [ - { - "store_file_id": f"{store_id}_{file_id}", - "store_id": store_id, - "file_id": file_id, - "file_info": json.dumps(file_info), - } - ] - await asyncio.to_thread( - self.client.upsert, - collection_name="openai_vector_store_files", - data=file_data, - ) - except Exception as e: - logger.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}") - raise async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None: """Delete vector store file metadata from Milvus database.""" diff --git a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py index 00040c46d..05a36149c 100644 --- a/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py +++ b/tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py @@ -335,3 +335,85 @@ async def test_save_openai_vector_store_file(milvus_vec_adapter, tmp_path_factor # validating we don't raise an exception await milvus_vec_adapter._save_openai_vector_store_file(store_id, file_id, file_info, file_contents) + +@pytest.mark.asyncio +async def test_update_openai_vector_store_file(milvus_vec_adapter, tmp_path_factory): + store_id = "vs_1234" + file_id = "file_1234" + + file_info = { + "id": file_id, + "status": "completed", + "vector_store_id": store_id, + "attributes": {}, + "filename": "test_file.txt", + "created_at": int(time.time()), + } + + file_contents = [ + {"content": "Test content", "chunk_metadata": {"chunk_id": "chunk_001"}, "metadata": {"file_id": file_id}} + ] + + await milvus_vec_adapter._save_openai_vector_store_file(store_id, file_id, file_info, file_contents) + + updated_file_info = file_info.copy() + updated_file_info["filename"] = "updated_test_file.txt" + + + await milvus_vec_adapter._update_openai_vector_store_file( + store_id, + file_id, + updated_file_info, + ) + + loaded_contents = await milvus_vec_adapter._load_openai_vector_store_file(store_id, file_id) + assert loaded_contents == updated_file_info + assert loaded_contents != file_info + +@pytest.mark.asyncio +async def test_load_openai_vector_store_file_contents(milvus_vec_adapter, tmp_path_factory): + store_id = "vs_1234" + file_id = "file_1234" + + file_info = { + "id": file_id, + "status": "completed", + "vector_store_id": store_id, + "attributes": {}, + "filename": "test_file.txt", + "created_at": int(time.time()), + } + + file_contents = [ + {"content": "Test content", "chunk_metadata": {"chunk_id": "chunk_001"}, "metadata": {"file_id": file_id}} + ] + + await milvus_vec_adapter._save_openai_vector_store_file(store_id, file_id, file_info, file_contents) + + loaded_contents = await milvus_vec_adapter._load_openai_vector_store_file_contents(store_id, file_id) + assert loaded_contents == file_contents + + +@pytest.mark.asyncio +async def test_delete_openai_vector_store_file_from_storage(milvus_vec_adapter, tmp_path_factory): + store_id = "vs_1234" + file_id = "file_1234" + + file_info = { + "id": file_id, + "status": "completed", + "vector_store_id": store_id, + "attributes": {}, + "filename": "test_file.txt", + "created_at": int(time.time()), + } + + file_contents = [ + {"content": "Test content", "chunk_metadata": {"chunk_id": "chunk_001"}, "metadata": {"file_id": file_id}} + ] + + await milvus_vec_adapter._save_openai_vector_store_file(store_id, file_id, file_info, file_contents) + await milvus_vec_adapter._delete_openai_vector_store_file_from_storage(store_id, file_id) + + loaded_contents = await milvus_vec_adapter._load_openai_vector_store_file_contents(store_id, file_id) + assert loaded_contents == [] \ No newline at end of file