fix: handle AI quota limits and configurable models

Resolves the "Resource Exhausted" error when using a standard API key
by defaulting to the correct embedding model (`embedding-001`). Also
introduces environment variables for configuring the chat and embedding
models, and improves exception handling to correctly catch and report
rate limiting and quota issues to the user.

Author: northdpole

.env.example

@@ -38,6 +38,9 @@ LOGIN_ALLOWED_DOMAINS=example.com
 # GCP
 
 GCP_NATIVE=false
+GEMINI_API_KEY=your-gemini-api-key
+VERTEX_CHAT_MODEL=gemini-2.0-flash
+VERTEX_EMBED_CONTENT_MODEL=embedding-001
 
 # Spreadsheet Auth
 

application/prompt_client/vertex_prompt_client.py

@@ -131,7 +131,7 @@ class VertexPromptClient:
 
     def __init__(self) -> None:
         self.client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
-        self.model_name = "gemini-2.0-flash"
+        self.model_name = os.environ.get("VERTEX_CHAT_MODEL", "gemini-2.0-flash")
 
     def get_model_name(self) -> str:
         """Return the model name being used."""
@@ -153,7 +153,11 @@ def _with_genai_rate_limit_retry(
         for attempt in range(max_retries + 1):
             try:
                 return fn()
-            except genai.errors.ClientError as e:
+            except (
+                genai.errors.ClientError,
+                googleExceptions.GoogleAPICallError,
+                grpc.RpcError,
+            ) as e:
                 _log_genai_client_error(context, e)
                 if not _is_genai_rate_limit_error(e) or attempt >= max_retries:
                     raise
@@ -205,9 +209,7 @@ def _truncate_one(t: str) -> str:
                 result = self.client.models.embed_content(
                     # Use a stable embeddings model that is supported for `embedContent`
                     # across API versions.
-                    model=os.environ.get(
-                        "VERTEX_EMBED_CONTENT_MODEL", "gemini-embedding-001"
-                    ),
+                    model=os.environ.get("VERTEX_EMBED_CONTENT_MODEL", "embedding-001"),
                     contents=texts if is_batch else texts[0],
                     config=types.EmbedContentConfig(task_type="SEMANTIC_SIMILARITY"),
                 )
@@ -217,7 +219,11 @@ def _truncate_one(t: str) -> str:
                 if is_batch:
                     return [emb.values for emb in result.embeddings]
                 return result.embeddings[0].values
-            except genai.errors.ClientError as e:
+            except (
+                genai.errors.ClientError,
+                googleExceptions.GoogleAPICallError,
+                grpc.RpcError,
+            ) as e:
                 _log_genai_client_error("embed_content", e)
                 if not _is_genai_rate_limit_error(e) or attempt >= max_retries:
                     raise
@@ -262,7 +268,7 @@ def create_chat_completion(self, prompt, closest_object_str) -> str:
 
         def _call() -> Any:
             response = self.client.models.generate_content(
-                model="gemini-2.0-flash",
+                model=self.model_name,
                 contents=msg,
                 config=types.GenerateContentConfig(
                     max_output_tokens=MAX_OUTPUT_TOKENS, temperature=0.5
@@ -279,7 +285,7 @@ def query_llm(self, raw_question: str) -> str:
 
         def _call() -> Any:
             response = self.client.models.generate_content(
-                model="gemini-2.0-flash",
+                model=self.model_name,
                 contents=msg,
                 config=types.GenerateContentConfig(
                     max_output_tokens=MAX_OUTPUT_TOKENS, temperature=0.5

application/tests/chat_completion_test.py

@@ -50,7 +50,7 @@ def test_completion_returns_503_json_on_gemini_429(self) -> None:
         self.assertIn("error", data)
         self.assertIn("rate-limited", data["error"])
 
-    def test_completion_propagates_non_429_genai_error(self) -> None:
+    def test_completion_returns_500_on_non_429_genai_error(self) -> None:
         os.environ["NO_LOGIN"] = "1"
         err = genai_errors.ClientError(
             400,
@@ -66,13 +66,15 @@ def test_completion_propagates_non_429_genai_error(self) -> None:
         with patch("application.prompt_client.prompt_client.PromptHandler") as mock_ph:
             mock_ph.return_value.generate_text.side_effect = err
             with self.app.test_client() as client:
-                with self.assertRaises(genai_errors.ClientError) as ctx:
-                    client.post(
-                        "/rest/v1/completion",
-                        json={"prompt": "test"},
-                        content_type="application/json",
-                    )
-                self.assertIs(ctx.exception, err)
+                response = client.post(
+                    "/rest/v1/completion",
+                    json={"prompt": "test"},
+                    content_type="application/json",
+                )
+        self.assertEqual(500, response.status_code)
+        data = json.loads(response.data)
+        self.assertIn("error", data)
+        self.assertIn("AI Service Error", data["error"])
 
 
 if __name__ == "__main__":

application/web/web_main.py

@@ -919,26 +919,40 @@ def chat_cre() -> Any:
     database = db.Node_collection()
     # Lazy import to avoid loading heavy prompt/ML dependencies at web boot.
     from google.genai import errors as genai_errors
-    from application.prompt_client import prompt_client
+    from google.api_core import exceptions as googleExceptions
+    import grpc
+    from application.prompt_client import prompt_client, vertex_prompt_client
 
     prompt = prompt_client.PromptHandler(database)
     try:
         response = prompt.generate_text(message.get("prompt"))
-    except genai_errors.ClientError as e:
-        # google.genai APIError uses ``code`` (HTTP status), not ``status_code``.
-        if getattr(e, "code", None) == 429:
+    except (
+        genai_errors.ClientError,
+        googleExceptions.GoogleAPICallError,
+        grpc.RpcError,
+    ) as e:
+        if vertex_prompt_client._is_genai_rate_limit_error(e):
             return (
                 jsonify(
                     {
                         "error": (
-                            "The AI service is temporarily rate-limited. "
+                            "The AI service is temporarily rate-limited or out of quota. "
                             "Please try again in a minute."
                         )
                     }
                 ),
                 503,
             )
-        raise
+        return (
+            jsonify({"error": f"AI Service Error: {str(e)}"}),
+            500,
+        )
+    except Exception as e:
+        logger.exception("Unexpected error during chatbot completion")
+        return (
+            jsonify({"error": f"An unexpected error occurred: {str(e)}"}),
+            500,
+        )
     return jsonify(response)