add budget/auth msg for API key

generation/iterative_local_analyzer.py

@@ -21,6 +21,8 @@ from generation.analyzer import Analyzer, CodeAct, CodeAction
 from retrieval.retriever import Metadata
 from utils import (
     clean,
+    is_budget_error,
+    API_MSG_BUDGET,
     generate_iterative_system_prompt,
     get_file_from_title,
     get_path_from_title,
@@ -454,7 +456,8 @@ class IterativeLocalAnalyzer(Analyzer):
                 thought_msg.metadata["status"] = "done"
                 thought_msg.metadata["title"] = "Error during analysis"
                 thought_msg.metadata["duration"] = time.time() - start_time
-                yield [thought_msg, "Unfortunately, there was an error during the LLM invocation. Please try again."]
+                error_msg = API_MSG_BUDGET if is_budget_error(e) else "Unfortunately, there was an error during the LLM invocation. Please try again."
+                yield [thought_msg, error_msg]
                 return
 
             success, response, is_final_answer = self.run_ai_generated_code(code)

generation/simple_local_analyzer_v2.py

@@ -20,6 +20,8 @@ from generation.analyzer import Analyzer, CodeAct, CodeAction
 from retrieval.retriever import Metadata
 from utils import (
     clean,
+    is_budget_error,
+    API_MSG_BUDGET,
     generate_system_prompt_v2,
     get_file_from_title,
     get_path_from_title,
@@ -424,7 +426,8 @@ class SimpleLocalAnalyzerV2(Analyzer):
                 thought_msg.metadata["status"] = "done"
                 thought_msg.metadata["title"] = "Error during analysis"
                 thought_msg.metadata["duration"] = time.time() - start_time
-                yield [thought_msg, "Unfortunately, there was an error during the LLM invocation. Please try again."]
+                error_msg = API_MSG_BUDGET if is_budget_error(e) else "Unfortunately, there was an error during the LLM invocation. Please try again."
+                yield [thought_msg, error_msg]
                 return
 
             success, response = self.run_ai_generated_code(code)

main.py

@@ -31,7 +31,7 @@ from retrieval.agentic_retriever import AgenticRetriever
 from retrieval.knn_retriever import KNNRetriever
 from retrieval.retriever import Retriever
 from retrieval.verified_retriever import VerifiedRetriever
-from utils import SUPPORTED_LLMS, get_llm_client, init_mappings, download_dataset_file, get_file_from_title, get_path_from_title
+from utils import SUPPORTED_LLMS, get_llm_client, init_mappings, download_dataset_file, get_file_from_title, get_path_from_title, is_budget_error, API_MSG_BUDGET
 
 
 class RetrievalCheck(BaseModel):
@@ -107,7 +107,7 @@ class OGD4All():
                     thought_msg_retrieval.content = "An error occurred while retrieving datasets."
                     thought_msg_retrieval.metadata["status"] = "done"
                     thought_msg_retrieval.metadata["title"] = "Retrieval failed"
-                    error_msg = "I'm sorry, an error occurred during the retrieval of relevant datasets. Please try again."
+                    error_msg = API_MSG_BUDGET if is_budget_error(e) else "I'm sorry, an error occurred during the retrieval of relevant datasets. Please try again."
                     self.reset = True
                     yield [thought_msg_retrieval, error_msg], updated_map
                     return
@@ -195,7 +195,8 @@ class OGD4All():
                     retrieval_check = self.retrieval_check_client.invoke(messages)
                 except Exception as e:
                     log.error("Error during retrieval check:", exc_info=True)
-                    yield "An error occurred while checking whether additional datasets are required.", updated_map
+                    error_msg = API_MSG_BUDGET if is_budget_error(e) else "An error occurred while checking whether additional datasets are required."
+                    yield error_msg, updated_map
                     return
 
                 if retrieval_check.retrievalRequired:
@@ -344,7 +345,8 @@ class OGD4All():
         except Exception as e:
             log.error("Caught an exception in chat_fn: %s", e, exc_info=True, backtrace=True, diagnose=True)
             self.finalize()
-            yield gr.ChatMessage(role="assistant", content="I am sorry, there has been an error processing your request. Please try again."), updated_map
+            error_msg = API_MSG_BUDGET if is_budget_error(e) else "I am sorry, there has been an error processing your request. Please try again."
+            yield gr.ChatMessage(role="assistant", content=error_msg), updated_map
             self.reset = True
             return
     

utils.py

@@ -3,6 +3,7 @@ import os
 import shutil
 import pandas as pd
 import logging
+import openai
 import base64
 import mimetypes
 import pymupdf4llm
@@ -554,6 +555,22 @@ class ChatOpenRouter(ChatOpenAI):
         super().__init__(base_url="https://openrouter.ai/api/v1", openai_api_key=openai_api_key, **kwargs)
 
 
+API_MSG_BUDGET = "The LLM API key powering this demo has exceeded its budget or is no longer valid. "
+
+
+def is_budget_error(e: Exception) -> bool:
+    """Return True if the exception is a permanent API failure (quota exhausted or key invalid/expired)."""
+    if isinstance(e, (openai.AuthenticationError, openai.PermissionDeniedError)):
+        return True
+
+    if isinstance(e, openai.RateLimitError):
+        body = getattr(e, "body", None) or {}
+        code = body.get("error", {}).get("code", "") if isinstance(body, dict) else ""
+        return code == "insufficient_quota" or "insufficient_quota" in str(e)
+
+    return False
+
+
 def get_llm_client(llm_name: str, temperature: float = 0.0):
     """
     Returns an LLM client based on the provided name.