Update app.py

app.py

@@ -68,7 +68,10 @@ def query_llm(model_choice, prompt):
         with st.spinner(f"🔍 Querying {model_choice}..."):
             r = requests.post(cfg["api_url"], headers=headers, json=payload, timeout=90)
         if r.status_code != 200:
-            st.error(f"🚨 API Error {r.status_code}: {r.text}")
+            if "No instances available" in r.text or r.status_code == 503:
+                st.error(f"{model_choice} is currently unavailable. Please try again later or select another model.")
+            else:
+                st.error(f"🚨 API Error {r.status_code}: {r.text}")
             return None
         content = r.json()["choices"][0]["message"]["content"]
         st.session_state.last_api = content
@@ -113,26 +116,14 @@ def fallback_supplier(text):
     return None
 
 def get_extraction_prompt(model_choice, txt):
-    if model_choice.startswith("DeepSeek"):
-        return (
-            "Extract full invoice info and RETURN ONLY a single-line json object with fields:\n"
-            '{"invoice_number":"string","invoice_date":"YYYY-MM-DD",'
-            '"po_number":"string|null","invoice_value":"string with currency",'
-            '"line_items":[{"description":"string","quantity":"number","unit_price":"string with currency","total_price":"string with currency"}]}\n'
-            "Use null for missing. NO extra text.\n\n"
-            f"Invoice Text:\n{txt}"
-        )
-    else:
-        return (
-            "Extract invoice data and RETURN ONLY a compact, one-line json object exactly:\n"
-            '{"invoice_header":{"invoice_number":"string","invoice_date":"YYYY-MM-DD",'
-            '"po_number":"string|null","invoice_value":"string with currency",'
-            '"supplier_name":"string|null","customer_name":"string|null"},'
-            '"line_items":[{"item_number":"string|null","description":"string","quantity":number,'
-            '"unit_price":"string with currency","total_price":"string with currency"}]}\n'
-            "Use null for missing. NO extras.\n\n"
-            f"Invoice Text:\n{txt}"
-        )
+    # New, broad prompt for all models:
+    return (
+        "Extract all possible metadata fields from the following invoice, including but not limited to header information, supplier and customer details, payment terms, tax details, references, and every possible line item with all available attributes. "
+        "Return a detailed JSON object containing every field you can identify, and make sure to include all line items as an array. "
+        "If any field is missing in the invoice, use null. Do not add any explanation or extra text outside the JSON. "
+        "\n\nInvoice Text:\n"
+        f"{txt}"
+    )
 
 def extract_invoice_info(model_choice, text):
     prompt = get_extraction_prompt(model_choice, text)
@@ -145,23 +136,32 @@ def extract_invoice_info(model_choice, text):
 
     # DeepSeek models: flat format
     if model_choice.startswith("DeepSeek"):
-        for k in ("invoice_number","invoice_date","po_number","invoice_value"):
-            data.setdefault(k, None)
-        items = data.setdefault("line_items", [])
-        for itm in items:
+        # Dynamically handle flat or semi-structured output (may contain any fields)
+        data.setdefault("line_items", [])
+        for itm in data["line_items"]:
+            if not isinstance(itm, dict):
+                continue
             for k in ("description","quantity","unit_price","total_price"):
                 itm.setdefault(k, None)
         return data
-    # Other models (OpenAI GPT-4.1, Mistral): nested format
-    hdr = data.setdefault("invoice_header", {})
+    # Other models (OpenAI GPT-4.1, Mistral): usually nested under invoice_header, but now prompt is broader, so handle flexibly
+    # Accepts a flexible schema as model may include extra keys
+    hdr = data.get("invoice_header", {})
+    if not hdr and any(k in data for k in ("invoice_number","supplier_name","customer_name")):
+        # If model returned flat, treat top-level keys as header
+        hdr = data
     for k in ("invoice_number","invoice_date","po_number","invoice_value","supplier_name","customer_name"):
         hdr.setdefault(k, None)
     if not hdr.get("supplier_name"):
         hdr["supplier_name"] = fallback_supplier(text)
-    items = data.setdefault("line_items", [])
+    items = data.get("line_items", [])
     for itm in items:
+        if not isinstance(itm, dict):
+            continue
         for k in ("item_number","description","quantity","unit_price","total_price"):
             itm.setdefault(k, None)
+    data["invoice_header"] = hdr
+    data["line_items"] = items
 
     return data
 
@@ -188,25 +188,31 @@ with tab2:
         info = extract_invoice_info(mdl, txt)
         if info:
             st.success("Extraction Complete")
+            # For DeepSeek, output may be flat; for others, prefer "invoice_header" nesting
             if mdl.startswith("DeepSeek"):
+                # Show all keys except line_items
+                non_items = {k: v for k, v in info.items() if k != "line_items"}
                 c1, c2 = st.columns(2)
-                c1.metric("Invoice #", info["invoice_number"])
-                c1.metric("PO #", info["po_number"])
-                c2.metric("Date", info["invoice_date"])
-                c2.metric("Value", info["invoice_value"])
+                for i, (k, v) in enumerate(non_items.items()):
+                    (c1 if i % 2 == 0 else c2).metric(k.replace("_", " ").title(), v)
                 st.subheader("Line Items")
-                st.table(info["line_items"])
+                st.table(info.get("line_items", []))
             else:
-                h = info["invoice_header"]
+                h = info.get("invoice_header", {})
                 c1, c2, c3 = st.columns(3)
-                c1.metric("Invoice #", h["invoice_number"])
-                c1.metric("Supplier", h["supplier_name"])
-                c2.metric("Date", h["invoice_date"])
-                c2.metric("Customer", h["customer_name"])
-                c3.metric("PO #", h["po_number"])
-                c3.metric("Total", h["invoice_value"])
+                c1.metric("Invoice #", h.get("invoice_number"))
+                c1.metric("Supplier", h.get("supplier_name"))
+                c2.metric("Date", h.get("invoice_date"))
+                c2.metric("Customer", h.get("customer_name"))
+                c3.metric("PO #", h.get("po_number"))
+                c3.metric("Total", h.get("invoice_value"))
+                # Show any additional header fields detected
+                extra_fields = {k: v for k, v in h.items() if k not in ("invoice_number", "supplier_name", "customer_name", "invoice_date", "po_number", "invoice_value")}
+                if extra_fields:
+                    st.subheader("Additional Header Metadata")
+                    st.json(extra_fields)
                 st.subheader("Line Items")
-                st.table(info["line_items"])
+                st.table(info.get("line_items", []))
 
     if "last_api" in st.session_state:
         with st.expander("Debug"):