Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion testing/constraints-3.14.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
pyarrow>=22.0.0
5 changes: 4 additions & 1 deletion tests/unit/aiplatform/test_prediction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
11 changes: 6 additions & 5 deletions tests/unit/vertex_rag/test_rag_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)


Expand Down Expand Up @@ -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
)


Expand Down
15 changes: 8 additions & 7 deletions tests/unit/vertex_rag/test_rag_data_preview.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
)


Expand Down
35 changes: 30 additions & 5 deletions vertexai/preview/rag/utils/_gapic_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -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(
Expand All @@ -314,16 +334,21 @@ 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


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(
Expand Down
8 changes: 5 additions & 3 deletions vertexai/rag/utils/_gapic_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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,
)
Expand Down