diff --git a/agentplatform/_genai/rag.py b/agentplatform/_genai/rag.py
index f00db97a5e..90cf95a112 100644
--- a/agentplatform/_genai/rag.py
+++ b/agentplatform/_genai/rag.py
@@ -2227,7 +2227,7 @@ def _update_config(
     def retrieve_contexts(
         self,
         *,
-        vertex_rag_store: Optional[types.VertexRagStoreOrDict] = None,
+        vertex_rag_store: types.VertexRagStoreOrDict,
         query: types.RagQueryOrDict,
         config: Optional[types.RetrieveContextsConfigOrDict] = None,
     ) -> types.RetrieveContextsResponse:
@@ -3258,7 +3258,7 @@ async def _update_config(
     async def retrieve_contexts(
         self,
         *,
-        vertex_rag_store: Optional[types.VertexRagStoreOrDict] = None,
+        vertex_rag_store: types.VertexRagStoreOrDict,
         query: types.RagQueryOrDict,
         config: Optional[types.RetrieveContextsConfigOrDict] = None,
     ) -> types.RetrieveContextsResponse: