src/inference.py

"""
Inference module for the Pleias-RAG system.
Orchestrates document retrieval from LanceDB and streamed response generation.
"""

import logging
import os
import re
import time
from typing import Iterator

import lancedb

from src.generation import GenerationEngine

logger = logging.getLogger(__name__)

# A table name maps to a folder under data/ (data/<table>). Restrict it to safe
# characters so a request cannot escape the data directory via path traversal.
_VALID_TABLE = re.compile(r"^[A-Za-z0-9_-]+$")


class PleiasBot:
    """
    Main orchestrator that combines document retrieval (LanceDB) with
    text generation (GenerationEngine) to answer user queries.

    A "table" is simply a folder under ``data/`` holding a LanceDB dataset.
    The shipped ``en``, ``fr`` and ``both`` folders are examples: drop your own
    folder in ``data/`` and reference it by name to query it instead.
    """

    def __init__(
        self,
        table_name: str = "both",
        model_path: str = "models/Pleias-RAG.gguf",
        temperature: float = 0.1,
        max_new_tokens: int = 2048,
        top_p: float = 0.95,
        repetition_penalty: float = 1.0,
        search_limit: int = 3,
        data_dir: str = "data",
    ):
        """
        Initialize the bot with model and database configurations.

        Args:
            table_name: Default table (folder under data/) to search when a
                request does not specify one.
            model_path: Path to the GGUF model file.
            temperature: Sampling temperature for generation.
            max_new_tokens: Maximum tokens to generate.
            top_p: Nucleus sampling probability.
            repetition_penalty: Penalty for repeated tokens.
            search_limit: Maximum number of sources to retrieve per query.
            data_dir: Directory that holds the table folders.
        """
        # Initialize the generation engine
        self.generation_engine = GenerationEngine(
            model_path_or_name=model_path,
            max_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
        )

        self.data_dir = data_dir
        self.default_table = table_name
        self.search_limit = search_limit
        self._tables = {}  # cache of opened LanceDB tables, keyed by name

        # Open the default table now so startup fails fast if it is missing.
        self.get_table(self.default_table)

    def get_table(self, name: str) -> lancedb.table.Table:
        """
        Open (and cache) the LanceDB table stored in ``data/<name>``.

        Prefers the conventional ``crsv`` dataset, but falls back to whatever
        dataset the folder contains, so any user-supplied table just works.

        Args:
            name: Table name, i.e. a folder under ``data_dir``.

        Returns:
            The opened LanceDB table.

        Raises:
            ValueError: If the name is invalid, the folder is missing, or the
                folder contains no LanceDB dataset.
        """
        if not name or not _VALID_TABLE.match(name):
            raise ValueError(
                f"Invalid table name: {name!r} (use letters, digits, '_' or '-')"
            )

        if name not in self._tables:
            path = os.path.join(self.data_dir, name)
            if not os.path.isdir(path):
                raise ValueError(f"Table '{name}' not found under {self.data_dir}/")

            db = lancedb.connect(path)
            available = db.table_names()
            if not available:
                raise ValueError(f"No LanceDB dataset found in {path}")

            inner = "crsv" if "crsv" in available else available[0]
            logger.info(f"Opening table '{name}' ({path}, dataset '{inner}')")
            self._tables[name] = db.open_table(inner)

        return self._tables[name]

    def search(self, text: str, table: lancedb.table.Table, limit: int = 3):
        """
        Perform full-text search on a LanceDB table.

        Args:
            text: The query text to search for.
            table: The LanceDB table to search in.
            limit: Maximum number of results to return.

        Returns:
            List of source dictionaries with keys:
                - "id": 1-based index
                - "text": The source content
                - "metadata": All other fields from the database
        """
        logger.info("Searching for text")
        start = time.time()

        results = (
            table.search(text, query_type="fts")
            .limit(limit)
            .to_pandas()
            .T.to_dict()
        )

        logger.info(f"Search time: {time.time() - start:.2f} seconds")

        # Reformat results into the expected structure
        sources = []
        for idx, key in enumerate(results.keys(), 1):
            sources.append(
                {
                    "id": idx,
                    "text": results[key]["text"],
                    "metadata": {
                        subkey: results[key][subkey]
                        for subkey in results[key].keys()
                        if subkey != "text"
                    },
                }
            )

        return sources

    def stream(self, user_message: str, table: str = None) -> Iterator[str]:
        """
        Retrieve relevant sources and stream the model's raw output.

        Args:
            user_message: The user's question.
            table: Table (folder under data/) to search. Defaults to the bot's
                default table when not given.

        Yields:
            Raw text pieces of the model output, in generation order.
        """
        # Step 1: Retrieve relevant sources from the requested table
        tbl = self.get_table(table or self.default_table)
        sources = self.search(user_message, table=tbl, limit=self.search_limit)

        # Step 2: Stream the generated response
        logger.info("Streaming response from model...")
        yield from self.generation_engine.stream_generate(user_message, sources)

    def stream_raw(self, user_message: str, sources: list) -> Iterator[str]:
        """
        Stream the model's raw output for directly-supplied sources (no retrieval).

        Intended for testing: the caller provides the query and sources, which
        are formatted into the prompt exactly as retrieved sources would be.

        Args:
            user_message: The user's question.
            sources: List of source dicts, each with a "text" key.

        Yields:
            Raw text pieces of the model output, in generation order.
        """
        logger.info(f"Streaming response from model (raw, {len(sources)} sources)...")
        yield from self.generation_engine.stream_generate(user_message, sources)