chatkit/backend/app/question_categorizer.py

"""Batch categorization of uploaded exam questions into the pre-defined taxonomy.

Called once after a questions PDF is processed. Assigns each question a
main_category and tags, updates ParsedData, and seeds the shared QuestionBank.
"""

from __future__ import annotations

import json
import logging
import re

from openai import AsyncOpenAI

from .config import get_settings
from .database import save_question_bank_batch, update_parsed_data_categories
from .taxonomy import MAIN_CATEGORIES, taxonomy_prompt_block

logger = logging.getLogger(__name__)

CATEGORIZE_SYSTEM_PROMPT = f"""You are ClassLens, an assistant for Taiwanese 國中 English exam analysis.

Given a list of English exam questions, classify each one into the pre-defined taxonomy below.
Return ONLY a JSON array — no markdown fences, no extra text:
[{{"question_num": 1, "main_category": "語法結構", "tags": ["時態", "過去簡單式"]}}, ...]

Rules:
- main_category MUST be exactly one of the 6 categories listed below.
- tags must be taken directly from the tag list for that category. Use an empty array [] if no tag fits.
- Classify based on the PRIMARY grammar or vocabulary skill being tested.

{taxonomy_prompt_block()}"""


async def categorize_questions(
    session_id: int,
    teacher_id: int,
    questions: list[dict],
    model: str = "gpt-4o-mini",
) -> None:
    """Classify questions, update ParsedData categories, and seed QuestionBank.

    `questions` is a list of dicts with at minimum {question_num, question_text, answer}.
    Errors are logged and swallowed — categorization failure must not block upload.
    """
    if not questions:
        return

    settings = get_settings()
    client = AsyncOpenAI(api_key=settings.openai_api_key)

    # Build compact input for the LLM
    items = "\n".join(
        f'{q["question_num"]}. {q.get("question_text") or q.get("question_str", "")}'
        for q in questions
    )
    user_msg = f"Classify these {len(questions)} questions:\n\n{items}"

    try:
        response = await client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": CATEGORIZE_SYSTEM_PROMPT},
                {"role": "user", "content": user_msg},
            ],
            temperature=0.0,
            max_tokens=1024,
        )
        content = response.choices[0].message.content or ""
        content = re.sub(r"^```json?\s*|\s*```$", "", content.strip())
        classifications: list[dict] = json.loads(content)
    except Exception as e:
        logger.warning("Question categorization failed for session %d: %s", session_id, e)
        return

    # Build lookup: question_num → classification
    class_map: dict[int, dict] = {c["question_num"]: c for c in classifications if "question_num" in c}

    # Validate main_category; fall back to empty string if not in taxonomy
    valid_cats = set(MAIN_CATEGORIES)
    for c in class_map.values():
        if c.get("main_category") not in valid_cats:
            c["main_category"] = ""
        if not isinstance(c.get("tags"), list):
            c["tags"] = []

    # 1. Update ParsedData rows with categories
    category_updates = [
        {
            "question_num": q["question_num"],
            "main_category": class_map.get(q["question_num"], {}).get("main_category", ""),
            "tags": class_map.get(q["question_num"], {}).get("tags", []),
        }
        for q in questions
    ]
    try:
        await update_parsed_data_categories(session_id, category_updates)
    except Exception as e:
        logger.warning("Failed to update ParsedData categories for session %d: %s", session_id, e)

    # 2. Seed shared QuestionBank
    bank_entries = [
        {
            "question_text": q.get("question_text") or q.get("question_str", ""),
            "answer": q.get("answer", ""),
            "main_category": class_map.get(q["question_num"], {}).get("main_category", ""),
            "tags": class_map.get(q["question_num"], {}).get("tags", []),
        }
        for q in questions
        if (q.get("question_text") or q.get("question_str", "")).strip()
    ]
    try:
        await save_question_bank_batch(session_id, teacher_id, bank_entries)
    except Exception as e:
        logger.warning("Failed to seed QuestionBank for session %d: %s", session_id, e)

    logger.info(
        "Categorized %d questions for session %d → %d in bank",
        len(questions), session_id, len(bank_entries),
    )