chatkit/backend/app/file_processor.py

"""File processing orchestrator.

Delegates extraction to a pluggable parser chosen by name or auto-picked.
Saves the parsed structured data to the database, along with metadata about
which parser produced it.
"""

from __future__ import annotations

import logging

from fastapi import UploadFile

from .answer_grid import seed_from_parsed, to_dict as grid_to_dict
from .database import delete_parsed_data, get_parsed_data, save_answer_grid, save_parsed_data
from .parsers import AUTO, get_parser, pick_auto
from .parsers.base import ParserFile, ParserResult

logger = logging.getLogger(__name__)


async def process_uploaded_files(
    files: list[UploadFile],
    data_type: str,
    session_id: int,
    description: str = "",
    model: str = "gpt-5.4",
    parser: str = AUTO,
) -> dict:
    """Parse uploaded files and persist the structured result.

    `data_type='answers'` triggers a dual extraction: the same files are parsed
    twice (once for student_answers, once for teacher_answers) and both rows are
    stored. The response merges both shapes under their respective keys.
    """
    parser_files: list[ParserFile] = []
    for file in files:
        content = await file.read()
        if not content:
            raise ValueError(f"File '{file.filename}' is empty.")
        parser_files.append(
            ParserFile(filename=file.filename or "unknown", content=content)
        )

    if not parser_files:
        raise ValueError("No files uploaded")

    if data_type == "answers":
        return await _process_combined_answers(
            parser_files, session_id, description, model, parser
        )

    result = await _run_with_fallback(parser_files, data_type, description, model, parser)
    await _persist(session_id, data_type, parser_files, result)
    return _response_payload(result)


async def _process_combined_answers(
    parser_files: list[ParserFile],
    session_id: int,
    description: str,
    model: str,
    parser: str,
) -> dict:
    student_result = await _run_with_fallback(
        parser_files, "student_answers", description, model, parser
    )
    teacher_result = await _run_with_fallback(
        parser_files, "teacher_answers", description, model, parser
    )

    await _persist(session_id, "student_answers", parser_files, student_result)
    await _persist(session_id, "teacher_answers", parser_files, teacher_result)

    # Seed and save a draft answer grid so step 2 has all data (students + correct answers).
    # Use only the answers extracted from teacher_result to avoid inheriting the question count
    # from a separately uploaded questions PDF (which may have a different total).
    t_data = teacher_result.data
    draft_grid = seed_from_parsed({}, student_result.data, t_data)
    await save_answer_grid(session_id, grid_to_dict(draft_grid))

    return {
        "student_answers": _response_payload(student_result),
        "teacher_answers": _response_payload(teacher_result),
    }


async def _persist(
    session_id: int,
    data_type: str,
    parser_files: list[ParserFile],
    result: ParserResult,
) -> None:
    if data_type == "questions":
        await delete_parsed_data(session_id, data_type)
        for q in result.data.get("questions", []):
            text = (q.get("text", "") or "").strip()
            options = q.get("options") or []
            if not text and not options:
                logger.warning(
                    "session %d: question %s extracted with no text and no "
                    "options — likely a failed extraction, not a legitimate "
                    "picture-only question",
                    session_id,
                    q.get("number", "?"),
                )
            await save_parsed_data(
                session_id,
                q.get("number", 0),
                text,
                "",
                "",
                [],
                options=options,
            )
    elif data_type == "teacher_answers":
        existing = {r["question_num"]: r for r in await get_parsed_data(session_id)}
        answers_map = {
            a.get("question_number", 0): (a.get("correct_answer") or a.get("answer") or "")
            for a in result.data.get("answers", [])
        }
        all_nums = sorted(set(existing) | set(answers_map))
        await delete_parsed_data(session_id, data_type)
        for qnum in all_nums:
            row = existing.get(qnum, {})
            await save_parsed_data(
                session_id,
                qnum,
                row.get("question_str", ""),
                answers_map.get(qnum, row.get("answer", "")),
                row.get("main_category", ""),
                row.get("tags", []),
                options=row.get("options", []),
            )
    # student_answers are not stored in ParsedData


def _response_payload(result: ParserResult) -> dict:
    return {
        **result.data,
        "_meta": {
            "parser": result.parser_name,
            "notes": list(result.notes),
        },
    }


async def _run_with_fallback(
    parser_files: list[ParserFile],
    data_type: str,
    description: str,
    model: str,
    parser_name: str,
) -> ParserResult:
    """Run the selected parser; in auto mode, fall back to LLM on failure."""
    if parser_name == AUTO:
        first_file = parser_files[0]
        chosen = pick_auto(first_file.content, first_file.filename, data_type)
        try:
            return await chosen.parse(parser_files, data_type, description, model)
        except Exception as primary_err:
            if chosen.name == "llm_vision":
                raise
            # Fallback to LLM vision on any native-parser failure
            llm = get_parser("llm_vision")
            result = await llm.parse(parser_files, data_type, description, model)
            return ParserResult(
                data=result.data,
                parser_name=result.parser_name,
                notes=(f"auto-fallback from {chosen.name}: {primary_err}", *result.notes),
            )

    chosen = get_parser(parser_name)
    return await chosen.parse(parser_files, data_type, description, model)