Create app.py

app.py

@@ -0,0 +1,504 @@
+from __future__ import annotations
+import os, io, re, json, time, mimetypes, tempfile, string
+from typing import List, Union, Tuple, Any, Iterable
+
+from PIL import Image
+import pandas as pd
+import gradio as gr
+import google.generativeai as genai
+import requests
+
+# ================== CONFIG ==================
+# KHÔNG hardcode key. YÊU CẦU đặt biến môi trường GOOGLE_API_KEY.
+DEFAULT_API_KEY = "AIzaSyCwyYCNqWWA7jqcc5WAG5jQhnGdWKslD4o"   # để trống. Nếu cần, bạn có thể set tạm thời ở ENV.
+
+INTERNAL_MODEL_MAP = {
+    "Gemini 2.5 Flash": "gemini-2.5-flash",
+    "Gemini 2.5 Pro":   "gemini-2.5-pro",
+}
+EXTERNAL_MODEL_NAME = "prithivMLmods/Camel-Doc-OCR-062825 (External)"
+
+try:
+    RESAMPLE = Image.Resampling.LANCZOS  # Pillow >= 10
+except AttributeError:
+    RESAMPLE = Image.LANCZOS             # Pillow < 10
+
+# ================== HELPERS ==================
+import fitz  # PyMuPDF
+
+def pdf_to_images(pdf_bytes: bytes) -> list[Image.Image]:
+    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
+    pages = []
+    for p in doc:
+        pix = p.get_pixmap(dpi=200)
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        pages.append(img)
+    return pages
+
+def ensure_rgb(im: Image.Image) -> Image.Image:
+    return im.convert("RGB") if im.mode != "RGB" else im
+
+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)):
+        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"):
+        return upload.read()
+    raise TypeError(f"Unsupported file object: {type(upload)}")
+
+def _make_previews(file_bytes: bytes, max_side: int = 2000) -> List[Image.Image]:
+    """Trả list PIL.Image đã RGB + resize theo max_side."""
+    if len(file_bytes) >= 4 and file_bytes[:4] == b"%PDF":
+        pages = pdf_to_images(file_bytes)
+    else:
+        pages = [Image.open(io.BytesIO(file_bytes))]
+    out = []
+    for im in pages:
+        im = ensure_rgb(im)
+        if max_side:
+            w, h = im.size
+            scale = min(max_side / float(w), max_side / float(h), 1.0)
+            if scale < 1.0:
+                im = im.resize((max(1, int(w*scale)), max(1, int(h*scale))), RESAMPLE)
+        out.append(im)
+    return out
+
+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"
+    mime, _ = mimetypes.guess_type(filename)
+    if not mime:
+        if len(file_bytes) >= 4 and file_bytes[:4] == b"%PDF":
+            mime = "application/pdf"
+            if not filename.lower().endswith(".pdf"):
+                filename += ".pdf"
+        else:
+            mime = "image/png"
+    return filename, mime
+
+def _extract_json_from_message(msg: str):
+    """Bóc JSON trong ```json ...``` nếu có. Trả về (obj, cleaned_string)."""
+    s = (msg or "").strip()
+    s = re.sub(r"^\s*```(?:json)?\s*", "", s, flags=re.IGNORECASE)
+    s = re.sub(r"\s*```\s*$", "", s)
+    try:
+        return json.loads(s), s
+    except Exception:
+        return None, s
+
+def _pretty_message(msg: str) -> str:
+    obj, s = _extract_json_from_message(msg)
+    return json.dumps(obj, ensure_ascii=False, indent=2) if obj is not None else s
+
+def _safe_text_from_gemini(resp):
+    try:
+        return resp.text
+    except Exception:
+        pass
+    texts = []
+    for c in getattr(resp, "candidates", []) or []:
+        content = getattr(c, "content", None)
+        parts = getattr(content, "parts", None) if content else None
+        if not parts:
+            continue
+        for p in parts:
+            t = getattr(p, "text", None)
+            if t:
+                texts.append(t)
+    return "\n".join(texts).strip()
+
+def _wait_file_active(file_obj, timeout_s: int = 60) -> object:
+    """Chờ file upload sang Gemini ở trạng thái ACTIVE, có timeout + backoff."""
+    start = time.time()
+    delay = 0.5
+    while hasattr(file_obj, "state") and getattr(file_obj.state, "name", "") == "PROCESSING":
+        if time.time() - start > timeout_s:
+            raise TimeoutError("Upload processing timeout.")
+        time.sleep(delay)
+        delay = min(delay * 1.5, 2.0)
+        file_obj = genai.get_file(file_obj.name)
+    if not hasattr(file_obj, "state") or file_obj.state.name != "ACTIVE":
+        st = getattr(file_obj, "state", None)
+        raise RuntimeError(f"Upload failed or not active. State={getattr(st, 'name', 'UNKNOWN')}")
+    return file_obj
+
+# ---------- JSON → Excel (schema-agnostic) ----------
+def _flatten_dict(d: dict, parent_key: str = "", sep: str = ".") -> dict:
+    """Flatten dict lồng nhau thành 1 level: {'a':{'b':1}} -> {'a.b':1}"""
+    items = []
+    for k, v in (d or {}).items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else str(k)
+        if isinstance(v, dict):
+            items.extend(_flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+def _sanitize_sheet_name(name: str, used: set[str]) -> str:
+    # Excel sheet name ≤ 31 chars, không chứa []:*?/\
+    invalid = set(r'[]:*?/\'' + '"')
+    clean = "".join(ch for ch in name if ch not in invalid)
+    clean = clean.strip()
+    if not clean:
+        clean = "sheet"
+    clean = clean[:31]
+    # đảm bảo unique
+    base, idx = clean, 1
+    while clean in used:
+        suffix = f"_{idx}"
+        clean = (base[: (31 - len(suffix))] + suffix)
+        idx += 1
+    used.add(clean)
+    return clean
+
+def _to_excel_generic(data: Any, path: str) -> str:
+    """
+    Quy tắc:
+      - Nếu là list[dict]  -> 1 sheet "data" (json_normalize)
+      - Nếu là dict:
+          + Tạo 1 sheet "summary" từ các field dạng scalar/dict (flatten)
+          + Với mỗi field là list:
+              · list[dict]  -> 1 sheet theo tên key (normalize)
+              · list[scalar]-> 1 sheet 1 cột 'value'
+              · list[mixed] -> chuyển thành cột 'value' dạng chuỗi
+    """
+    with pd.ExcelWriter(path) as writer:
+        used_names = set()
+
+        def add_df(df: pd.DataFrame, sheet: str):
+            sheetname = _sanitize_sheet_name(sheet, used_names)
+            df.to_excel(writer, index=False, sheet_name=sheetname)
+
+        if isinstance(data, list):
+            # list tổng quát
+            try:
+                df = pd.json_normalize(data, sep=".")
+            except Exception:
+                df = pd.DataFrame({"value": [json.dumps(x, ensure_ascii=False) for x in data]})
+            add_df(df, "data")
+            return path
+
+        if isinstance(data, dict):
+            scalars = {}
+            list_sheets: list[tuple[str, pd.DataFrame]] = []
+
+            for k, v in data.items():
+                if isinstance(v, list):
+                    if len(v) == 0:
+                        list_sheets.append((k, pd.DataFrame()))
+                    elif isinstance(v[0], dict):
+                        try:
+                            df = pd.json_normalize(v, sep=".")
+                        except Exception:
+                            df = pd.DataFrame({"value": [json.dumps(x, ensure_ascii=False) for x in v]})
+                        list_sheets.append((k, df))
+                    elif not isinstance(v[0], (list, dict)):
+                        df = pd.DataFrame({"value": v})
+                        list_sheets.append((k, df))
+                    else:
+                        df = pd.DataFrame({"value": [json.dumps(x, ensure_ascii=False) for x in v]})
+                        list_sheets.append((k, df))
+                elif isinstance(v, dict):
+                    scalars.update(_flatten_dict({k: v}))
+                else:
+                    scalars[k] = v
+
+            # summary sheet
+            if len(scalars) > 0:
+                add_df(pd.DataFrame([scalars]), "summary")
+
+            # each list -> one sheet
+            for k, df in list_sheets:
+                add_df(df, k if k else "list")
+
+            # nếu dict chỉ có list, không có summary => vẫn OK (chỉ có các sheet list)
+            return path
+
+        # kiểu khác: ghi thành 1 cột value
+        add_df(pd.DataFrame({"value": [json.dumps(data, ensure_ascii=False)]}), "data")
+        return path
+
+# ================== HANDLERS ==================
+def preview_process(file):
+    """Trả list đường dẫn ảnh PNG tạm cho Gallery (ổn định hơn list PIL)."""
+    if file is None:
+        return []
+    try:
+        file_bytes = _read_file_bytes(file)
+        images = _make_previews(file_bytes, max_side=2000)
+        paths = []
+        for i, im in enumerate(images):
+            fd, path = tempfile.mkstemp(suffix=f"_preview_{i}.png")
+            os.close(fd)
+            im.save(path, format="PNG")
+            paths.append(path)
+        return paths
+    except Exception as e:
+        print(f"Preview error: {e}")
+        return []
+
+# -------- Internal (Gemini) - Base (1 lượt, không thinking) --------
+def run_process_internal_base(file_bytes, filename, mime, question, model_choice,
+                              temperature, top_p):
+    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, INTERNAL_MODEL_MAP["Gemini 2.5 Flash"])
+    gen_config = {"temperature": float(temperature), "top_p": float(top_p)}
+    model = genai.GenerativeModel(model_name=model_name, generation_config=gen_config)
+
+    uploaded = None
+    tmp_path = None
+    try:
+        if file_bytes:
+            suffix = os.path.splitext(filename)[1] or ".bin"
+            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                tmp.write(file_bytes)
+                tmp_path = tmp.name
+            uploaded = genai.upload_file(path=tmp_path, mime_type=mime)
+            uploaded = _wait_file_active(uploaded, timeout_s=60)
+
+        user_prompt = (question or "").strip()
+        if not user_prompt:
+            user_prompt = (
+                "Perform high-quality OCR on the provided file. If PDF: read all pages in order. "
+                "Return clean plain text. If structure is obvious (tables, key:value), preserve it. "
+                "If you can, output JSON that captures the structure."
+            )
+
+        # Gọi model
+        if uploaded:
+            resp = model.generate_content([user_prompt, uploaded])
+        else:
+            resp = model.generate_content(user_prompt)
+
+        # Lấy đúng message LLM (pretty nếu là JSON)
+        answer_raw = _safe_text_from_gemini(resp)
+        message = _pretty_message(answer_raw)
+
+        # Parse JSON (nếu có) để export. Không validate schema.
+        parsed_obj, _ = _extract_json_from_message(answer_raw)
+
+        return message, parsed_obj
+    finally:
+        if tmp_path and os.path.exists(tmp_path):
+            try: os.remove(tmp_path)
+            except Exception: pass
+        try:
+            if uploaded and hasattr(uploaded, "name"):
+                genai.delete_file(uploaded.name)
+        except Exception:
+            pass
+
+# -------- External API --------
+def run_process_external(file_bytes, filename, mime, question, api_url,
+                         temperature, top_p):
+    if not api_url or not str(api_url).strip():
+        return "ERROR: Missing external API endpoint (hãy dán URL).", None
+    try:
+        user_prompt = (question or "").strip()
+        if not user_prompt:
+            user_prompt = (
+                "Perform high-quality OCR on the provided file. If PDF: read all pages in order. "
+                "Return clean plain text. If structure is obvious (tables, key:value), preserve it. "
+                "If you can, output JSON that captures the structure."
+            )
+
+        data = {"prompt": user_prompt, "temperature": str(temperature), "top_p": str(top_p)}
+
+        if file_bytes:
+            files = {"file": (filename, file_bytes, mime)}
+            r = requests.post(api_url, files=files, data=data, timeout=60)
+        else:
+            r = requests.post(api_url, json=data, timeout=60)
+
+        if r.status_code >= 400:
+            return f"ERROR: External API HTTP {r.status_code}: {r.text[:300]}", None
+
+        answer = None
+        try:
+            j = r.json()
+            answer = j.get("message") or j.get("text") or j.get("data")
+            if isinstance(answer, (dict, list)):
+                answer = json.dumps(answer, ensure_ascii=False)
+        except Exception:
+            answer = r.text
+
+        answer = (answer or "").strip()
+        message = _pretty_message(answer)
+        parsed_obj, _ = _extract_json_from_message(answer)
+
+        return message, parsed_obj
+    except Exception as e:
+        return f"ERROR: {type(e).__name__}: {str(e) or repr(e)}", None
+
+# -------- Router --------
+def run_process(file, question, model_choice, temperature, top_p, external_api_url):
+    """
+    Router (không Agent, không thinking):
+      - Nếu chọn External model -> run_process_external
+      - Ngược lại -> Gemini nội bộ (Base 1 lượt)
+    """
+    try:
+        has_file = file is not None
+        file_bytes = filename = mime = None
+        if has_file:
+            file_bytes = _read_file_bytes(file)
+            filename, mime = _guess_name_and_mime(file, file_bytes)
+
+        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
+            )
+
+        return run_process_internal_base(
+            file_bytes=file_bytes, filename=filename, mime=mime,
+            question=question, model_choice=model_choice,
+            temperature=temperature, top_p=top_p
+        )
+    except Exception as e:
+        return f"ERROR: {type(e).__name__}: {str(e) or repr(e)}", None
+
+def on_export_excel(parsed_obj):
+    try:
+        if not parsed_obj:
+            return None  # UI sẽ không hiện file nếu None
+        tmp_path = os.path.join(tempfile.gettempdir(), f"export_{int(time.time())}.xlsx")
+        return _to_excel_generic(parsed_obj, tmp_path)
+    except Exception as e:
+        print(f"Export error: {e}")
+        return None
+
+def clear_all():
+    # file, preview, output_text, question, model, parsed_state, download,
+    # temperature, top_p, external_api_url
+    return (
+        None, [], "", "",
+        "Gemini 2.5 Flash", None, None,
+        0.2, 0.95, ""
+    )
+
+# ================== UI ==================
+def _toggle_external_visibility(selected: str):
+    return gr.update(visible=(selected == EXTERNAL_MODEL_NAME))
+
+def main():
+    custom_css = """
+    .gradio-container { max-width: 1400px !important; margin: 0 auto; }
+    #main-row { display: flex; gap: 20px; align-items: flex-start; }
+    #left-column { flex: 1; min-width: 400px; max-width: 600px; }
+    #right-column { flex: 1; min-width: 400px; }
+    #file-upload { border: 2px dashed #d1d5db; border-radius: 12px; padding: 20px; text-align: center; transition: border-color 0.3s ease; }
+    #file-upload:hover { border-color: #3b82f6; }
+    #preview-gallery { max-height: 600px; overflow-y: auto; border: 1px solid #e5e7eb; border-radius: 12px; background: #f9fafb; padding: 10px; }
+    #preview-gallery .grid { grid-template-columns: 1fr !important; gap: 10px !important; }
+    #preview-gallery img { width: 100% !important; height: auto !important; object-fit: contain !important; background: white; }
+    #controls-section { background: #f8fafc; padding: 20px; border-radius: 12px; margin-bottom: 20px; }
+    #results-section { background: #ffffff; border: 1px solid #e5e7eb; border-radius: 12px; padding: 20px; }
+    #llm-output { max-height: 500px; overflow-y: auto; font-family: monospace; font-size: 13px; }
+    .primary-button { background: linear-gradient(90deg, #3b82f6, #1d4ed8) !important; color: white !important; border: none !important; border-radius: 8px !important; padding: 10px 20px !important; font-weight: 500 !important; }
+    .primary-button:hover { transform: translateY(-1px) !important; box-shadow: 0 4px 12px rgba(59, 130, 246, 0.3) !important; }
+    .secondary-button { background: #f3f4f6 !important; color: #374151 !important; border: 1px solid #d1d5db !important; border-radius: 8px !important; padding: 8px 16px !important; }
+    @media (max-width: 1024px) { #main-row { flex-direction: column; } #left-column, #right-column { min-width: 100%; max-width: 100%; } }
+    """
+
+    with gr.Blocks(title="OCR Multi-Agent System", css=custom_css, theme=gr.themes.Soft()) as demo:
+        gr.HTML("""
+        <div style="text-align: center; padding: 20px 0; margin-bottom: 30px;">
+          <h1 style="color:#1f2937; font-size: 2.5rem; font-weight: bold; margin-bottom: 8px;">📄 OCR Extraction (LLM-first)</h1>
+          <p style="color:#6b7280; font-size: 1.1rem; margin: 0;">Upload PDF/images → LLM produces raw text/JSON → Export Excel (schema-agnostic)</p>
+        </div>
+        """)
+
+        last_parsed_state = gr.State(value=None)
+
+        with gr.Row(elem_id="main-row"):
+            # Left
+            with gr.Column(elem_id="left-column"):
+                gr.Markdown("### 📁 Upload Document")
+                file = gr.File(
+                    label="Choose PDF or Image file",
+                    file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
+                    type="filepath",
+                    elem_id="file-upload"
+                )
+                gr.Markdown("### 👁️ Document Preview")
+                preview = gr.Gallery(columns=1, height=None, show_label=False, elem_id="preview-gallery", allow_preview=True)
+
+            # Right
+            with gr.Column(elem_id="right-column"):
+                with gr.Group(elem_id="controls-section"):
+                    gr.Markdown("### ⚙️ Processing Options")
+                    with gr.Row():
+                        model_choice = gr.Dropdown(
+                            choices=[*INTERNAL_MODEL_MAP.keys(), EXTERNAL_MODEL_NAME],
+                            value="Gemini 2.5 Flash",
+                            label="Model"
+                        )
+
+                    with gr.Row():
+                        temperature = gr.Slider(0.0, 2.0, value=0.2, step=0.05, label="temperature")
+                        top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="top_p")
+
+                    external_api_url = gr.Textbox(
+                        label="External API endpoint (URL)",
+                        placeholder="https://your-host/path/to/ocr",
+                        visible=False
+                    )
+
+                    question = gr.Textbox(
+                        label="Custom Prompt (optional)",
+                        placeholder="Leave blank for default OCR; or ask model to output JSON by your own schema...",
+                        lines=3
+                    )
+                    with gr.Row():
+                        run_btn = gr.Button("🚀 Process Document", elem_classes=["primary-button"])
+                        clear_btn = gr.Button("🗑️ Clear All", elem_classes=["secondary-button"])
+
+                with gr.Group(elem_id="results-section"):
+                    gr.Markdown("### 📊 LLM Message (raw/pretty)")
+                    output_text = gr.Code(label="LLM Message", language="json", elem_id="llm-output")
+                    with gr.Row():
+                        export_btn = gr.Button("⬇️ Export to Excel", elem_classes=["secondary-button"])
+                        download_file = gr.File(label="Download Excel", interactive=False)
+
+        # Events
+        file.change(preview_process, inputs=[file], outputs=[preview])
+        model_choice.change(_toggle_external_visibility, inputs=[model_choice], outputs=[external_api_url])
+
+        run_btn.click(
+            run_process,
+            inputs=[file, question, model_choice, temperature, top_p, external_api_url],
+            outputs=[output_text, last_parsed_state]
+        )
+
+        export_btn.click(on_export_excel, inputs=[last_parsed_state], outputs=[download_file])
+
+        clear_btn.click(
+            clear_all,
+            inputs=[],
+            outputs=[file, preview, output_text, question, model_choice, last_parsed_state,
+                     download_file, temperature, top_p, external_api_url]
+        )
+
+    return demo
+
+demo = main()
+
+if __name__ == "__main__":
+    demo.launch()