diff --git a/google/generativeai/client.py b/google/generativeai/client.py
index bd45201d..3d639239 100644
--- a/google/generativeai/client.py
+++ b/google/generativeai/client.py
@@ -27,7 +27,6 @@
     __version__ = "0.0.0"
 
 USER_AGENT = "genai-py"
-GENAI_API_DISCOVERY_URL = "https://generativelanguage.googleapis.com/$discovery/rest"
 
 
 class FileServiceClient(glm.FileServiceClient):
@@ -43,7 +42,7 @@ def _setup_discovery_api(self):
         request = googleapiclient.http.HttpRequest(
             http=httplib2.Http(),
             postproc=lambda resp, content: (resp, content),
-            uri=f"{GENAI_API_DISCOVERY_URL}?version=v1beta&key={api_key}",
+            uri=f"https://{self.api_endpoint}/$discovery/rest?version=v1beta&key={api_key}",
         )
         response, content = request.execute()