agents/document_agent.py

# -----------------------------------------------
# document_agent.py
# Reads uploaded medical document (PDF or image)
# Extracts all text and medical data using Gemini
# Returns structured text for next agents
# -----------------------------------------------

from google import genai
from google.genai import types
from pypdf import PdfReader
from dotenv import load_dotenv
import os
import base64

load_dotenv()
client = genai.Client()

# -----------------------------------------------
# Extract text from PDF using pypdf
# -----------------------------------------------
def extract_pdf_text(file_path):
    try:
        reader = PdfReader(file_path)
        text = ""
        for page in reader.pages:
            text += page.extract_text() + "\n"
        return text.strip()
    except Exception as e:
        return None

# -----------------------------------------------
# Convert image to base64 for Gemini Vision
# -----------------------------------------------
def image_to_base64(file_path):
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")

# -----------------------------------------------
# Get mime type from file extension
# -----------------------------------------------
def get_mime_type(file_path):
    ext = os.path.splitext(file_path)[1].lower()
    mime_types = {
        ".jpg":  "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png":  "image/png",
        ".pdf":  "application/pdf",
    }
    return mime_types.get(ext, "image/jpeg")

# -----------------------------------------------
# Analyze document with Gemini Vision
# -----------------------------------------------
def analyze_with_gemini(content, is_image=False, mime_type=None):
    prompt = """You are a medical document analyzer.
Analyze this medical document carefully and extract ALL information.

Return your response in this EXACT format:

DOCUMENT_TYPE: [blood report / X-ray report / MRI report / prescription / discharge summary / other]
PATIENT_INFO: [any patient details found — name, age, gender, date]
FINDINGS: [all medical values, test results, or observations found]
ABNORMALITIES: [any values marked as high/low/abnormal, or any concerning findings]
DOCTOR_NOTES: [any doctor observations or impressions]
RAW_TEXT: [complete extracted text from document]

Be thorough. Extract every single medical value and finding you can see.
If information is not available for a field, write: Not found"""

    for attempt in range(3):
        try:
            if is_image:
                response = client.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=[
                        prompt,
                        types.Part.from_bytes(
                            data=base64.b64decode(content),
                            mime_type=mime_type
                        )
                    ]
                )
            else:
                response = client.models.generate_content(
                    model="gemini-2.5-flash",
                    contents=f"{prompt}\n\nDOCUMENT TEXT:\n{content}"
                )
            return response.text
        except Exception as e:
            error_msg = str(e)
            if "503" in error_msg or "429" in error_msg or "UNAVAILABLE" in error_msg or "exhausted" in error_msg.lower():
                wait_time = (attempt + 1) * 5
                print(f"[Document Agent] ⚠️  Rate limit — retrying in {wait_time}s... (attempt {attempt+1}/3)")
                import time
                time.sleep(wait_time)
            else:
                return f"ERROR: {e}"
    return "ERROR: Gemini API unavailable after 3 retries"

# -----------------------------------------------
# Check if Gemini extracted meaningful content
# -----------------------------------------------
def is_valid_extraction(parsed: dict) -> bool:
    findings = parsed.get("FINDINGS", "")
    raw_text = parsed.get("RAW_TEXT", "")
    doc_type = parsed.get("DOCUMENT_TYPE", "")

    if len(findings) < 20 and len(raw_text) < 50:
        return False

    if doc_type.lower() in ["unknown", "not found", ""]:
        return False

    return True

# -----------------------------------------------
# Main document agent function
# Called by LangGraph with state
# -----------------------------------------------
def run_document_agent(state: dict) -> dict:
    file_path = state.get("file_path")
    print(f"\n[Document Agent] Processing: {file_path}")

    if not file_path or not os.path.exists(file_path):
        state["error"]         = "File not found"
        state["raw_text"]      = ""
        state["document_type"] = "unknown"
        return state

    ext = os.path.splitext(file_path)[1].lower()

    # --- Handle PDF ---
    if ext == ".pdf":
        print("[Document Agent] PDF detected — extracting text...")
        pdf_text = extract_pdf_text(file_path)

        if pdf_text and len(pdf_text) > 100:
            print(f"[Document Agent] Extracted {len(pdf_text)} characters from PDF")
            gemini_output = analyze_with_gemini(pdf_text, is_image=False)
        else:
            print("[Document Agent] Scanned PDF detected — using Gemini Vision...")
            image_data = image_to_base64(file_path)
            gemini_output = analyze_with_gemini(
                image_data,
                is_image=True,
                mime_type="application/pdf"
            )

    # --- Handle Image ---
    elif ext in [".jpg", ".jpeg", ".png"]:
        print("[Document Agent] Image detected — using Gemini Vision...")
        image_data = image_to_base64(file_path)
        mime_type  = get_mime_type(file_path)
        gemini_output = analyze_with_gemini(
            image_data,
            is_image=True,
            mime_type=mime_type
        )

        # Quality check for phone-clicked photos
        parsed_check = parse_gemini_output(gemini_output)
        if not is_valid_extraction(parsed_check):
            state["error"] = (
                "Could not read the document clearly. Please:\n"
                "• Ensure good lighting\n"
                "• Place report flat on a surface\n"
                "• Take photo directly from above\n"
                "• Make sure text is clearly visible"
            )
            state["raw_text"]      = ""
            state["document_type"] = "unknown"
            return state

    else:
        state["error"]         = f"Unsupported format: {ext}"
        state["raw_text"]      = ""
        state["document_type"] = "unknown"
        return state

    # --- Parse Gemini output ---
    print("[Document Agent] Parsing Gemini output...")
    parsed = parse_gemini_output(gemini_output)

    # --- Update state ---
    state["raw_text"]      = parsed.get("RAW_TEXT", gemini_output)
    state["document_type"] = parsed.get("DOCUMENT_TYPE", "unknown")
    state["findings"]      = parsed.get("FINDINGS", "")
    state["abnormalities"] = parsed.get("ABNORMALITIES", "")
    state["patient_info"]  = parsed.get("PATIENT_INFO", "")
    state["doctor_notes"]  = parsed.get("DOCTOR_NOTES", "")
    state["gemini_raw"]    = gemini_output

    # Flag imaging documents for stronger disclaimer
    if any(word in state["document_type"].lower()
           for word in ["x-ray", "xray", "mri", "ct", "scan", "ultrasound"]):
        state["is_imaging"] = True
    else:
        state["is_imaging"] = False

    print(f"[Document Agent] ✅ Done — Document type: {state['document_type']}")
    print(f"[Document Agent] Is imaging: {state['is_imaging']}")
    return state

# -----------------------------------------------
# Parse Gemini structured output into dict
# -----------------------------------------------
def parse_gemini_output(text):
    result = {}
    current_key   = None
    current_value = []

    for line in text.split("\n"):
        found_key = False
        for key in ["DOCUMENT_TYPE", "PATIENT_INFO", "FINDINGS",
                    "ABNORMALITIES", "DOCTOR_NOTES", "RAW_TEXT"]:
            if line.startswith(f"{key}:"):
                if current_key:
                    result[current_key] = " ".join(current_value).strip()
                current_key   = key
                current_value = [line[len(key)+1:].strip()]
                found_key     = True
                break

        if not found_key and current_key:
            current_value.append(line.strip())

    if current_key:
        result[current_key] = " ".join(current_value).strip()

    return result