Update app.py

app.py

@@ -196,9 +196,62 @@ def extract_pdf_note(file_path: str) -> str:
     except Exception as e:
         print(f"⚠️ extract_pdf_note lỗi: {e}")
         return ""
+def extract_airline_header_via_ocr(file_path: str) -> str:
+    """
+    Dùng Gemini OCR nhận diện hãng bay ở trang đầu PDF.
+    ⚡ Tối ưu: chỉ lấy 1 trang đầu, DPI=120, JPEG quality=60 để giảm dung lượng.
+    """
+    import google.generativeai as genai
+    from PIL import Image
+    import fitz, io, tempfile, os
+
+    api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
+    genai.configure(api_key=api_key)
+    model = genai.GenerativeModel("gemini-2.5-flash")
+
+    # --- Chuyển trang đầu PDF thành ảnh (giảm DPI và nén) ---
+    pdf = fitz.open(file_path)
+    pix = pdf[0].get_pixmap(dpi=120)  # ⚡ DPI thấp hơn giúp nhẹ hơn nhiều
+    img = Image.open(io.BytesIO(pix.tobytes("png"))).convert("RGB")
+
+    # Nén ảnh JPEG chất lượng thấp hơn để nhẹ KB
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
+        img.save(tmp.name, format="JPEG", quality=60, optimize=True)  # ⚡ chỉ còn ~150–250KB
+        img_path = tmp.name
+
+    # --- Upload nhẹ hơn nhiều ---
+    uploaded = genai.upload_file(path=img_path, mime_type="image/jpeg")
+
+    # --- Prompt yêu cầu nhận diện header ---
+    prompt = """
+    Identify from this airline rate sheet:
+    - Airline name (e.g. Qatar Airways, Turkish Airlines)
+    - Airline code (e.g. QR, TK, EK, VN)
+    - Title (e.g. SGN PRICING NOV25)
+    - Validity info (e.g. Effective from 01 Nov 2025, Until Further Notice)
+    Return JSON with fields: airline_name, airline_code, title, valid_from, valid_to.
+    """
 
-def call_gemini_with_prompt(content_text: str, note_text: str, question: str, model_choice: str, temperature: float, top_p: float):
-    """Gửi bảng + note vào Gemini (ưu tiên prompt tùy chỉnh nếu có)"""
+    resp = model.generate_content([prompt, uploaded])
+    genai.delete_file(uploaded.name)
+
+    result = getattr(resp, "text", "").strip()
+    print("🛫 OCR header (compressed):", result)
+    return result
+def call_gemini_with_prompt(
+    header: str,
+    content_text: str,
+    note_text: str,
+    question: str,
+    model_choice: str,
+    temperature: float,
+    top_p: float
+):
+    """
+    Gửi header + bảng CSV + note vào Gemini.
+    Ưu tiên: nếu user nhập prompt riêng → dùng prompt đó, ngược lại dùng PROMPT_FREIGHT_JSON.
+    Header (nếu có) sẽ được chèn thêm vào đầu để giúp model nhận diện hãng bay, thời gian hiệu lực, v.v.
+    """
     api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY)
     genai.configure(api_key=api_key)
 
@@ -210,23 +263,39 @@ def call_gemini_with_prompt(content_text: str, note_text: str, question: str, mo
         }
     )
 
-    # Nếu user không nhập câu hỏi riêng, dùng prompt chuẩn FREIGHT_JSON
+    # --- Xác định prompt chính ---
     base_prompt = question.strip() if question and question.strip() else PROMPT_FREIGHT_JSON
 
-    prompt = f"""
-                {base_prompt}
-                
-                Below is the extracted CSV data:
-                {content_text}
-                
-                Below are the notes extracted from the PDF (e.g. Valid From, Origin, Remark, Package Type rules):
-                {note_text}
-                
-                Please analyze all data and generate the JSON output following the schema above.
-                """
-
-    print("🧠 Sending prompt to Gemini...")
-    response = model.generate_content(prompt)
+    # --- Ghép nội dung ---
+    prompt_parts = [base_prompt]
+
+    if header and header.strip():
+        prompt_parts.append(f"""
+        ### Header information (from first page OCR or PDF header):
+        {header}
+        """)
+
+    prompt_parts.append(f"""
+    ### Extracted table data (CSV format):
+    {content_text}
+    """)
+
+    if note_text and note_text.strip():
+        prompt_parts.append(f"""
+        ### Notes or remarks extracted from the PDF:
+        {note_text}
+        """)
+
+    prompt_parts.append("""
+    Please analyze all data (header + table + notes) and generate the final JSON output
+    following the defined schema above. Ensure that any airline, date, or rule from header/note
+    is merged into the JSON result (e.g. shipping_line, valid_from, valid_to, remarks, etc.).
+    """)
+
+    full_prompt = "\n".join(prompt_parts)
+
+    print("🧠 Sending full prompt (with header if available) to Gemini...")
+    response = model.generate_content(full_prompt)
     result_text = getattr(response, "text", str(response))
 
     return result_text
@@ -251,11 +320,11 @@ def run_process(file, question, model_choice, temperature, top_p, external_api_u
             # 1️⃣ Extract bảng bằng Camelot
             df = extract_pdf_tables(tmp_path)
             note_text = extract_pdf_note(tmp_path)
-
+            header = extract_airline_header_via_ocr()
             if not df.empty:
                 csv_text = df.to_csv(index=False)
                 print("✅ Gửi Gemini để sinh JSON...")
-                message = call_gemini_with_prompt(csv_text, note_text, question, model_choice, temperature, top_p)
+                message = call_gemini_with_prompt(header, csv_text, note_text, question, model_choice, temperature, top_p)
                 return message, None
             else:
                 print("⚠️ Không có bảng hợp lệ, fallback OCR Gemini.")