Update app.py

app.py

@@ -1,11 +1,12 @@
-# app.py — MCP server (single-file, refined)
-# - Robust JSON extraction/repair
-# - Safer model loading (use dtype param)
-# - Better logging and error handling
-# - Uses prompts.py functions
+# app.py — MCP server (refined)
+# Key improvements:
+# - Robust JSON extraction & repair
+# - Detailed debug logging, write raw LLM output to /tmp when parse fails
+# - Defensive LLM handling
+# - Uses your ocr_engine.extract_text_and_conf
 
 from mcp.server.fastmcp import FastMCP
-from typing import Optional, Any, Dict, List
+from typing import Optional, Any, Dict
 import requests
 import os
 import gradio as gr
@@ -13,324 +14,277 @@ import json
 import re
 import logging
 import gc
-import math
+import time
+import traceback
 
-# --- Import OCR Engine & Prompts ---
+# imports from local modules (these must exist)
+from ocr_engine import extract_text_and_conf
+from prompts import get_ocr_extraction_prompt, get_agent_prompt
+
+# config (must exist)
 try:
-    from ocr_engine import extract_text_and_conf
-    from prompts import get_ocr_extraction_prompt, get_agent_prompt
-except ImportError:
-    def extract_text_and_conf(path): return ("", 0.0)
-    def get_ocr_extraction_prompt(txt, page_count=1): return txt
-    def get_agent_prompt(h, u): return u
+    from config import CLIENT_ID, CLIENT_SECRET, REFRESH_TOKEN, API_BASE, INVOICE_API_BASE, ORGANIZATION_ID, LOCAL_MODEL
+except Exception as e:
+    raise SystemExit("Missing config.py or required keys. Error: " + str(e))
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("mcp_server")
 
-# --- Load Config ---
-try:
-    from config import (
-        CLIENT_ID, CLIENT_SECRET, REFRESH_TOKEN, API_BASE,
-        INVOICE_API_BASE, ORGANIZATION_ID, LOCAL_MODEL
-    )
-except Exception:
-    raise SystemExit("Config missing. Please create config.py with required keys.")
-
 mcp = FastMCP("ZohoCRMAgent")
 
-# --- Globals ---
 LLM_PIPELINE = None
 TOKENIZER = None
 
-# --- NEW: Evaluation / KPI Logic (Integrated OCR Score) ---
-def calculate_extraction_confidence(data: dict, ocr_score: float) -> dict:
-    semantic_score = 0
-    issues = []
-
-    # Structure baseline
-    semantic_score += 10
-
-    # Total Amount Check (30)
-    grand = data.get("totals", {}).get("grand_total") if isinstance(data.get("totals"), dict) else None
-    if grand is not None:
-        try:
-            float(str(grand))
-            semantic_score += 30
-        except:
-            issues.append("Missing/Invalid Total Amount")
-    else:
-        issues.append("Missing/Invalid Total Amount")
-
-    # Date Check (20)
-    date_str = data.get("invoice_date")
-    if date_str and isinstance(date_str, str) and len(date_str) >= 8:
-        semantic_score += 20
-    else:
-        issues.append("Missing Invoice Date")
-
-    # Line Items Check (30)
-    items = data.get("line_items", [])
-    if isinstance(items, list) and len(items) > 0:
-        if any((isinstance(i, dict) and (i.get("name") or i.get("description"))) for i in items):
-            semantic_score += 30
-        else:
-            semantic_score += 10
-            issues.append("Line Items missing descriptions")
-    else:
-        issues.append("No Line Items detected")
-
-    # Contact Name (10)
-    buyer = data.get("buyer", {}) or {}
-    if buyer.get("contact_name") or buyer.get("company_name"):
-        semantic_score += 10
-    else:
-        issues.append("Missing Buyer / Contact Name")
-
-    final_score = (semantic_score * 0.8) + (ocr_score * 0.2)
-    rating = "High" if final_score > 80 else ("Medium" if final_score > 50 else "Low")
-    if ocr_score < 60:
-        issues.append(f"Low OCR Confidence ({ocr_score}%) - Check image quality")
-
-    return {
-        "score": int(round(final_score)),
-        "ocr_score": int(round(ocr_score)),
-        "semantic_score": semantic_score,
-        "rating": rating,
-        "issues": issues
-    }
-
-# --- Robust JSON extraction & repair helpers ---
+# ---------------- JSON extraction helpers ----------------
 def _try_json_loads(text: str) -> Optional[Any]:
     try:
         return json.loads(text)
     except Exception:
         return None
 
-def _remove_wrapping_code_blocks(s: str) -> str:
+def _remove_code_fences(s: str) -> str:
     s = re.sub(r"```(?:json)?\s*", "", s, flags=re.IGNORECASE)
     s = re.sub(r"\s*```$", "", s, flags=re.IGNORECASE)
     return s.strip()
 
 def _attempt_simple_repairs(s: str) -> str:
-    # 1) Keep printable characters (allow newline/tab)
+    # keep printable chars
     s = "".join(ch for ch in s if (ch == "\n" or ch == "\t" or (32 <= ord(ch) <= 0x10FFFF)))
-    # 2) Remove trailing commas before } or ]
+    # remove trailing commas
     s = re.sub(r",\s*(\}|])", r"\1", s)
-    # 3) Convert lone single-quoted JSON to double quotes only if safe
+    # convert single quotes if double quotes not present
     if '"' not in s and "'" in s:
         s = s.replace("'", '"')
-    # 4) Remove assistant labels
-    s = re.sub(r"^(assistant:|response:)\s*", "", s, flags=re.IGNORECASE)
     return s
 
+def _dump_raw_llm_output(text: str) -> str:
+    """Dump raw LLM output to a timestamped file for debugging and return path."""
+    try:
+        ts = int(time.time())
+        path = f"/tmp/llm_output_{ts}.txt"
+        with open(path, "w", encoding="utf-8") as f:
+            f.write(text)
+        logger.info("Wrote raw LLM output to %s for debugging", path)
+        return path
+    except Exception as e:
+        logger.exception("Failed to write raw llm output: %s", e)
+        return ""
+
 def extract_json_safely(text: str) -> Optional[Any]:
     """
-    Robustly extract JSON from an LLM output string.
-    Returns Python object or None.
+    Robustly extract JSON from LLM output.
+    1) Try direct loads
+    2) Try marker extraction <<<JSON>>> ... <<<END_JSON>>>
+    3) Try largest balanced { ... } block
+    4) Try array [...]
+    On failure, write raw text to /tmp and return None.
     """
     if not text:
         return None
 
-    # 1) Direct parse
+    # direct
     parsed = _try_json_loads(text)
     if parsed is not None:
         return parsed
 
-    # 2) Marker-based extraction
-    marker_pattern = re.compile(r"<<<JSON>>>\s*([\s\S]*?)\s*<<<END_JSON>>>", re.IGNORECASE)
-    m = marker_pattern.search(text)
+    # marker-based extraction
+    marker_re = re.compile(r"<<<JSON>>>\s*([\s\S]*?)\s*<<<END_JSON>>>", re.IGNORECASE)
+    m = marker_re.search(text)
     if m:
-        candidate = _remove_wrapping_code_blocks(m.group(1))
-        parsed = _try_json_loads(candidate)
-        if parsed is not None:
-            return parsed
-        repaired = _attempt_simple_repairs(candidate)
+        cand = _remove_code_fences(m.group(1))
+        p = _try_json_loads(cand)
+        if p is not None:
+            return p
+        cand2 = _attempt_simple_repairs(cand)
         try:
-            return json.loads(repaired)
+            return json.loads(cand2)
         except Exception as e:
-            logger.warning("Marker JSON parse/repair failed: %s", e)
+            logger.warning("Marker JSON repair failed: %s", e)
 
-    # 3) Fallback: try to extract largest balanced {...} span
-    brace_spans = []
+    # fallback: largest balanced {...}
     stack = []
+    spans = []
     for i, ch in enumerate(text):
         if ch == "{":
             stack.append(i)
         elif ch == "}" and stack:
             start = stack.pop()
-            brace_spans.append((start, i))
-    # sort by length desc
-    brace_spans = sorted(brace_spans, key=lambda t: t[1] - t[0], reverse=True)
-    seen = set()
-    for start, end in brace_spans:
-        if (start, end) in seen:
+            spans.append((start, i))
+    spans = sorted(spans, key=lambda t: t[1]-t[0], reverse=True)
+    for start, end in spans:
+        cand = text[start:end+1].strip()
+        if len(cand) < 20:
             continue
-        seen.add((start, end))
-        candidate = text[start:end+1].strip()
-        if len(candidate) < 20:
-            continue
-        candidate = _remove_wrapping_code_blocks(candidate)
-        parsed = _try_json_loads(candidate)
-        if parsed is not None:
-            return parsed
-        repaired = _attempt_simple_repairs(candidate)
+        cand = _remove_code_fences(cand)
+        p = _try_json_loads(cand)
+        if p is not None:
+            return p
+        cand2 = _attempt_simple_repairs(cand)
         try:
-            return json.loads(repaired)
+            return json.loads(cand2)
         except Exception:
             continue
 
-    # 4) Try arrays as last resort
-    arr_match = re.search(r"(\[[\s\S]*\])", text)
-    if arr_match:
-        candidate = _remove_wrapping_code_blocks(arr_match.group(1))
-        parsed = _try_json_loads(candidate)
-        if parsed is not None:
-            return parsed
-        repaired = _attempt_simple_repairs(candidate)
+    # try array
+    arr = re.search(r"(\[[\s\S]*\])", text)
+    if arr:
+        cand = _remove_code_fences(arr.group(1))
+        p = _try_json_loads(cand)
+        if p is not None:
+            return p
+        cand2 = _attempt_simple_repairs(cand)
         try:
-            return json.loads(repaired)
-        except:
+            return json.loads(cand2)
+        except Exception:
             pass
 
-    # 5) Give up
-    excerpt = text[:2000] if len(text) > 2000 else text
-    logger.error("Failed to extract JSON. Excerpt:\n%s", excerpt)
+    # failed: dump raw text and log traceback
+    dump_path = _dump_raw_llm_output(text)
+    logger.error("extract_json_safely: failed to parse JSON. Raw output saved to: %s", dump_path)
     return None
 
-def _normalize_local_path_args(args: Any) -> Any:
-    if not isinstance(args, dict):
-        return args
-    fp = args.get("file_path") or args.get("path")
-    if isinstance(fp, str) and fp.startswith("/mnt/data/") and os.path.exists(fp):
-        args["file_url"] = f"file://{fp}"
-    return args
-
-# --- Model Loading (safer) ---
+# ---------------- Model helpers (defensive) ----------------
 def init_local_model():
     global LLM_PIPELINE, TOKENIZER
     if LLM_PIPELINE is not None:
         return
-
     try:
         from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
         import torch
-
-        logger.info("Loading local model: %s", LOCAL_MODEL)
         TOKENIZER = AutoTokenizer.from_pretrained(LOCAL_MODEL)
-        # prefer dtype over deprecated torch_dtype
-        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+        dtype = None
+        # choose dtype depending on CUDA availability
+        if torch.cuda.is_available():
+            dtype = torch.float16
         model = AutoModelForCausalLM.from_pretrained(LOCAL_MODEL, device_map="auto", torch_dtype=dtype)
         LLM_PIPELINE = pipeline("text-generation", model=model, tokenizer=TOKENIZER)
-        logger.info("Local model loaded.")
+        logger.info("Local model initialized.")
     except Exception as e:
-        logger.exception("Model load error: %s", e)
+        logger.exception("Failed to load local model: %s", e)
         LLM_PIPELINE = None
 
 def local_llm_generate(prompt: str, max_tokens: int = 512) -> Dict[str, Any]:
-    """
-    Generate text using local pipeline. Returns dict: { "text": <str>, "raw": <pipeline output> }
-    """
     if LLM_PIPELINE is None:
         init_local_model()
     if LLM_PIPELINE is None:
         return {"text": "Model not loaded.", "raw": None}
-
     try:
         out = LLM_PIPELINE(prompt, max_new_tokens=max_tokens, return_full_text=False, do_sample=False)
-        # pipeline output shapes vary; try to extract generated text robustly
-        if isinstance(out, list) and len(out) > 0:
+        # defensively extract text
+        text = ""
+        if isinstance(out, list) and out:
             first = out[0]
             if isinstance(first, dict) and "generated_text" in first:
                 text = first["generated_text"]
             elif isinstance(first, str):
                 text = first
             else:
-                # fallback: join values
                 text = str(first)
         elif isinstance(out, str):
             text = out
-        else:
-            text = ""
         return {"text": text, "raw": out}
     except Exception as e:
-        logger.exception("LLM generation failed: %s", e)
-        return {"text": f"Error: {e}", "raw": None}
+        logger.exception("LLM generation error: %s", e)
+        return {"text": f"LLM error: {e}", "raw": None}
 
-# --- Tools (Zoho) ---
+# ---------------- Zoho token utility ----------------
 def _get_valid_token_headers() -> dict:
     try:
         r = requests.post("https://accounts.zoho.in/oauth/v2/token", params={
             "refresh_token": REFRESH_TOKEN, "client_id": CLIENT_ID,
             "client_secret": CLIENT_SECRET, "grant_type": "refresh_token"
-        }, timeout=10)
+        }, timeout=15)
         if r.status_code == 200:
-            return {"Authorization": f"Zoho-oauthtoken {r.json().get('access_token')}"}
-        logger.error("Zoho token error: %s", r.text)
+            tok = r.json().get("access_token")
+            return {"Authorization": f"Zoho-oauthtoken {tok}"}
+        else:
+            logger.error("Token refresh failed: %s", r.text)
+            return {}
     except Exception as e:
-        logger.exception("Token request failed: %s", e)
-    return {}
+        logger.exception("Token refresh exception: %s", e)
+        return {}
 
+# ---------------- MCP tool implementations ----------------
 @mcp.tool()
 def create_record(module_name: str, record_data: dict) -> str:
-    h = _get_valid_token_headers()
-    if not h:
-        return json.dumps({"status": "error", "message": "Auth Failed"})
+    headers = _get_valid_token_headers()
+    if not headers:
+        return json.dumps({"status": "error", "message": "Auth failed"})
     try:
-        r = requests.post(f"{API_BASE}/{module_name}", headers=h, json={"data": [record_data]}, timeout=15)
-        return json.dumps(r.json()) if r.status_code in (200, 201) else json.dumps({"status": "error", "http_status": r.status_code, "text": r.text})
+        r = requests.post(f"{API_BASE}/{module_name}", headers=headers, json={"data": [record_data]}, timeout=15)
+        return json.dumps(r.json()) if r.status_code in (200,201) else json.dumps({"status":"error","http_status":r.status_code,"text":r.text})
     except Exception as e:
         logger.exception("create_record failed: %s", e)
-        return json.dumps({"status": "error", "message": str(e)})
+        return json.dumps({"status":"error","message": str(e)})
 
 @mcp.tool()
 def create_invoice(data: dict) -> str:
-    h = _get_valid_token_headers()
-    if not h:
-        return json.dumps({"status": "error", "message": "Auth Failed"})
+    headers = _get_valid_token_headers()
+    if not headers:
+        return json.dumps({"status": "error", "message": "Auth failed"})
     try:
-        r = requests.post(f"{INVOICE_API_BASE}/invoices", headers=h,
-                          params={"organization_id": ORGANIZATION_ID}, json=data, timeout=15)
-        return json.dumps(r.json()) if r.status_code in (200, 201) else json.dumps({"status": "error", "http_status": r.status_code, "text": r.text})
+        r = requests.post(f"{INVOICE_API_BASE}/invoices", headers=headers, params={"organization_id": ORGANIZATION_ID}, json=data, timeout=15)
+        return json.dumps(r.json()) if r.status_code in (200,201) else json.dumps({"status":"error","http_status": r.status_code, "text": r.text})
     except Exception as e:
         logger.exception("create_invoice failed: %s", e)
-        return json.dumps({"status": "error", "message": str(e)})
+        return json.dumps({"status":"error","message": str(e)})
 
+# ---------------- Document processing ----------------
 @mcp.tool()
 def process_document(file_path: str, target_module: Optional[str] = "Contacts") -> dict:
+    """Full flow: OCR -> LLM extraction -> KPI -> result with raw llm text for debugging"""
     if not os.path.exists(file_path):
-        return {"status": "error", "error": f"File not found at path: {file_path}"}
+        return {"status": "error", "error": f"File not found: {file_path}"}
 
-    # 1) OCR (returns text + confidence)
     raw_text, ocr_score = extract_text_and_conf(file_path)
     if not raw_text:
-        return {"status": "error", "error": "OCR empty or failed."}
+        return {"status": "error", "error": "OCR returned empty text."}
 
-    # 2) LLM extraction
     prompt = get_ocr_extraction_prompt(raw_text, page_count=1)
-    res = local_llm_generate(prompt, max_tokens=512)
-    llm_text = res.get("text", "")
-    data = extract_json_safely(llm_text)
+    llm_res = local_llm_generate(prompt, max_tokens=512)
+    llm_text = llm_res.get("text", "")
 
-    # 3) KPI calculation
-    kpis = {"score": 0, "rating": "Fail", "issues": ["Extraction Failed"]}
-    if data:
-        kpis = calculate_extraction_confidence(data, ocr_score)
+    parsed = extract_json_safely(llm_text)
+    kpis = {"score": 0, "rating": "Fail", "issues": ["Extraction failed"]}
+    if parsed:
+        # compute kpis basic heuristics (simple)
+        try:
+            total = parsed.get("totals", {}).get("grand_total")
+            semantic_ok = 1 if total else 0
+            kpis = {
+                "score": 80 if semantic_ok else 40,
+                "rating": "High" if semantic_ok else "Low",
+                "ocr_score": ocr_score,
+                "issues": [] if semantic_ok else ["grand_total missing"]
+            }
+        except Exception:
+            kpis["issues"].append("Error computing KPIs")
+
+    # If parse failed, persist raw LLM output path for debugging
+    raw_dump = None
+    if not parsed:
+        raw_dump = _dump_raw_llm_output(llm_text)
 
     return {
-        "status": "success",
+        "status": "success" if parsed else "partial",
         "file": os.path.basename(file_path),
-        "extracted_data": data if data else {"raw": llm_text},
+        "extracted_data": parsed if parsed else None,
         "raw_llm_output": llm_text,
+        "raw_llm_dump_path": raw_dump,
         "kpis": kpis
     }
 
-# --- Executor for tool-calls ---
+# ---------------- Agent orchestration and chat ----------------
 def parse_and_execute(model_text: str, history: list) -> str:
     payload = extract_json_safely(model_text)
     if not payload:
-        return "No valid tool call found in model output."
+        return "No valid tool JSON found in model output. Raw output saved for debugging."
+
+    if isinstance(payload, dict):
+        cmds = [payload]
+    else:
+        cmds = payload
 
-    cmds = [payload] if isinstance(payload, dict) else payload
     results = []
     last_contact_id = None
 
@@ -338,21 +292,18 @@ def parse_and_execute(model_text: str, history: list) -> str:
         if not isinstance(cmd, dict):
             continue
         tool = cmd.get("tool")
-        args = _normalize_local_path_args(cmd.get("args", {}))
+        args = cmd.get("args", {})
 
         if tool == "create_record":
-            module_name = args.get("module_name", "Contacts")
-            record_data = args.get("record_data", {})
-            res = create_record(module_name, record_data)
+            module = args.get("module_name", "Contacts")
+            record = args.get("record_data", {})
+            res = create_record(module, record)
             results.append(f"create_record -> {res}")
+            # attempt to capture id
             try:
                 rj = json.loads(res)
-                # try to locate id in common locations
-                if isinstance(rj, dict):
-                    if "data" in rj and isinstance(rj["data"], list) and rj["data"] and "details" in rj["data"][0]:
-                        last_contact_id = rj["data"][0]["details"].get("id")
-                    elif "id" in rj:
-                        last_contact_id = rj.get("id")
+                if isinstance(rj, dict) and "data" in rj and isinstance(rj["data"], list) and rj["data"]:
+                    last_contact_id = rj["data"][0].get("details", {}).get("id")
             except Exception:
                 pass
 
@@ -366,47 +317,51 @@ def parse_and_execute(model_text: str, history: list) -> str:
         else:
             results.append(f"Unknown tool: {tool}")
 
-    return "\n".join(results) if results else "No actionable tool commands executed."
+    return "\n".join(results) if results else "No actionable tool calls executed."
 
-# --- Chat Core ---
 def chat_logic(message: str, file_path: Optional[str], history: list) -> str:
-    # PHASE: File Upload -> Extraction -> KPI Report
     if file_path:
-        logger.info("Processing uploaded file: %s", file_path)
+        logger.info("chat_logic: processing file %s", file_path)
         doc = process_document(file_path)
-        if doc.get("status") == "success":
-            data = doc.get("extracted_data")
-            kpi = doc.get("kpis", {})
-            extracted_json = json.dumps(data, indent=2) if not isinstance(data, str) else data
-            rating_emoji = "🟢" if kpi.get('rating') == 'High' else ("🟡" if kpi.get('rating') == 'Medium' else "🔴")
-            issues_txt = "\n".join([f"- {i}" for i in kpi.get('issues', [])]) if kpi.get('issues') else "None"
-
-            return (
-                f"### 📄 Extraction Complete: **{doc.get('file')}**\n"
-                f"**Combined Confidence:** {rating_emoji} {kpi.get('score')}/100\n"
-                f"*(OCR Signal: {kpi.get('ocr_score')}% | Data Quality: {kpi.get('semantic_score')}%)*\n\n"
-                f"**Issues Detected:**\n{issues_txt}\n\n"
-                f"```json\n{extracted_json}\n```\n\n"
-                "If you want to persist this to Zoho, type **Create Invoice** or ask me to create the contact/item first."
+        status = doc.get("status")
+        if status in ("success", "partial"):
+            extracted = doc.get("extracted_data")
+            raw_llm = doc.get("raw_llm_output")
+            dump_path = doc.get("raw_llm_dump_path")
+            kpis = doc.get("kpis", {})
+            extracted_pretty = json.dumps(extracted, indent=2) if extracted else "(no structured JSON parsed)"
+            msg = (
+                f"### 📄 Extraction Result for **{doc.get('file')}**\n"
+                f"Status: {status}\n"
+                f"KPI Score: {kpis.get('score')} Rating: {kpis.get('rating')}\n"
+                f"OCR Confidence: {kpis.get('ocr_score', 'N/A')}\n\n"
+                f"Extracted JSON:\n```json\n{extracted_pretty}\n```\n"
             )
+            if dump_path:
+                msg += f"\n⚠️ The model output could not be parsed into strict JSON. Raw LLM output saved to: `{dump_path}`\n"
+                msg += "You can inspect that file to debug the model response or prompt."
+            msg += "\nType 'Create Invoice' to persist when ready."
+            return msg
         else:
-            return f"OCR/Extraction error: {doc.get('error')}"
+            return f"Error during processing: {doc.get('error')}"
 
-    # PHASE: Text Interaction -> Agent orchestrator
+    # text-only interaction
     hist_txt = "\n".join([f"U: {h[0]}\nA: {h[1]}" for h in history]) if history else ""
     prompt = get_agent_prompt(hist_txt, message)
     gen = local_llm_generate(prompt, max_tokens=256)
     gen_text = gen.get("text", "")
 
-    # If LLM returned a tool JSON, execute it
     tool_payload = extract_json_safely(gen_text)
     if tool_payload:
         return parse_and_execute(gen_text, history)
 
-    # Otherwise return the LLM text (human facing)
-    return gen_text or "No response from model."
+    # if not a tool call, return the LLM text (or clear error)
+    if gen_text:
+        return gen_text
+    else:
+        return "No response from model."
 
-# --- UI wrapper ---
+# ---------------- Gradio wrapper ----------------
 def chat_handler(msg, hist):
     txt = msg.get("text", "")
     files = msg.get("files", [])