Update app.py

app.py

@@ -14,7 +14,7 @@ import base64
 import asyncio
 import gc
 
-# --- NEW: Import OCR Engine & Prompts ---
+# --- Import OCR Engine & Prompts ---
 try:
     from ocr_engine import extract_text_from_file
     from prompts import get_ocr_extraction_prompt, get_agent_prompt
@@ -31,7 +31,7 @@ logger = logging.getLogger("mcp_server")
 # Attempt to import transformers
 TRANSFORMERS_AVAILABLE = False
 try:
-    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+    from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
     TRANSFORMERS_AVAILABLE = True
 except Exception as e:
     logger.warning("transformers not available: %s", e)
@@ -63,14 +63,10 @@ _init_analytics()
 # Regex JSON Extractor
 # ----------------------------
 def extract_json_safely(text: str) -> Optional[Any]:
-    """
-    Extracts JSON from text even if the model adds conversational filler.
-    """
     try:
         return json.loads(text)
     except:
         pass
-    
     try:
         match = re.search(r'(\{.*\}|\[.*\])', text, re.DOTALL)
         if match:
@@ -81,7 +77,7 @@ def extract_json_safely(text: str) -> Optional[Any]:
     return None
 
 # ----------------------------
-# Local LLM loader
+# Local LLM loader (Lazy Loading Fix)
 # ----------------------------
 LLM_PIPELINE = None
 TOKENIZER = None
@@ -89,8 +85,14 @@ LOADED_MODEL_NAME = None
 
 def init_local_model():
     global LLM_PIPELINE, TOKENIZER, LOADED_MODEL_NAME
+    
+    # FIX 1: Check if already loaded to prevent double-memory usage
+    if LLM_PIPELINE is not None:
+        return
+
     if not LOCAL_MODEL or not TRANSFORMERS_AVAILABLE:
         return
+
     try:
         logger.info(f"Loading model: {LOCAL_MODEL}...")
         TOKENIZER = AutoTokenizer.from_pretrained(LOCAL_MODEL, trust_remote_code=True)
@@ -98,23 +100,28 @@ def init_local_model():
         
         LLM_PIPELINE = pipeline("text-generation", model=model, tokenizer=TOKENIZER)
         LOADED_MODEL_NAME = LOCAL_MODEL
-        logger.info("Model loaded.")
+        logger.info("Model loaded successfully.")
     except Exception as e:
         logger.error(f"Model load failed: {e}")
 
-init_local_model()
+# FIX 2: Removed global call to init_local_model() here.
+# It will now be called only in __main__ or lazily.
 
-# --- FIX APPLIED HERE ---
 def local_llm_generate(prompt: str, max_tokens: int = 512) -> Dict[str, Any]:
+    # FIX 3: Lazy load if accessed before main (safety net)
+    if LLM_PIPELINE is None:
+        init_local_model()
+        
     if LLM_PIPELINE is None:
         return {"text": "LLM not loaded.", "raw": None}
+    
     try:
-        # FIX: Added `use_cache=False` to resolve 'DynamicCache' object has no attribute 'seen_tokens'
+        # Includes Cache Fix from previous step
         out = LLM_PIPELINE(
             prompt, 
             max_new_tokens=max_tokens, 
             return_full_text=False, 
-            use_cache=False
+            use_cache=False 
         )
         text = out[0]["generated_text"] if out else ""
         return {"text": text, "raw": out}
@@ -144,15 +151,19 @@ def _get_valid_token_headers() -> dict:
     r = requests.post(token_url, params=params, timeout=20)
     if r.status_code == 200:
         return {"Authorization": f"Zoho-oauthtoken {r.json().get('access_token')}"}
-    raise RuntimeError(f"Token refresh failed: {r.text}")
+    # Don't crash entire app on token fail, just return None so tool can report it
+    logger.error(f"Token Refresh Failed: {r.text}")
+    return {}
 
 @mcp.tool()
 def authenticate_zoho() -> str:
-    _get_valid_token_headers(); return "Zoho token refreshed (ok)."
+    h = _get_valid_token_headers()
+    return "Zoho token refreshed." if h else "Failed to refresh token."
 
 @mcp.tool()
 def create_record(module_name: str, record_data: dict) -> str:
     headers = _get_valid_token_headers()
+    if not headers: return "Auth Error"
     url = f"{API_BASE}/{module_name}"
     r = requests.post(url, headers=headers, json={"data": [record_data]}, timeout=20)
     if r.status_code in (200, 201): return json.dumps(r.json(), ensure_ascii=False)
@@ -161,18 +172,21 @@ def create_record(module_name: str, record_data: dict) -> str:
 @mcp.tool()
 def get_records(module_name: str, page: int = 1, per_page: int = 200) -> list:
     headers = _get_valid_token_headers()
+    if not headers: return []
     r = requests.get(f"{API_BASE}/{module_name}", headers=headers, params={"page": page, "per_page": per_page})
     return r.json().get("data", []) if r.status_code == 200 else []
 
 @mcp.tool()
 def update_record(module_name: str, record_id: str, data: dict) -> str:
     headers = _get_valid_token_headers()
+    if not headers: return "Auth Error"
     r = requests.put(f"{API_BASE}/{module_name}/{record_id}", headers=headers, json={"data": [data]})
     return json.dumps(r.json()) if r.status_code == 200 else r.text
 
 @mcp.tool()
 def delete_record(module_name: str, record_id: str) -> str:
     headers = _get_valid_token_headers()
+    if not headers: return "Auth Error"
     r = requests.delete(f"{API_BASE}/{module_name}/{record_id}", headers=headers)
     return json.dumps(r.json()) if r.status_code == 200 else r.text
 
@@ -183,6 +197,7 @@ def _ensure_invoice_config():
 def create_invoice(data: dict) -> str:
     _ensure_invoice_config()
     headers = _get_valid_token_headers()
+    if not headers: return "Auth Error"
     params = {"organization_id": ORGANIZATION_ID}
     r = requests.post(f"{INVOICE_API_BASE}/invoices", headers=headers, params=params, json=data)
     if r.status_code in (200, 201): return json.dumps(r.json(), ensure_ascii=False)
@@ -191,6 +206,7 @@ def create_invoice(data: dict) -> str:
 def upload_invoice_attachment(invoice_id: str, file_path: str) -> str:
     if not os.path.exists(file_path): return "File not found"
     headers = _get_valid_token_headers()
+    if not headers: return "Auth Error"
     headers.pop("Content-Type", None)
     url = f"{INVOICE_API_BASE}/invoices/{invoice_id}/attachments"
     with open(file_path, "rb") as f:
@@ -200,9 +216,6 @@ def upload_invoice_attachment(invoice_id: str, file_path: str) -> str:
 
 @mcp.tool()
 def process_document(file_path: str, target_module: Optional[str] = "Contacts") -> dict:
-    """
-    Extracts data from file using OCR + LLM.
-    """
     try:
         if not os.path.exists(file_path):
              return {"status": "error", "error": "file not found"}
@@ -215,7 +228,7 @@ def process_document(file_path: str, target_module: Optional[str] = "Contacts")
         # 2. Use Prompt Template
         prompt = get_ocr_extraction_prompt(raw_text)
         
-        # 3. Generate (with cache fix applied in local_llm_generate)
+        # 3. Generate
         llm_out = local_llm_generate(prompt, max_tokens=300)
         extracted_text = llm_out.get("text", "")
         
@@ -334,7 +347,7 @@ def deepseek_response(message: str, file_path: Optional[str] = None, history: li
     history_text = "\n".join([f"User: {h[0]}\nBot: {h[1]}" for h in history])
     prompt = get_agent_prompt(history_text, ocr_context, message)
     
-    # 3. Generate (Cache Fix applies here too)
+    # 3. Generate
     gen = local_llm_generate(prompt, max_tokens=256)
     response_text = gen["text"]
     
@@ -384,6 +397,9 @@ def cleanup_event_loop():
 if __name__ == "__main__":
     cleanup_event_loop()
     
+    # FIX: Explicitly load model once here to prevent concurrent load attempts by Gradio
+    init_local_model()
+
     demo = gr.ChatInterface(
         fn=chat_handler, 
         multimodal=True,