src/generation.py

import json
import os
import re
from dataclasses import dataclass
from pathlib import Path

from llama_cpp import Llama
from llama_index.core import VectorStoreIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

GENERATOR_MODEL = "qwen2.5-3b-instruct-q4_k_m.gguf"
VALID_MODES     = {"flashcard", "summary"}

SYSTEM_PROMPT_FLASHCARD = """You are an expert study assistant that creates structured flash cards for students.
IMPORTANT: Always respond in the same language as the topic query. If they ask in Spanish, answer in Spanish. If they ask in English, answer in English.

Given a context extracted from a textbook and a topic query, generate a flash card in JSON format.
The JSON must follow this exact schema with no additional fields:

{{
    "concept": "name of the concept",
    "definition": "clear and concise definition in 2-3 sentences",
    "key_points": ["point 1", "point 2", "point 3"],
    "examples": ["example 1", "example 2"],
    "suggested_image": "brief description of an image that would illustrate this concept"
}}

Rules:
- Respond ONLY with the JSON object, no preamble, no explanation, no markdown backticks
- Base your response strictly on the provided context
- If the context does not contain enough information, still follow the schema but indicate the limitation in the definition field
- All fields are required"""

SYSTEM_PROMPT_SUMMARY = """You are an expert study assistant that creates structured summaries for students.
IMPORTANT: Always respond in the same language as the topic query. If they ask in Spanish, answer in Spanish. If they ask in English, answer in English.

Given a context extracted from a textbook and a topic query, generate a consolidated summary in JSON format.
The JSON must follow this exact schema with no additional fields:

{{
    "topic": "main topic name",
    "overview": "2-3 sentence overview of the topic",
    "concepts": [
        {{
            "name": "concept name",
            "description": "brief description in 1-2 sentences"
        }}
    ]
}}

Rules:
- Respond ONLY with the JSON object, no preamble, no explanation, no markdown backticks
- Base your response strictly on the provided context
- Include between 3 and 8 concepts
- All fields are required"""

FLASHCARD_EXAMPLE = """{
    "concept": "Amida",
    "definition": "Una amida es un compuesto derivado de un ácido carboxílico donde el grupo hidroxilo es reemplazado por un grupo amino. Las amidas se encuentran en proteínas y muchas moléculas biológicas.",
    "key_points": ["Derivada de ácidos carboxílicos", "Contiene un grupo carbonilo unido a nitrógeno", "Presente en proteínas como enlaces peptídicos"],
    "examples": ["Acetamida (CH3CONH2)", "Nylon (poliamida sintética)"],
    "suggested_image": "Fórmula estructural de una amida mostrando el grupo carbonilo unido a nitrógeno"
}"""

SUMMARY_EXAMPLE = """{
    "topic": "Derivados de ácidos carboxílicos",
    "overview": "Los derivados de ácidos carboxílicos son compuestos que pueden hidrolizarse para dar ácidos carboxílicos.",
    "concepts": [
        {"name": "Éster", "description": "Formado por la reacción de un ácido carboxílico con un alcohol"},
        {"name": "Amida", "description": "Formada por la reacción de un ácido carboxílico con una amina"}
    ]
}"""


@dataclass
class FlashCard:
    concept         : str
    definition      : str
    key_points      : list
    examples        : list
    suggested_image : str


@dataclass
class ConsolidatedSummary:
    topic    : str
    overview : str
    concepts : list


def _truncate_at_complete_json(raw: str) -> str:
    """Walk the string tracking depth and string context so brackets inside
    string values don't fool the truncation point detection."""
    depth         = 0
    in_string     = False
    escape        = False
    last_complete = -1

    for i, ch in enumerate(raw):
        if escape:
            escape = False
            continue
        if ch == "\\" and in_string:
            escape = True
            continue
        if ch == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if ch in "{[":
            depth += 1
        elif ch in "}]":
            depth -= 1
            if depth == 0:
                last_complete = i

    return raw[:last_complete + 1] if last_complete != -1 else raw


def _clean_json_output(raw: str) -> str:
    raw = re.sub(r"```json|```", "", raw).strip()
    raw = re.sub(r",\s*}", "}", raw)   # trailing comma before }
    raw = re.sub(r",\s*]", "]", raw)   # trailing comma before ]
    raw = re.sub(r"\{\s*,", "{", raw)  # leading comma after {
    raw = _truncate_at_complete_json(raw)
    return raw


_llm = None


def load_language_model(model_path: Path, n_ctx: int = 4096) -> Llama:
    global _llm
    if _llm is None:
        _llm = Llama(
            model_path=str(model_path),
            n_ctx=n_ctx,
            n_threads=os.cpu_count(),
            verbose=False,
        )
    return _llm


def generate_flashcard(
    query: str,
    index: VectorStoreIndex,
    llm: Llama,
    embed_model: HuggingFaceEmbedding,
    mode: str = "flashcard",
    k: int = 3,
) -> "FlashCard | ConsolidatedSummary":
    if mode not in VALID_MODES:
        raise ValueError(f"Invalid mode '{mode}'. Must be one of {VALID_MODES}")

    system_prompt = SYSTEM_PROMPT_FLASHCARD if mode == "flashcard" else SYSTEM_PROMPT_SUMMARY
    example       = FLASHCARD_EXAMPLE if mode == "flashcard" else SUMMARY_EXAMPLE
    max_tokens    = 2048 if mode == "flashcard" else 2048

    retriever = index.as_retriever(similarity_top_k=k, embed_model=embed_model)
    nodes     = retriever.retrieve(query)
    if not nodes:
        raise ValueError(f"No chunks retrieved for query: '{query}'")

    context = "\n\n".join(f"[Chunk {i+1}]\n{node.text}" for i, node in enumerate(nodes))

    last_error = None
    for attempt, prompt in enumerate([
        system_prompt,
        system_prompt + f"\n\nExample of expected output:\n{example}",
    ]):
        try:
            response = llm.create_chat_completion(
                messages=[
                    {"role": "system", "content": prompt},
                    {"role": "user",   "content": f"Context:\n{context}\n\nTopic: {query}"},
                ],
                max_tokens=max_tokens,
                temperature=0.1,
            )
            raw = json.loads(_clean_json_output(response["choices"][0]["message"]["content"]))
            if mode == "flashcard":
                return FlashCard(**raw)
            return ConsolidatedSummary(**raw)
        except (json.JSONDecodeError, KeyError, TypeError) as e:
            last_error = e

    raise ValueError(
        f"Generation failed after 2 attempts for query '{query}'. "
        f"Last error: {last_error}"
    )