pdf-quiz-generator/backend/app/services/ai_service.py

import json
import logging

import litellm

from app.config import settings

logger = logging.getLogger(__name__)

EXTRACTION_PROMPT = """You are a quiz question extractor. The following content is from a PDF that contains quiz/exam questions with answers and explanations.

Your job is to EXTRACT (not generate) all the questions, their options, correct answers, and explanations exactly as they appear in the content.

Return a JSON object with a "questions" key containing an array where each object has:
- "question_text": the full question text as it appears
- "question_type": "mcq" if it has multiple choice options, "true_false" if True/False, "fill_blank" if fill-in-the-blank
- "options": array of option strings exactly as written (for mcq), ["True", "False"] (for true_false), or null (for fill_blank)
- "correct_answer": the correct answer exactly as marked in the source
- "explanation": the explanation text if provided, or "" if none
- "page_reference": {page_ref}

Important:
- Extract ALL questions found in the content — do not skip any
- Preserve the original wording exactly
- If a question has an image reference, include "[IMAGE]" in the question_text where the image would be
- If no clear correct answer is marked, use the best answer based on the explanation

Content from page(s) {page_info}:
{content}

Return ONLY valid JSON. No markdown formatting."""


def get_model_for_task(db, task: str = "extraction") -> tuple[str, str | None]:
    """Get the configured model for a specific task from DB, or fall back to settings."""
    try:
        from app.models.ai_model_config import AIModelConfig
        config = db.query(AIModelConfig).filter(
            AIModelConfig.task == task,
            AIModelConfig.is_active == True,
            AIModelConfig.is_default == True,
        ).first()
        if config:
            return config.model_id, config.api_key
    except Exception:
        pass
    return settings.LITELLM_MODEL, settings.LITELLM_API_KEY or None


def _truncate_content(content: str, max_chars: int = 100000) -> str:
    if len(content) <= max_chars:
        return content
    half = max_chars // 2
    return content[:half] + "\n\n... [content truncated] ...\n\n" + content[-half:]


def extract_questions(
    content: str,
    page_info: str = "unknown",
    page_ref: int | None = None,
    model_id: str | None = None,
    api_key: str | None = None,
) -> list[dict]:
    """Extract quiz questions from PDF content using LiteLLM."""
    content = _truncate_content(content)

    prompt = EXTRACTION_PROMPT.format(
        content=content,
        page_info=page_info,
        page_ref=page_ref if page_ref else "null",
    )

    use_model = model_id or settings.LITELLM_MODEL
    use_key = api_key or settings.LITELLM_API_KEY

    last_error = None
    for attempt in range(3):
        try:
            kwargs = {
                "model": use_model,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.1,  # low temp for faithful extraction
            }
            if use_key:
                kwargs["api_key"] = use_key

            # Try JSON mode if supported
            try:
                kwargs["response_format"] = {"type": "json_object"}
                response = litellm.completion(**kwargs)
            except Exception:
                kwargs.pop("response_format", None)
                response = litellm.completion(**kwargs)

            response_text = response.choices[0].message.content

            # Try to parse JSON, handle markdown code blocks
            text = response_text.strip()
            if text.startswith("```"):
                text = text.split("\n", 1)[1] if "\n" in text else text[3:]
                if text.endswith("```"):
                    text = text[:-3]
                text = text.strip()

            data = json.loads(text)
            questions = data.get("questions", data) if isinstance(data, dict) else data
            if not isinstance(questions, list):
                raise ValueError("Response is not a list of questions")

            validated = []
            for q in questions:
                if not all(k in q for k in ("question_text", "correct_answer")):
                    continue
                qtype = q.get("question_type", "mcq")
                if qtype not in ("mcq", "true_false", "fill_blank"):
                    qtype = "mcq"
                validated.append({
                    "question_text": q["question_text"],
                    "question_type": qtype,
                    "options": q.get("options"),
                    "correct_answer": q["correct_answer"],
                    "explanation": q.get("explanation", ""),
                    "page_reference": q.get("page_reference"),
                })

            if validated:
                return validated

            raise ValueError("No valid questions extracted from content")

        except Exception as e:
            last_error = e
            logger.warning(f"Extraction attempt {attempt + 1} failed: {e}")

    raise RuntimeError(f"Failed to extract questions after 3 attempts: {last_error}")


def generate_tts_audio(
    text: str,
    model_id: str | None = None,
    api_key: str | None = None,
) -> bytes | None:
    """Generate TTS audio using LiteLLM (supports Google Vertex, OpenAI, etc.)."""
    use_model = model_id or "vertex_ai/google/cloud-tts"
    use_key = api_key or settings.LITELLM_API_KEY

    try:
        # LiteLLM speech endpoint
        response = litellm.speech(
            model=use_model,
            input=text,
            api_key=use_key if use_key else None,
        )
        # Response is audio bytes
        if hasattr(response, "content"):
            return response.content
        if hasattr(response, "read"):
            return response.read()
        return bytes(response)
    except Exception as e:
        logger.error(f"TTS generation failed: {e}")
        return None