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

import json
import logging

import litellm

from app.config import settings

logger = logging.getLogger(__name__)


def _proxy_model(model_id: str) -> str:
    """Prefix model with openai/ if using a LiteLLM proxy and no provider is specified."""
    if settings.LITELLM_API_BASE and "/" not in model_id:
        return f"openai/{model_id}"
    return model_id

EXTRACTION_PROMPT = """You are extracting questions from a PREP (Pediatric Review and Education Program) exam PDF.

These PDFs follow a strict format:
1. A numbered question with a clinical vignette (patient scenario)
2. Five answer options labeled A, B, C, D, E
3. A line "Correct Answer: X" where X is the letter of the correct option
4. An explanation paragraph
5. A "Critique:" section with detailed reasoning
6. A "Content Specifications:" section listing the learning objectives

Your task: extract every question and return ONLY a JSON object in this exact format:

{{"questions": [
  {{
    "question_text": "<full question stem including any patient vignette>",
    "question_type": "mcq",
    "options": ["<option A text>", "<option B text>", "<option C text>", "<option D text>", "<option E text>"],
    "correct_answer": "<full text of the correct option, NOT the letter>",
    "explanation": "<explanation paragraph>\\n\\nCritique: <critique section verbatim>\\n\\nContent Specifications: <content spec section verbatim>",
    "page_reference": {page_ref}
  }}
]}}

CRITICAL RULES — the correct_answer field is the most important:
- Find the line "Correct Answer: X" after each question. X is the letter (A, B, C, D, or E).
- Look up the full text of option X from the options list and store it as correct_answer.
- Example: if options are ["alpha","beta","gamma","delta","epsilon"] and Correct Answer is C, then correct_answer = "gamma"
- NEVER store just a letter like "C" — always store the FULL OPTION TEXT.
- If you cannot find "Correct Answer:" for a question, set correct_answer to null (do not guess).
- Do not omit any question — extract ALL questions even if partially cut off.
- Preserve ALL text verbatim — do not summarize or paraphrase anything.
- Return ONLY the JSON — no markdown, no preamble, no explanation.

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


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 = _proxy_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
            if settings.LITELLM_API_BASE:
                kwargs["api_base"] = settings.LITELLM_API_BASE

            # Don't force JSON mode — let the model respond naturally and we parse it
            response = litellm.completion(**kwargs)

            response_text = response.choices[0].message.content
            logger.info(f"Model raw response (first 500 chars): {response_text[:500]!r}")

            # 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)

            # Handle all common response shapes
            if isinstance(data, list):
                questions = data
            elif isinstance(data, dict):
                # Try common keys
                for key in ("questions", "items", "results", "data"):
                    if isinstance(data.get(key), list):
                        questions = data[key]
                        break
                else:
                    # Maybe the dict itself is a single question
                    if "question_text" in data:
                        questions = [data]
                    else:
                        raise ValueError(f"Unexpected response shape: {list(data.keys())}")
            else:
                raise ValueError("Response is not a list of questions")

            validated = []
            skipped = []
            for q in questions:
                if "question_text" not in q:
                    continue
                correct = q.get("correct_answer")
                if not correct:
                    skipped.append(q.get("question_text", "")[:120])
                    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": correct,
                    "explanation": q.get("explanation", ""),
                    "page_reference": q.get("page_reference"),
                    "skipped": [],
                })

            # Attach skipped list to first question so caller can surface it
            if validated and skipped:
                validated[0]["skipped"] = skipped

            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!r}")

    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. Supports OpenAI, ElevenLabs, Google Cloud TTS, and AWS Polly.

    model_id conventions:
      tts-1:alloy          → OpenAI TTS (voice after colon)
      tts-1-hd:nova        → OpenAI TTS HD
      elevenlabs/<voice>   → ElevenLabs
      google/<voice_name>  → Google Cloud TTS (e.g. google/en-US-Wavenet-D)
      polly/<VoiceId>      → AWS Polly Neural (e.g. polly/Joanna)
    """
    import httpx, base64

    use_model = model_id or "tts-1:alloy"

    # ── ElevenLabs ─────────────────────────────────────────────
    if use_model.startswith("elevenlabs/") or use_model.startswith("eleven_labs/"):
        voice = use_model.split("/", 1)[1]
        key = api_key or settings.ELEVENLABS_API_KEY
        if not key:
            logger.error("ElevenLabs API key not configured")
            return None
        try:
            resp = httpx.post(
                f"https://api.elevenlabs.io/v1/text-to-speech/{voice}",
                headers={"xi-api-key": key, "Content-Type": "application/json"},
                json={"text": text, "model_id": "eleven_turbo_v2_5"},
                timeout=30,
            )
            resp.raise_for_status()
            return resp.content
        except Exception as e:
            logger.error(f"ElevenLabs TTS failed: {e}")
            return None

    # ── Google Cloud TTS ────────────────────────────────────────
    if use_model.startswith("google/"):
        voice_name = use_model[len("google/"):]
        key = api_key or settings.GOOGLE_TTS_API_KEY
        if not key:
            logger.error("Google TTS API key not configured (GOOGLE_TTS_API_KEY)")
            return None
        # Parse language code from voice name (e.g. "en-US-Wavenet-D" → "en-US")
        parts = voice_name.split("-")
        lang_code = f"{parts[0]}-{parts[1]}" if len(parts) >= 2 else "en-US"
        try:
            resp = httpx.post(
                f"https://texttospeech.googleapis.com/v1/text:synthesize?key={key}",
                json={
                    "input": {"text": text},
                    "voice": {"languageCode": lang_code, "name": voice_name},
                    "audioConfig": {"audioEncoding": "MP3"},
                },
                timeout=30,
            )
            resp.raise_for_status()
            return base64.b64decode(resp.json()["audioContent"])
        except Exception as e:
            logger.error(f"Google TTS failed: {e}")
            return None

    # ── AWS Polly ───────────────────────────────────────────────
    if use_model.startswith("polly/"):
        voice_id = use_model[len("polly/"):]
        access_key = api_key or settings.AWS_ACCESS_KEY_ID
        secret_key = settings.AWS_SECRET_ACCESS_KEY
        region = settings.AWS_REGION or "us-east-1"
        if not access_key or not secret_key:
            logger.error("AWS credentials not configured (AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY)")
            return None
        try:
            import boto3
            polly = boto3.client(
                "polly",
                aws_access_key_id=access_key,
                aws_secret_access_key=secret_key,
                region_name=region,
            )
            response = polly.synthesize_speech(
                Text=text,
                OutputFormat="mp3",
                VoiceId=voice_id,
                Engine="neural",
            )
            return response["AudioStream"].read()
        except Exception as e:
            logger.error(f"AWS Polly TTS failed: {e}")
            return None

    # ── OpenAI (default) ────────────────────────────────────────
    # model_id may encode voice as "tts-1:nova", "tts-1-hd:alloy", etc.
    clean_model = use_model.replace("openai/", "")
    oai_voice = "alloy"
    if ":" in clean_model:
        clean_model, oai_voice = clean_model.split(":", 1)

    # Per-model key > OPENAI_API_KEY (direct) > LITELLM_API_KEY (proxy)
    if api_key:
        key = api_key
        base = (settings.LITELLM_API_BASE or "https://api.openai.com").rstrip("/")
    elif settings.OPENAI_API_KEY:
        key = settings.OPENAI_API_KEY
        base = "https://api.openai.com"
    else:
        key = settings.LITELLM_API_KEY
        base = (settings.LITELLM_API_BASE or "https://api.openai.com").rstrip("/")

    try:
        resp = httpx.post(
            f"{base}/v1/audio/speech",
            headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
            json={"model": clean_model, "input": text, "voice": oai_voice},
            timeout=60,
        )
        resp.raise_for_status()
        return resp.content
    except Exception as e:
        logger.error(f"OpenAI TTS failed: {e}")
        return None