Upload folder using huggingface_hub

sanatan_assistant.py

@@ -3,6 +3,10 @@ from typing import Any, Literal
 from dotenv import load_dotenv
 from config import SanatanConfig
 from db import MetadataWhereClause, SanatanDatabase
+import time
+import random
+import logging
+from typing import Optional
 
 load_dotenv(override=True)
 logger = logging.getLogger(__name__)
@@ -17,6 +21,72 @@ allowedScriptureTitles = Literal[
     *[scripture["title"] for scripture in sanatanConfig.scriptures]
 ]
 
+MAX_RETRIES = 3
+BASE_BACKOFF = 0.8  # seconds
+MAX_BACKOFF = 5.0
+
+class RateLimitedError(Exception):
+    pass
+
+
+def _is_rate_limit_error(exc: Exception) -> bool:
+    msg = str(exc).lower()
+    return (
+        "429" in msg
+        or "rate limit" in msg
+        or "<html" in msg  # HF sends HTML error pages
+    )
+
+
+def safe_search(
+    *,
+    collection_name,
+    query,
+    metadata_where_clause,
+    n_results,
+    search_type,
+):
+    last_exc: Optional[Exception] = None
+
+    for attempt in range(1, MAX_RETRIES + 1):
+        try:
+            return sanatanDatabase.search(
+                collection_name=collection_name,
+                query=query,
+                metadata_where_clause=metadata_where_clause,
+                n_results=n_results,
+                search_type=search_type,
+            )
+
+        except Exception as exc:
+            last_exc = exc
+
+            if not _is_rate_limit_error(exc):
+                # Non-transient error → fail fast
+                logger.exception("Non-retriable DB error")
+                raise
+
+            # Rate limit / transient error
+            if attempt >= MAX_RETRIES:
+                break
+
+            backoff = min(
+                BASE_BACKOFF * (2 ** (attempt - 1)) + random.uniform(0, 0.2),
+                MAX_BACKOFF,
+            )
+
+            logger.warning(
+                "DB rate-limited (attempt %s/%s). Retrying in %.2fs",
+                attempt,
+                MAX_RETRIES,
+                backoff,
+            )
+            time.sleep(backoff)
+
+    logger.error("DB search failed after retries")
+    raise RateLimitedError("Database temporarily unavailable") from last_exc
+
+
 def format_scripture_answer(
     collection_name: allowedCollections, question: str, query_tool_output: str
 ):
@@ -146,13 +216,21 @@ def query(
     except:
         raise
 
-    response = sanatanDatabase.search(
-        collection_name=collection_name,
-        query=query,
-        metadata_where_clause=metadata_where_clause,
-        n_results=n_results,
-        search_type=search_type,
-    )
+    try:
+        response = safe_search(
+            collection_name=collection_name,
+            query=query,
+            metadata_where_clause=metadata_where_clause,
+            n_results=n_results,
+            search_type=search_type,
+        )
+    except RateLimitedError:
+        logger.warning(
+            "Falling back due to DB rate limit: collection=%s search_type=%s",
+            collection_name,
+            search_type,
+        )
+        return ""
 
     return "\n\n".join(
         f"Document: {doc}\nMetadata: {meta}\nID: {id_}"