chatkit/backend/app/file_processor.py

"""File processing: upload files and extract structured Q&A via GPT-4 Vision / GPT-4o."""

from __future__ import annotations

import base64
import json
import mimetypes
from pathlib import Path
from typing import Optional

import fitz  # pymupdf
from openai import AsyncOpenAI
from fastapi import UploadFile

from .config import get_settings
from .database import save_parsed_data, delete_parsed_data


# Extensions that should be read as text, not sent as images
TEXT_EXTENSIONS = {".txt", ".csv", ".tsv", ".json", ".xml", ".md", ".html", ".htm", ".log"}

# Image extensions with their MIME types
IMAGE_MIME = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".bmp": "image/bmp",
    ".webp": "image/webp",
    ".tiff": "image/tiff",
    ".tif": "image/tiff",
}

NO_FABRICATION_RULE = """
CRITICAL: Only extract data that is ACTUALLY present in the provided content.
Do NOT invent, fabricate, or hallucinate any data.
If the content is empty, unreadable, or does not contain the expected data, return an empty result like {"questions": []} or {"students": []} or {"answers": []}.
"""

IMAGE_DESCRIPTION_RULE = """
IMPORTANT - Images, Diagrams, Charts, and Figures:
- If a question contains or references an image, diagram, chart, graph, map, table, or any visual element, you MUST describe it in detail as part of the question text.
- Use the format: [圖片描述: ...detailed description...] inserted where the image appears in the question.
- The description must be detailed enough that someone who CANNOT see the image can still fully understand and answer the question.
- Include all relevant data points, labels, axes, values, shapes, positions, colors, relationships, and any text visible in the image.
- For charts/graphs: describe the type, axes labels, data values, trends, and all visible data points.
- For diagrams: describe all components, connections, labels, measurements, and spatial relationships.
- For maps: describe locations, boundaries, labels, and any marked features.
- For tables embedded as images: transcribe the full table content.
"""

EXTRACTION_PROMPTS = {
    "questions": """You are an exam data extraction expert. Extract all exam questions from the provided content.

Return a JSON object with this structure:
{
  "questions": [
    {
      "number": 1,
      "text": "the full question text (including [圖片描述: ...] for any images/diagrams)",
      "type": "multiple_choice" or "short_answer" or "essay" or "true_false",
      "options": ["A) ...", "B) ...", "C) ...", "D) ..."],  // null if not multiple choice
      "points": 10  // point value if visible, null otherwise
    }
  ]
}

Rules:
- Extract EVERY question you can see
- Preserve the original language of the questions
- If options are present, include them exactly as written
- If point values are shown, include them
""" + IMAGE_DESCRIPTION_RULE + NO_FABRICATION_RULE + "- Return valid JSON only",

    "student_answers": """You are an exam data extraction expert. Extract all student answers from the provided content.

Return a JSON object with this structure:
{
  "students": [
    {
      "name": "student name or ID",
      "id": "student ID if visible",
      "answers": [
        {
          "question_number": 1,
          "answer": "the student's answer text"
        }
      ]
    }
  ]
}

Rules:
- Extract answers for EVERY student visible
- Preserve exact answer text
- If a student left a question blank, set answer to null
- Use the student's name or ID as shown in the document
- If a student's answer includes a drawing or diagram, describe it as text: [圖片描述: ...]
""" + IMAGE_DESCRIPTION_RULE + NO_FABRICATION_RULE + "- Return valid JSON only",

    "teacher_answers": """You are an exam data extraction expert. Extract the correct/model answers (answer key) from the provided content.

Return a JSON object with this structure:
{
  "answers": [
    {
      "question_number": 1,
      "correct_answer": "the correct answer",
      "explanation": "explanation if provided, otherwise null"
    }
  ]
}

Rules:
- Extract the correct answer for EVERY question
- Include explanations if they are provided
- Preserve exact answer text
- If the answer references or includes an image/diagram, describe it as text: [圖片描述: ...]
""" + IMAGE_DESCRIPTION_RULE + NO_FABRICATION_RULE + "- Return valid JSON only",
}


def pdf_to_images(pdf_bytes: bytes) -> list[bytes]:
    """Convert PDF bytes to a list of PNG image bytes, one per page."""
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    images = []
    for page in doc:
        # Render at 2x resolution for better OCR quality
        pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
        images.append(pix.tobytes("png"))
    doc.close()
    return images


def try_read_as_text(file_bytes: bytes, filename: str) -> Optional[str]:
    """Try to decode file bytes as text. Returns None if not text."""
    ext = Path(filename).suffix.lower()
    if ext in TEXT_EXTENSIONS:
        for encoding in ("utf-8", "utf-8-sig", "big5", "gb2312", "shift_jis", "latin-1"):
            try:
                return file_bytes.decode(encoding)
            except (UnicodeDecodeError, ValueError):
                continue
    return None


def get_image_mime(filename: str) -> Optional[str]:
    """Get MIME type for image files. Returns None if not a known image type."""
    ext = Path(filename).suffix.lower()
    return IMAGE_MIME.get(ext)


async def process_uploaded_files(
    files: list[UploadFile],
    data_type: str,
    session_id: int,
    description: str = "",
    model: str = "gpt-5.4",
) -> dict:
    """
    Process uploaded files: convert to images/text, then extract data via GPT-4o.
    Returns the structured data extracted.
    """
    image_parts: list[dict] = []      # image content parts for GPT
    text_parts: list[str] = []         # text content from text files
    source_files: list[str] = []

    for file in files:
        file_bytes = await file.read()
        filename = file.filename or "unknown"

        if not file_bytes:
            raise ValueError(f"File '{filename}' is empty (0 bytes). Please upload a valid file.")

        ext = Path(filename).suffix.lower()
        source_files.append(filename)

        if ext == ".pdf":
            page_images = pdf_to_images(file_bytes)
            for i, img in enumerate(page_images):
                b64 = base64.b64encode(img).decode()
                image_parts.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}", "detail": "high"},
                })
        else:
            # Try to read as text first
            text_content = try_read_as_text(file_bytes, filename)
            if text_content is not None:
                text_parts.append(f"--- Content from {filename} ---\n{text_content}")
            else:
                # Try as image with correct MIME type
                mime = get_image_mime(filename) or "image/png"
                b64 = base64.b64encode(file_bytes).decode()
                image_parts.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:{mime};base64,{b64}", "detail": "high"},
                })

    if not image_parts and not text_parts:
        raise ValueError("No valid content found in uploaded files")

    # Extract data using selected model
    structured = await extract_data(
        image_parts=image_parts,
        text_parts=text_parts,
        data_type=data_type,
        description=description,
        model=model,
    )

    # Clear old data of this type for the session, then save new
    await delete_parsed_data(session_id, data_type)

    raw_text = json.dumps({"source_files": source_files}, ensure_ascii=False)
    await save_parsed_data(session_id, data_type, files[0].filename or "upload", raw_text, structured)

    return structured


async def extract_data(
    image_parts: list[dict],
    text_parts: list[str],
    data_type: str,
    description: str = "",
    model: str = "gpt-5.4",
) -> dict:
    """Call GPT-4o to extract structured data from images and/or text."""
    settings = get_settings()
    client = AsyncOpenAI(api_key=settings.openai_api_key)

    system_prompt = EXTRACTION_PROMPTS.get(data_type)
    if not system_prompt:
        raise ValueError(f"Unknown data_type: {data_type}")

    # Build user message content
    user_text = f"Please extract the {data_type.replace('_', ' ')} from the provided content."
    if description.strip():
        user_text = f"Context about this data: {description.strip()}\n\n{user_text}"

    # If we have text content, include it
    if text_parts:
        combined_text = "\n\n".join(text_parts)
        user_text += f"\n\n--- FILE CONTENT ---\n{combined_text}"

    content_parts: list[dict] = [{"type": "text", "text": user_text}]
    content_parts.extend(image_parts)

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content_parts},
    ]

    # Build kwargs — some models don't support all parameters
    kwargs: dict = {
        "model": model,
        "messages": messages,
        "max_completion_tokens": 16384,
    }
    # json_object response_format may not work on all models; try with it first
    try:
        response = await client.chat.completions.create(
            **kwargs,
            response_format={"type": "json_object"},
        )
    except Exception:
        # Fallback: no response_format, rely on prompt to return JSON
        response = await client.chat.completions.create(**kwargs)

    choice = response.choices[0]
    content = choice.message.content
    if not content:
        reason = getattr(choice, "finish_reason", "unknown")
        refusal = getattr(choice.message, "refusal", None)
        detail = f"finish_reason={reason}"
        if refusal:
            detail += f", refusal={refusal}"
        raise ValueError(f"Model returned empty content ({detail}). Try a different model.")
    return json.loads(content)