Update app.py

app.py

@@ -22,7 +22,81 @@ try:
     RESAMPLE = Image.Resampling.LANCZOS  # Pillow >= 10
 except AttributeError:
     RESAMPLE = Image.LANCZOS             # Pillow < 10
+PROMPT_FREIGHT_JSON = """
+Please analyze the freight rate table in the file I provide and convert it into JSON in the following structure:
+{
+  "shipping_line": "...",
+  "shipping_line_code": "...",
+  "shipping_line_reason": "Why this carrier is chosen?",
+  "fee_type": "Air Freight",
+  "valid_from": ...,
+  "valid_to": ...,
+  "charges": [
+    {
+      "frequency": "...",
+      "package_type": "...",
+      "aircraft_type": "...",
+      "direction": "Export or Import or null",
+      "origin": "...",
+      "destination": "...",
+      "charge_name": "...",
+      "charge_code": "...",
+      "charge_code_reason": "...",
+      "cargo_type": "...",
+      "currency": "...",
+      "transit": "...",
+      "transit_time": "...",
+      "weight_breaks": {
+        "min": ...,
+        "-45kg": ...,
+        "45kg": ...,
+        "100kg": ...,
+        "300kg": ...,
+        "500kg": ...,
+        "1000kg": ...,
+        "other": {
+          key: value
+        }
+      },
+      "remark": "..."
+    }
+  ],
+  "local_charges": [
+    {
+      "charge_name": "...",
+      "charge_code": "...",
+      "unit": "...",
+      "amount": ...,
+      "remark": "..."
+    }
+  ]
+}
 
+### Date rules
+- valid_from format:
+  - `DD/MM/YYYY` (if full date)
+  - `01/MM/YYYY` (if month+year only)
+  - `01/01/YYYY` (if year only)
+  - `UFN` if missing
+- valid_to:
+  - exact `DD/MM/YYYY` if present
+  - else `UFN`
+
+STRICT RULES:
+- ONLY return a single JSON object as specified above.
+- All rates must exactly match the corresponding weight break columns (45kg, 100kg, 300kg, 500kg, 1000kg, etc.). set null if N/A. No assumptions or interpolations.
+- If the table shows "RQ" or similar, set value as "RQST".
+- Group same-price destinations into one record separated by "/".
+- Always use IATA code for origin and destination.
+- Flight number (e.g. ZH118) is not charge code.
+- Frequency: D[1-7]; 'Daily' = D1234567. Join multiple (e.g. D3,D4→D34).
+- If local charges exist, list them.
+- If validity missing, set null.
+- Direction: Export if origin is Vietnam (SGN, HAN, DAD...), else Import.
+- Provide short plain English reasons for "shipping_line_reason" & "charge_code_reason".
+- Replace commas in remarks with semicolons.
+- Only return JSON.
+"""
 # ================== HELPERS ==================
 import fitz  # PyMuPDF
 
@@ -248,6 +322,41 @@ def preview_process(file):
         print(f"Preview error: {e}")
         return []
 
+def _merge_freight_objects(objs: list[dict]) -> dict | None:
+    if not objs: return None
+    base = {}
+    for k in ["shipping_line","shipping_line_code","shipping_line_reason","fee_type","valid_from","valid_to"]:
+        for o in objs:
+            if isinstance(o, dict) and o.get(k):
+                base[k] = o[k]
+                break
+        base.setdefault(k, None)
+
+    seen = set()
+    merged_charges, merged_local = [], []
+    def norm(v): return v.replace(",", ";") if isinstance(v, str) else v
+
+    for o in objs:
+        for c in (o.get("charges") or []):
+            wb = json.dumps(c.get("weight_breaks", {}), sort_keys=True, ensure_ascii=False)
+            key = (c.get("origin"), c.get("destination"), c.get("charge_name"), c.get("charge_code"), c.get("currency"), wb)
+            if key in seen: continue
+            c["remark"] = norm(c.get("remark"))
+            merged_charges.append(c)
+            seen.add(key)
+        for lc in (o.get("local_charges") or []):
+            lc["remark"] = norm(lc.get("remark"))
+            merged_local.append(lc)
+
+    base["charges"] = merged_charges
+    base["local_charges"] = merged_local
+    return base
+def _coerce_only_json(text: str) -> str:
+    obj, s = _extract_json_from_message(text)
+    if obj is not None:
+        return json.dumps(obj, ensure_ascii=False)
+    m = re.search(r"\{.*\}\s*$", text, flags=re.DOTALL)
+    return m.group(0) if m else text.strip()
 # -------- Internal (Gemini) - Base (1 lượt, không thinking) --------
 def run_process_internal_base(file_bytes, filename, mime, question, model_choice,
                               temperature, top_p):
@@ -302,6 +411,57 @@ def run_process_internal_base(file_bytes, filename, mime, question, model_choice
                 genai.delete_file(uploaded.name)
         except Exception:
             pass
+# ================== MAIN OCR FUNCTION ==================
+def run_process_internal_base_v2(file_bytes, filename, mime, question, model_choice, temperature, top_p, batch_size=3):
+    api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
+    if not api_key:
+        return "ERROR: Missing GOOGLE_API_KEY.", None
+    genai.configure(api_key=api_key)
+
+    model_name = INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash")
+    model = genai.GenerativeModel(model_name=model_name, generation_config={"temperature": float(temperature), "top_p": float(top_p)})
+
+    if file_bytes[:4] == b"%PDF":
+        pages = pdf_to_images(file_bytes)
+    else:
+        pages = [Image.open(io.BytesIO(file_bytes))]
+
+    user_prompt = (question or "").strip() or PROMPT_FREIGHT_JSON
+    all_json_results, all_text_results = [], []
+
+    for i in range(0, len(pages), batch_size):
+        batch = pages[i:i+batch_size]
+        uploaded = []
+        for im in batch:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+                im.save(tmp.name)
+                up = genai.upload_file(path=tmp.name, mime_type="image/png")
+                up = _wait_file_active(up)
+                uploaded.append(up)
+        try:
+            prompt = f"{user_prompt}\n(This is batch {i//batch_size+1})"
+            resp = model.generate_content([prompt] + uploaded)
+            text = _safe_text_from_gemini(resp)
+            json_text = _coerce_only_json(text)
+            try:
+                parsed = json.loads(json_text)
+                all_json_results.append(parsed)
+            except Exception:
+                all_text_results.append(text)
+        finally:
+            for up in uploaded:
+                try: genai.delete_file(up.name)
+                except: pass
+
+    if all_json_results:
+        merged_json = _merge_freight_objects(all_json_results)
+        message = json.dumps(merged_json, ensure_ascii=False, indent=2)
+        return message, merged_json
+
+    combined_text = "\n\n".join(all_text_results)
+    message = _pretty_message(combined_text)
+    parsed_obj, _ = _extract_json_from_message(combined_text)
+    return message, parsed_obj
 
 # -------- External API --------
 def run_process_external(file_bytes, filename, mime, question, api_url,
@@ -366,7 +526,7 @@ def run_process(file, question, model_choice, temperature, top_p, external_api_u
                 temperature=temperature, top_p=top_p
             )
 
-        return run_process_internal_base(
+        return run_process_internal_base_v2(
             file_bytes=file_bytes, filename=filename, mime=mime,
             question=question, model_choice=model_choice,
             temperature=temperature, top_p=top_p