Update app.py

app.py

@@ -36,6 +36,16 @@ MODELS = {
             "HTTP-Referer": "https://huggingface.co",
             "X-Title": "Invoice Extractor"
         }
+    },
+    "Mistral Small": {
+        "api_url": "https://openrouter.ai/api/v1/chat/completions",
+        "model_name": "mistralai/mistral-small-3.1-24b-instruct:free",
+        "api_key_env": "OPENROUTER_API_KEY",
+        "response_format": {"type": "json_object"},
+        "extra_headers": {
+            "HTTP-Referer": "https://huggingface.co",
+            "X-Title": "Invoice Extractor"
+        }
     }
 }
 
@@ -142,6 +152,37 @@ def clean_json_response(text):
 
 def get_extraction_prompt(model_choice, text):
     """Return the appropriate prompt based on model choice"""
+    base_prompt = """Extract complete invoice information and return a VALID JSON object with these fields:
+{
+  "invoice_header": {
+    "invoice_number": "string",
+    "invoice_date": "YYYY-MM-DD",
+    "po_number": "string or null",
+    "invoice_value": "string with currency",
+    "supplier_name": "string or null",
+    "customer_name": "string or null"
+  },
+  "line_items": [
+    {
+      "item_number": "string or null",
+      "description": "string",
+      "quantity": "number",
+      "unit_price": "string with currency",
+      "total_price": "string with currency"
+    }
+  ]
+}
+Rules:
+1. Return ONLY valid JSON (no additional text or markdown)
+2. Use null for missing fields
+3. Date format must be YYYY-MM-DD
+4. All currency values must include currency symbol or code
+5. Include all line items found in the invoice
+6. For line items, quantity should be a number, prices as strings with currency
+7. Do not include any explanations or notes
+Invoice Text:
+""" + text
+
     if model_choice == "DeepSeek v3":
         return f"""Extract complete invoice information from the text below and return ONLY a valid JSON object with these fields:
 {{
@@ -186,37 +227,8 @@ Invoice Text:
 Invoice Text:
 """ + text
     
-    else:  # Llama 4 Mavericks
-        return f"""Extract complete invoice information and return a VALID JSON object with these fields:
-{{
-  "invoice_header": {{
-    "invoice_number": "string",
-    "invoice_date": "YYYY-MM-DD",
-    "po_number": "string or null",
-    "invoice_value": "string with currency",
-    "supplier_name": "string or null",
-    "customer_name": "string or null"
-  }},
-  "line_items": [
-    {{
-      "item_number": "string or null",
-      "description": "string",
-      "quantity": "number",
-      "unit_price": "string with currency",
-      "total_price": "string with currency"
-    }}
-  ]
-}}
-Rules:
-1. Return ONLY valid JSON (no additional text or markdown)
-2. Use null for missing fields
-3. Date format must be YYYY-MM-DD
-4. All currency values must include currency symbol or code
-5. Include all line items found in the invoice
-6. For line items, quantity should be a number, prices as strings with currency
-7. Do not include any explanations or notes
-Invoice Text:
-""" + text
+    else:  # For Llama 4 and Mistral
+        return base_prompt
 
 def format_currency(value):
     """Helper function to format currency values consistently"""
@@ -234,8 +246,8 @@ def display_line_items(line_items, model_choice="DeepSeek v3"):
 
     st.subheader("📋 Line Items")
     
-    if model_choice == "Llama 4 Mavericks":
-        # Display as a table for Llama
+    if model_choice in ["Llama 4 Mavericks", "Mistral Small"]:
+        # Display as a table for Llama/Mistral
         items_display = []
         for idx, item in enumerate(line_items, 1):
             items_display.append({
@@ -267,7 +279,7 @@ def display_invoice_data(model_choice, invoice_data):
     if not invoice_data:
         return
     
-    if model_choice == "Llama 4 Mavericks":
+    if model_choice in ["Llama 4 Mavericks", "Mistral Small"]:
         # Display header information
         st.subheader("Invoice Summary")
         header = invoice_data.get("invoice_header", {})
@@ -326,7 +338,7 @@ def extract_invoice_info(model_choice, text):
         return None
     
     # Normalize data structure based on model
-    if model_choice == "Llama 4 Mavericks":
+    if model_choice in ["Llama 4 Mavericks", "Mistral Small"]:
         if "invoice_header" not in parsed_data:
             parsed_data["invoice_header"] = {}
         if "line_items" not in parsed_data: