Update app.py

app.py

@@ -22,6 +22,7 @@ import pytesseract
 # Bypasses router.huggingface.co (PAID) entirely.
 # Uses api-inference.huggingface.co — completely free, same HF_TOKEN, no credits.
 import requests as _req
+import time as _time
 
 _FREE_API = "https://api-inference.huggingface.co/models/{model}"
 _MODELS = [
@@ -35,44 +36,58 @@ def call_llama(prompt: str, max_tokens: int = 2000) -> str:
     headers = {"Authorization": f"Bearer {token}"} if token else {}
 
     for model in _MODELS:
-        try:
-            # Format prompt correctly for each model family
-            if "zephyr" in model:
-                fmt = f"<|system|>\nYou are a precise legal document analyst.</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
-            elif "Phi-3" in model:
-                fmt = f"<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
-            else:
-                fmt = f"[INST] {prompt} [/INST]"
-
-            def _post(f):
-                return _req.post(
-                    _FREE_API.format(model=model), headers=headers,
-                    json={"inputs": f,
-                          "parameters": {"max_new_tokens": max_tokens,
-                                         "temperature": 0.7,
-                                         "return_full_text": False,
-                                         "do_sample": True}},
-                    timeout=120
+        # Try each model up to 4 times to handle cold starts
+        for attempt in range(4):
+            try:
+                if "zephyr" in model:
+                    fmt = f"<|system|>\nYou are a precise legal document analyst.</s>\n<|user|>\n{prompt}</s>\n<|assistant|>\n"
+                elif "Phi-3" in model:
+                    fmt = f"<|user|>\n{prompt}<|end|>\n<|assistant|>\n"
+                else:
+                    fmt = f"[INST] {prompt} [/INST]"
+
+                resp = _req.post(
+                    _FREE_API.format(model=model),
+                    headers=headers,
+                    json={
+                        "inputs": fmt,
+                        "parameters": {
+                            "max_new_tokens": max_tokens,
+                            "temperature": 0.7,
+                            "return_full_text": False,
+                            "do_sample": True,
+                        },
+                        "options": {"wait_for_model": True}
+                    },
+                    timeout=180
                 )
 
-            resp = _post(fmt)
+                if resp.status_code == 200:
+                    result = resp.json()
+                    if isinstance(result, list) and result:
+                        text = result[0].get("generated_text", "").strip()
+                        if text:
+                            return text
 
-            # Model loading — wait and retry once
-            if resp.status_code == 503:
-                import time; time.sleep(20)
-                resp = _post(fmt)
+                elif resp.status_code == 503:
+                    # Still loading — wait longer each attempt
+                    _time.sleep(30 * (attempt + 1))
+                    continue
 
-            if resp.status_code == 200:
-                result = resp.json()
-                if isinstance(result, list) and result:
-                    text = result[0].get("generated_text", "").strip()
-                    if text:
-                        return text
+                elif resp.status_code in [429, 402]:
+                    # Rate limit or payment — try next model immediately
+                    break
 
-        except Exception:
-            continue
+                else:
+                    # Other error — short wait then retry
+                    _time.sleep(10)
+                    continue
+
+            except Exception:
+                _time.sleep(10)
+                continue
 
-    return "Could not reach AI. Make sure HF_TOKEN is set in your Space secrets."
+    return "Could not reach AI. The models may be loading — please try again in 30 seconds."
 
 
 # ── FILE EXTRACTION ───────────────────────────────────────────────────────────