generator.py

"""
generator.py — Qwen2.5-3B-Instruct-Q4_K_M generation singleton.

Timeout fix
-----------
"signal timed out" is a SIGALRM raised by llama-cpp's internal C-level
watchdog. Two facts from the server logs drove this design:

  1. signal.signal() can ONLY be called from the main thread. FastAPI runs
     /query in a threadpool worker, so calling _suppress_sigalrm() inside
     generate() raises:
       ValueError: signal only works in main thread of the main interpreter

  2. Therefore SIGALRM is suppressed ONCE in load(), which is called from
     the lifespan startup (main thread). This is enough because llama-cpp
     sets the handler at model-load time; we override it immediately after.

  3. The ThreadPoolExecutor wrapper in generate() provides a hard 110s
     wall-clock timeout as a safety net, catching any runaway inference and
     returning a graceful message instead of a 500 crash.

Timing budget on Qwen2.5-3B-Q4 @ ~8 tok/s:
  prefill  ~2s  (800 prompt tokens)
  generate ~50s (400 output tokens)
  total    ~52s — well under the 110s thread timeout
"""
from __future__ import annotations

import concurrent.futures
import logging
import signal
import threading
from typing import Iterator, List, Optional

from config import (
    GGUF_MODEL_PATH,
    LLM_CONTEXT_LENGTH,
    LLM_MAX_TOKENS,
    LLM_N_GPU_LAYERS,
    LLM_N_THREADS,
    LLM_TEMPERATURE,
    LLM_TOP_P,
)
from models import RetrievedChunk

logger = logging.getLogger(__name__)

_SYSTEM_PROMPT = """\
You are a precise document question-answering assistant.
You will be given numbered context passages extracted from a document.
Your job is to answer the question using ONLY information present in these passages.

Rules:
- Read EVERY passage completely before forming your answer — the answer may appear in any passage, including later ones.
- If ANY passage contains the answer or part of the answer, state it directly. Cite the passage number and page.
- If multiple passages are needed for a complete answer, combine them.
- NEVER say the context does not contain information if you have not fully read all passages. If you find the answer, state it.
- Do not fabricate numbers, names, or facts not present in the passages.
- When a passage says X performs well but a LATER sentence in the SAME passage qualifies or contradicts this (e.g. "however Y performs better"), include the qualification in your answer.
- Keep answers concise and factual. A one-word or one-number answer is almost always incomplete — include units and context."""

# Hard cap on chars per chunk in the prompt.
# Increased from 600→800 to avoid truncating the exact sentence containing
# the answer. At 150-word chunks (~900 chars avg), 600 was cutting mid-chunk.
_MAX_CHUNK_CHARS = 800

# Wall-clock timeout for the ThreadPoolExecutor safety net (seconds).
# Must be under the proxy/HTTP timeout (usually 120s).
_INFERENCE_TIMEOUT_S = 110


def _build_context(chunks: List[RetrievedChunk]) -> str:
    """
    Build the context block for the LLM prompt.

    Space 1 changes reflected here:
      - chunk.text is now CLEAN (no [Context:] prefix) — FIX #1. The manual
        stripping that was here is removed; it was a workaround that is no
        longer needed and would corrupt legitimate text starting with "[".
      - The prompt header now shows the full outline breadcrumb
        (e.g. "Methods > Data Collection") instead of just section_title.
        This gives the LLM richer location context for attribution.
      - chunk.text no longer starts with a duplicate section title so the
        duplicate-stripping guard is also removed.
    """
    ordered = sorted(chunks, key=lambda c: (c.doc_name, c.page_num))
    parts = []
    for i, chunk in enumerate(ordered, 1):
        # Use outline_path breadcrumb if available, fall back to section_title.
        outline = getattr(chunk, "outline_path", None) or []
        if outline:
            location = " > ".join(outline)
        elif chunk.section_title:
            location = chunk.section_title
        else:
            location = ""

        header = (
            f"[{i}] {chunk.doc_name} | p.{chunk.page_num}"
            + (f" | {location}" if location else "")
            + (f" | {chunk.region_type}" if chunk.region_type != "paragraph" else "")
        )

        if chunk.table_html:
            # Tables: use structured HTML so the LLM can read rows/columns properly.
            text = chunk.table_html
        else:
            # chunk.text is clean (Space 1 FIX #1) — no stripping needed.
            text = chunk.text or "(no text content)"
            if len(text) > _MAX_CHUNK_CHARS:
                text = text[:_MAX_CHUNK_CHARS] + "…"

        parts.append(f"{header}\n{text}")
    return "\n\n".join(parts)


def _build_messages(question: str, chunks: List[RetrievedChunk]) -> List[dict]:
    context = _build_context(chunks)
    user_content = f"Context passages:\n{context}\n\nQuestion: {question}"
    return [
        {"role": "system", "content": _SYSTEM_PROMPT},
        {"role": "user",   "content": user_content},
    ]


class Generator:
    def __init__(self):
        self._llm      = None
        self._ready    = False
        self._executor: Optional[concurrent.futures.ThreadPoolExecutor] = None

    def load(self) -> None:
        """
        Called once from the lifespan startup — always runs in the main thread.
        This is the only safe place to call signal.signal().
        """
        if self._ready:
            return
        try:
            from llama_cpp import Llama
        except ImportError:
            raise RuntimeError("llama-cpp-python is not installed.")
        import os
        if not os.path.isfile(GGUF_MODEL_PATH):
            raise FileNotFoundError(
                f"GGUF model not found at {GGUF_MODEL_PATH}. "
                "Run download_models.py first."
            )
        logger.info("Loading GGUF model from %s …", GGUF_MODEL_PATH)
        self._llm = Llama(
            model_path=GGUF_MODEL_PATH,
            n_ctx=LLM_CONTEXT_LENGTH,
            n_threads=LLM_N_THREADS,
            n_gpu_layers=LLM_N_GPU_LAYERS,
            verbose=False,
        )

        # Suppress SIGALRM here — in the main thread, right after model load.
        # signal.signal() raises ValueError in worker threads, so this is the
        # only safe call site. llama-cpp sets the alarm at load time; we
        # override it immediately after.
        if threading.current_thread() is threading.main_thread():
            try:
                signal.signal(signal.SIGALRM, signal.SIG_IGN)
                logger.info("SIGALRM watchdog disabled (main thread)")
            except (OSError, AttributeError):
                # SIGALRM not available on Windows — safe to ignore.
                pass
        else:
            logger.warning(
                "load() called outside main thread — SIGALRM not suppressed. "
                "The ThreadPoolExecutor timeout will still protect /query."
            )

        self._ready = True
        # Create the executor once here (main thread) rather than per generate()
        # call. ThreadPoolExecutor construction involves OS-level thread spawning;
        # doing it on every inference call wastes 5-30ms per request.
        if self._executor is not None:
            self._executor.shutdown(wait=False)
        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
        logger.info("Generator ready (Qwen2.5-3B-Instruct-Q4_K_M)")

    def generate(self, question: str, chunks: List[RetrievedChunk]) -> str:
        """
        Blocking generation — called from a FastAPI threadpool worker.
        DO NOT call signal.signal() here; it raises ValueError in worker threads.
        The ThreadPoolExecutor provides a hard wall-clock timeout instead.
        """
        if not self._ready:
            raise RuntimeError("Generator.load() has not been called")
        if not chunks:
            return "No relevant passages were found in the document for that question."

        messages = _build_messages(question, chunks)

        def _run():
            return self._llm.create_chat_completion(
                messages=messages,
                max_tokens=LLM_MAX_TOKENS,
                temperature=LLM_TEMPERATURE,
                top_p=LLM_TOP_P,
                stop=["<|im_end|>", "<|endoftext|>"],
            )

        # Reuse the executor created in load() — avoids thread-spawn overhead
        # (~5-30ms) that occurred when creating a new ThreadPoolExecutor per call.
        future = self._executor.submit(_run)
        try:
            response = future.result(timeout=_INFERENCE_TIMEOUT_S)
        except concurrent.futures.TimeoutError:
            logger.warning("Inference timed out after %ds", _INFERENCE_TIMEOUT_S)
            return (
                "Answer generation timed out. The document may be too large "
                "for a single query. Try asking a more specific question."
            )
        except Exception as e:
            logger.error("Inference error: %s", e)
            return f"An error occurred during generation: {e}"

        answer = response["choices"][0]["message"]["content"].strip()
        logger.debug("Generated answer (%d chars)", len(answer))
        return answer

    def generate_stream(self, question: str, chunks: List[RetrievedChunk]) -> Iterator[str]:
        """
        Streaming generation — also runs in a worker thread via run_in_executor.
        Same rule: no signal.signal() calls here.
        """
        if not self._ready:
            yield "Generator not ready."
            return
        if not chunks:
            yield "No relevant passages were found in the document for that question."
            return

        messages = _build_messages(question, chunks)
        try:
            stream = self._llm.create_chat_completion(
                messages=messages,
                max_tokens=LLM_MAX_TOKENS,
                temperature=LLM_TEMPERATURE,
                top_p=LLM_TOP_P,
                stop=["<|im_end|>", "<|endoftext|>"],
                stream=True,
            )
            for chunk in stream:
                delta = chunk["choices"][0].get("delta", {})
                text = delta.get("content", "")
                if text:
                    yield text
        except Exception as e:
            logger.error("Stream inference error: %s", e)
            yield f"\n\n[Generation error: {e}]"

    @property
    def is_ready(self) -> bool:
        return self._ready


generator = Generator()