Update app.py

app.py

@@ -1,36 +1,39 @@
 from __future__ import annotations
-import os, io, re, json, mimetypes, tempfile
-from typing import List, Union, Tuple
+import os, io, re, json, time, mimetypes, tempfile
+from typing import List, Union, Tuple, Any
 from PIL import Image
 import pandas as pd
 import gradio as gr
 import google.generativeai as genai
 import requests
+import pdfplumber
+import fitz  # PyMuPDF
 
 # ================== CONFIG ==================
 DEFAULT_API_KEY = "AIzaSyBbK-1P3JD6HPyE3QLhkOps6_-Xo3wUFbs"
+
 INTERNAL_MODEL_MAP = {
     "Gemini 2.5 Flash": "gemini-2.5-flash",
-    "Gemini 2.5 Pro": "gemini-2.5-pro",
+    "Gemini 2.5 Pro":   "gemini-2.5-pro",
 }
 EXTERNAL_MODEL_NAME = "prithivMLmods/Camel-Doc-OCR-062825 (External)"
 
 PROMPT_FREIGHT_JSON = """
-Please analyze the freight rate table in the CSV file I provide and convert it into JSON with this structure:
+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": "...",
+  "valid_from": ...,
+  "valid_to": ...,
   "charges": [
     {
       "frequency": "...",
       "package_type": "...",
       "aircraft_type": "...",
       "direction": "Export or Import or null",
-      "origin": "...",       # detect automatically from header, filename, or text (e.g. SGN/HAN/DAD)
+      "origin": "...",
       "destination": "...",
       "charge_name": "...",
       "charge_code": "...",
@@ -47,8 +50,10 @@ Please analyze the freight rate table in the CSV file I provide and convert it i
         "+300kg": ...,
         "+500kg": ...,
         "+1000kg": ...,
-        "other": { key: value },
-        "weight_breaks_reason": "Why chosen weight_breaks?"
+        "other": {
+          key: value
+        },
+        "weight_breaks_reason":"Why chosen weight_breaks?"
       },
       "remark": "..."
     }
@@ -63,148 +68,159 @@ Please analyze the freight rate table in the CSV file I provide and convert it i
     }
   ]
 }
-
-Rules:
-- If filename or top text includes "Origin: SGN", "SGN", "HAN", or "DAD" → use as origin.
-- If missing, infer origin from file name (e.g., "TK - SGN Rate Sheet.csv" → SGN).
-- All rates must match the weight break columns (M, N, 45, 100, 300, 500, 1000, etc.).
-- No assumptions; set null if missing.
-- Only return valid JSON object as above.
+### 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 (M,N,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 ==================
-def _read_file_bytes(upload: Union[str, os.PathLike, dict, object] | None) -> bytes:
-    if upload is None:
-        raise ValueError("No file uploaded.")
-    if isinstance(upload, (str, os.PathLike)):
+def _read_file_bytes(upload):
+    if isinstance(upload, str):
         with open(upload, "rb") as f:
             return f.read()
-    if isinstance(upload, dict) and "path" in upload:
-        with open(upload["path"], "rb") as f:
-            return f.read()
-    if hasattr(upload, "read"):
+    elif hasattr(upload, "read"):
         return upload.read()
-    raise TypeError(f"Unsupported file object: {type(upload)}")
+    raise TypeError("Unsupported file input")
 
 def _guess_name_and_mime(file, file_bytes: bytes) -> Tuple[str, str]:
-    if isinstance(file, (str, os.PathLike)):
-        filename = os.path.basename(str(file))
-    elif isinstance(file, dict) and "name" in file:
-        filename = os.path.basename(file["name"])
-    elif isinstance(file, dict) and "path" in file:
-        filename = os.path.basename(file["path"])
-    else:
-        filename = "upload.bin"
+    filename = os.path.basename(file.name if hasattr(file, "name") else str(file))
     mime, _ = mimetypes.guess_type(filename)
-    if not mime:
-        if filename.lower().endswith(".csv"):
-            mime = "text/csv"
-        elif len(file_bytes) >= 4 and file_bytes[:4] == b"%PDF":
-            mime = "application/pdf"
-        else:
-            mime = "image/png"
-    return filename, mime
-
-
-# ================== GEMINI PROCESS ==================
-def run_gemini_text(file_bytes, filename, mime, model_choice, question, temperature, top_p):
-    """Gemini đọc CSV/text → sinh JSON"""
-    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)})
-
-    # đọc CSV nếu có
-    csv_text = None
-    if mime == "text/csv" or filename.lower().endswith(".csv"):
-        try:
-            df = pd.read_csv(io.BytesIO(file_bytes))
-            csv_text = df.to_csv(index=False)
-        except Exception:
-            csv_text = file_bytes.decode("utf-8", errors="ignore")
-
-    # prompt chính
-    user_prompt = question.strip() if question else PROMPT_FREIGHT_JSON
-    full_prompt = (
-        f"{user_prompt}\n\n"
-        f"Filename: {filename}\n\n"
-        f"Below is the table text extracted from your CSV file:\n{csv_text or file_bytes.decode('utf-8', errors='ignore')}\n\n"
-        "Please analyze and return valid JSON only."
-    )
-
-    resp = model.generate_content(full_prompt)
-    return resp.text.strip(), None
+    if not mime and file_bytes[:4] == b"%PDF":
+        mime = "application/pdf"
+    return filename, mime or "application/octet-stream"
 
+def pdf_to_images(pdf_bytes: bytes) -> list[Image.Image]:
+    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+    return [Image.frombytes("RGB", [p.get_pixmap(dpi=200).width, p.get_pixmap(dpi=200).height], p.get_pixmap(dpi=200).samples) for p in doc]
 
-# ================== EXTERNAL API (nếu có) ==================
-def run_process_external(file_bytes, filename, mime, question, api_url, temperature, top_p):
-    if not api_url:
-        return "ERROR: Missing external API endpoint.", None
-    data = {"prompt": question or "", "temperature": str(temperature), "top_p": str(top_p)}
-    files = {"file": (filename, file_bytes, mime)}
-    r = requests.post(api_url, files=files, data=data, timeout=60)
-    if r.status_code >= 400:
-        return f"ERROR: External API HTTP {r.status_code}: {r.text[:200]}", None
-    return r.text, None
+# ================== PDF CHECK ==================
+def check_pdf_structure(file_bytes: bytes) -> bool:
+    """Trả về True nếu PDF có nhiều trang và dạng bảng."""
+    try:
+        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
+            if len(pdf.pages) <= 2:
+                return False
+            for page in pdf.pages[:3]:
+                tables = page.find_tables()
+                if tables:
+                    return True
+        return False
+    except Exception as e:
+        print("PDF check error:", e)
+        return False
 
+# ================== GEMINI CALL ==================
+def call_gemini_with_prompt(content_text: str, question: str, model_choice: str, temperature: float, top_p: float):
+    api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
+    genai.configure(api_key=api_key)
+    model = genai.GenerativeModel(
+        model_name=INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash"),
+        generation_config={"temperature": temperature, "top_p": top_p}
+    )
+    prompt = f"{PROMPT_FREIGHT_JSON}\n{question or ''}\n\nBelow is the extracted CSV data:\n{content_text}"
+    response = model.generate_content(prompt)
+    return getattr(response, "text", str(response))
 
-# ================== MAIN ROUTER ==================
+# ================== MAIN LOGIC ==================
 def run_process(file, question, model_choice, temperature, top_p, external_api_url):
     try:
         if file is None:
-            return "ERROR: No file uploaded.", None
+            return "❌ No file uploaded.", None
 
         file_bytes = _read_file_bytes(file)
         filename, mime = _guess_name_and_mime(file, file_bytes)
-
-        print(f"[INFO] Processing {filename} ({mime})...")
-
-        # Nếu là CSV → đọc text & gửi Gemini
-        if mime == "text/csv" or filename.lower().endswith(".csv"):
-            print("🟢 Detected CSV file → Sending to Gemini for JSON conversion.")
-            return run_gemini_text(file_bytes, filename, mime, model_choice, question, temperature, top_p)
-
-        # Nếu chọn external
-        if model_choice == EXTERNAL_MODEL_NAME:
-            return run_process_external(
-                file_bytes=file_bytes, filename=filename, mime=mime,
-                question=question, api_url=external_api_url,
-                temperature=temperature, top_p=top_p
-            )
-
-        # fallback: PDF / image
-        return "⚠️ Only CSV supported in this version. Please upload .csv file.", None
+        print(f"[UPLOAD] {filename} ({mime})")
+
+        # 1️⃣ Nếu là PDF và có nhiều trang dạng bảng
+        if mime == "application/pdf" and check_pdf_structure(file_bytes):
+            print("➡️ PDF nhiều trang & dạng bảng → trích xuất CSV trước khi gọi Gemini.")
+            all_dfs, saved_header = [], None
+            with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
+                for idx, page in enumerate(pdf.pages, start=1):
+                    table = page.extract_table({
+                        "vertical_strategy": "lines",
+                        "horizontal_strategy": "text",
+                        "snap_tolerance": 3,
+                        "intersection_tolerance": 5,
+                    })
+                    if not table or len(table) < 2:
+                        continue
+                    header, rows = table[0], table[1:]
+                    if saved_header is None:
+                        saved_header = header
+                    elif len(header) < len(saved_header):
+                        header = saved_header
+                    try:
+                        df = pd.DataFrame(rows, columns=header)
+                        all_dfs.append(df)
+                    except Exception as e:
+                        print(f"⚠️ Trang {idx} lỗi DataFrame: {e}")
+
+            if all_dfs:
+                final_df = pd.concat(all_dfs, ignore_index=True).dropna(how="all")
+                csv_text = final_df.to_csv(index=False)
+                print(f"✅ Trích xuất {len(final_df)} dòng, gửi Gemini xử lý JSON.")
+                message = call_gemini_with_prompt(csv_text, question, model_choice, temperature, top_p)
+                return message, None
+            else:
+                print("⚠️ Không có bảng hợp lệ, fallback qua OCR bình thường.")
+
+        # 2️⃣ Các loại file còn lại → xử lý như cũ
+        api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel(
+            model_name=INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash"),
+            generation_config={"temperature": temperature, "top_p": top_p}
+        )
+        uploaded = genai.upload_file(path=file.name)
+        resp = model.generate_content([question or PROMPT_FREIGHT_JSON, uploaded])
+        genai.delete_file(uploaded.name)
+        return getattr(resp, "text", str(resp)), None
 
     except Exception as e:
-        return f"ERROR: {type(e).__name__}: {str(e)}", None
-
+        return f"ERROR: {type(e).__name__}: {e}", None
 
 # ================== UI ==================
 def main():
-    with gr.Blocks(title="CSV → JSON Converter (Gemini)") as demo:
-        gr.Markdown("## 📦 Upload CSV → Gemini generates structured JSON")
-
-        file = gr.File(label="Upload CSV file")
-        question = gr.Textbox(label="Custom Prompt (optional)", lines=2)
+    with gr.Blocks(title="OCR + Table Extraction for Gemini") as demo:
+        file = gr.File(label="📂 Upload PDF / Image / CSV")
+        question = gr.Textbox(label="Prompt", lines=2)
         model_choice = gr.Dropdown(choices=[*INTERNAL_MODEL_MAP.keys(), EXTERNAL_MODEL_NAME],
                                    value="Gemini 2.5 Flash", label="Model")
         temperature = gr.Slider(0.0, 2.0, value=0.2, step=0.05)
         top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01)
         external_api_url = gr.Textbox(label="External API URL", visible=False)
         output_text = gr.Code(label="Gemini Output", language="json")
-        run_btn = gr.Button("🚀 Convert to JSON")
+        run_btn = gr.Button("🚀 Run Extraction")
 
-        run_btn.click(run_process,
-                      inputs=[file, question, model_choice, temperature, top_p, external_api_url],
-                      outputs=[output_text, gr.State()])
+        run_btn.click(
+            run_process,
+            inputs=[file, question, model_choice, temperature, top_p, external_api_url],
+            outputs=[output_text, gr.State()]
+        )
 
     return demo
 
-
 demo = main()
 if __name__ == "__main__":
     demo.launch()