diff --git a/agents/rag/setup_corpus.py b/agents/rag/setup_corpus.py index e0005dd..c2bc549 100644 --- a/agents/rag/setup_corpus.py +++ b/agents/rag/setup_corpus.py @@ -23,6 +23,10 @@ GCS_SOURCE = os.environ.get( f"gs://{PROJECT_ID}-agent-staging/rag-docs/", ) +# RagManagedDbConfig.Tier enum may not exist in older SDK versions. +# tier=1 == BASIC, tier=2 == ENTERPRISE +_TIER_BASIC = 1 + def ensure_serverless_engine_config() -> None: """ @@ -76,10 +80,9 @@ def get_or_create_corpus() -> RagCorpus: print(f"Creating corpus '{CORPUS_DISPLAY_NAME}' in {LOCATION} (serverless)...") corpus = rag.create_corpus( display_name=CORPUS_DISPLAY_NAME, - # Explicitly request serverless / basic managed DB tier backend_config=rag.RagVectorDbConfig( rag_managed_db=rag.RagManagedDbConfig( - tier=rag.RagManagedDbConfig.Tier.BASIC, + tier=_TIER_BASIC, # 1 = BASIC (serverless), 2 = ENTERPRISE ) ), )