diff --git a/testing/constraints-3.14.txt b/testing/constraints-3.14.txt index a1fae03754..e0347baa13 100644 --- a/testing/constraints-3.14.txt +++ b/testing/constraints-3.14.txt @@ -13,4 +13,4 @@ ipython==8.22.2 # Pinned to unbreak TypeAliasType import error google-adk==0.0.2 google-genai>=1.10.0 google-vizier==0.1.21 -pyarrow>=18.0.0 \ No newline at end of file +pyarrow>=22.0.0 \ No newline at end of file diff --git a/tests/unit/aiplatform/test_prediction.py b/tests/unit/aiplatform/test_prediction.py index a1b49f9862..2b380f95fb 100644 --- a/tests/unit/aiplatform/test_prediction.py +++ b/tests/unit/aiplatform/test_prediction.py @@ -3287,7 +3287,10 @@ def test_health(self, model_server_env_mock, importlib_import_module_mock_twice) assert response.status_code == 200 - def test_predict(self, model_server_env_mock, importlib_import_module_mock_twice): + @pytest.mark.asyncio + async def test_predict( + self, model_server_env_mock, importlib_import_module_mock_twice + ): model_server = CprModelServer() client = TestClient(model_server.app) diff --git a/tests/unit/vertex_rag/test_rag_data.py b/tests/unit/vertex_rag/test_rag_data.py index 94f6c35bf9..67d3bfdd82 100644 --- a/tests/unit/vertex_rag/test_rag_data.py +++ b/tests/unit/vertex_rag/test_rag_data.py @@ -419,9 +419,10 @@ def create_transformation_config( def rag_corpus_eq(returned_corpus, expected_corpus): assert returned_corpus.name == expected_corpus.name assert returned_corpus.display_name == expected_corpus.display_name - assert returned_corpus.backend_config.__eq__(expected_corpus.backend_config) - assert returned_corpus.vertex_ai_search_config.__eq__( - expected_corpus.vertex_ai_search_config + assert returned_corpus.backend_config == expected_corpus.backend_config + assert ( + returned_corpus.vertex_ai_search_config + == expected_corpus.vertex_ai_search_config ) @@ -464,8 +465,8 @@ def import_files_request_eq(returned_request, expected_request): def rag_engine_config_eq(returned_config, expected_config): assert returned_config.name == expected_config.name - assert returned_config.rag_managed_db_config.__eq__( - expected_config.rag_managed_db_config + assert ( + returned_config.rag_managed_db_config == expected_config.rag_managed_db_config ) diff --git a/tests/unit/vertex_rag/test_rag_data_preview.py b/tests/unit/vertex_rag/test_rag_data_preview.py index cbd6bdd30c..1999ad3d0e 100644 --- a/tests/unit/vertex_rag/test_rag_data_preview.py +++ b/tests/unit/vertex_rag/test_rag_data_preview.py @@ -674,12 +674,13 @@ def create_transformation_config( def rag_corpus_eq(returned_corpus, expected_corpus): assert returned_corpus.name == expected_corpus.name assert returned_corpus.display_name == expected_corpus.display_name - assert returned_corpus.vector_db.__eq__(expected_corpus.vector_db) - assert returned_corpus.backend_config.__eq__(expected_corpus.backend_config) - assert returned_corpus.vertex_ai_search_config.__eq__( - expected_corpus.vertex_ai_search_config + assert returned_corpus.vector_db == expected_corpus.vector_db + assert returned_corpus.backend_config == expected_corpus.backend_config + assert ( + returned_corpus.vertex_ai_search_config + == expected_corpus.vertex_ai_search_config ) - assert returned_corpus.corpus_type_config.__eq__(expected_corpus.corpus_type_config) + assert returned_corpus.corpus_type_config == expected_corpus.corpus_type_config def rag_file_eq(returned_file, expected_file): @@ -729,8 +730,8 @@ def import_files_request_eq(returned_request, expected_request): def rag_engine_config_eq(returned_config, expected_config): assert returned_config.name == expected_config.name - assert returned_config.rag_managed_db_config.__eq__( - expected_config.rag_managed_db_config + assert ( + returned_config.rag_managed_db_config == expected_config.rag_managed_db_config ) diff --git a/vertexai/preview/rag/utils/_gapic_utils.py b/vertexai/preview/rag/utils/_gapic_utils.py index 4730c37b95..76cb4686c9 100644 --- a/vertexai/preview/rag/utils/_gapic_utils.py +++ b/vertexai/preview/rag/utils/_gapic_utils.py @@ -185,6 +185,24 @@ def _check_rag_embedding_model_config( return gapic_vector_db.rag_embedding_model_config.ByteSize() > 0 +def _check_document_corpus( + gapic_corpus_type_config: GapicRagCorpus.CorpusTypeConfig, +) -> bool: + try: + return gapic_corpus_type_config.__contains__("document_corpus") + except AttributeError: + return gapic_corpus_type_config.document_corpus.ByteSize() > 0 + + +def _check_memory_corpus( + gapic_corpus_type_config: GapicRagCorpus.CorpusTypeConfig, +) -> bool: + try: + return gapic_corpus_type_config.__contains__("memory_corpus") + except AttributeError: + return gapic_corpus_type_config.memory_corpus.ByteSize() > 0 + + def _convert_gapic_to_rag_managed_db( gapic_rag_managed_db: GapicRagVectorDbConfig.RagManagedDb, ) -> RagManagedDb: @@ -252,10 +270,10 @@ def convert_gapic_to_vector_db( def convert_gapic_to_vertex_ai_search_config( - gapic_vertex_ai_search_config: VertexAiSearchConfig, -) -> VertexAiSearchConfig: + gapic_vertex_ai_search_config: GapicVertexAiSearchConfig, +) -> Optional[VertexAiSearchConfig]: """Convert Gapic VertexAiSearchConfig to VertexAiSearchConfig.""" - if gapic_vertex_ai_search_config.serving_config: + if gapic_vertex_ai_search_config.ByteSize() > 0: return VertexAiSearchConfig( serving_config=gapic_vertex_ai_search_config.serving_config, ) @@ -293,6 +311,8 @@ def convert_gapic_to_backend_config( gapic_vector_db: GapicRagVectorDbConfig, ) -> RagVectorDbConfig: """Convert Gapic RagVectorDbConfig to VertexVectorSearch, Pinecone, or RagManagedDb.""" + if not gapic_vector_db or not gapic_vector_db.ByteSize(): + return None vector_config = RagVectorDbConfig() if _check_pinecone(gapic_vector_db): vector_config.vector_db = Pinecone( @@ -314,6 +334,11 @@ def convert_gapic_to_backend_config( gapic_vector_db.rag_embedding_model_config ) ) + if ( + vector_config.vector_db is None + and vector_config.rag_embedding_model_config is None + ): + return None return vector_config @@ -321,9 +346,9 @@ def convert_gapic_to_rag_corpus_type_config( gapic_rag_corpus_type_config: GapicRagCorpus.CorpusTypeConfig, ) -> RagCorpusTypeConfig: """Convert GapicRagCorpus.CorpusTypeConfig to RagCorpusTypeConfig.""" - if gapic_rag_corpus_type_config.document_corpus: + if _check_document_corpus(gapic_rag_corpus_type_config): return RagCorpusTypeConfig(corpus_type_config=DocumentCorpus()) - elif gapic_rag_corpus_type_config.memory_corpus: + elif _check_memory_corpus(gapic_rag_corpus_type_config): return RagCorpusTypeConfig( corpus_type_config=MemoryCorpus( llm_parser=LlmParserConfig( diff --git a/vertexai/rag/utils/_gapic_utils.py b/vertexai/rag/utils/_gapic_utils.py index 3ee39a7a0f..d44f00dd73 100644 --- a/vertexai/rag/utils/_gapic_utils.py +++ b/vertexai/rag/utils/_gapic_utils.py @@ -175,6 +175,8 @@ def convert_gapic_to_backend_config( gapic_vector_db: GapicRagVectorDbConfig, ) -> RagVectorDbConfig: """Convert Gapic RagVectorDbConfig to VertexVectorSearch, Pinecone, or RagManagedDb.""" + if not gapic_vector_db: + return None vector_config = RagVectorDbConfig() if _check_pinecone(gapic_vector_db): vector_config.vector_db = Pinecone( @@ -198,10 +200,10 @@ def convert_gapic_to_backend_config( def convert_gapic_to_vertex_ai_search_config( - gapic_vertex_ai_search_config: VertexAiSearchConfig, -) -> VertexAiSearchConfig: + gapic_vertex_ai_search_config: GapicVertexAiSearchConfig, +) -> Optional[VertexAiSearchConfig]: """Convert Gapic VertexAiSearchConfig to VertexAiSearchConfig.""" - if gapic_vertex_ai_search_config.serving_config: + if gapic_vertex_ai_search_config.ByteSize() > 0: return VertexAiSearchConfig( serving_config=gapic_vertex_ai_search_config.serving_config, )