src/main.py

"""
Flask API server for the Pleias-RAG system.
Exposes a single /ask_stream endpoint that streams the model's raw output.
"""

import argparse
import logging

from flask import Flask, Response, jsonify, request, stream_with_context

import src.inference as inference

bot = None
app = Flask(__name__)


def configure_logging(debug: bool = False):
    """
    Set up logging configuration for the application.
    Suppresses verbose output from llama_cpp and werkzeug.
    """
    level = logging.DEBUG if debug else logging.INFO
    logging.basicConfig(
        level=level,
        format="%(asctime)s %(name)s %(levelname)s: %(message)s",
        handlers=[logging.StreamHandler()],
        force=True,
    )
    logging.getLogger("llama_cpp").setLevel(logging.WARNING)
    logging.getLogger("werkzeug").setLevel(logging.WARNING)


@app.route("/ask_stream", methods=["POST"])
def handle_ask_stream():
    """
    Streaming endpoint: streams the model's raw output token-by-token as plain text.

    Expects JSON payload:
        - "query" (required): The user's question
        - "table" (optional): Which table (folder under data/) to search.
          Defaults to the server's table. Shipped examples: "en", "fr", "both".

    Example:
        {"query": "What is CRSV?", "table": "en"}

    Returns:
        A text/plain stream of the raw model output.
    """
    if not request.is_json:
        return jsonify({"error": "Request must be JSON"}), 400

    data = request.get_json()
    user_query = data.get("query")
    table = data.get("table")  # Optional; defaults to the server's table

    if not user_query:
        return jsonify({"error": "Missing 'query' key in JSON payload"}), 400

    # Validate/open the requested table up front so we can return a proper error
    # status (once streaming starts the response is already committed to 200).
    try:
        bot.get_table(table or bot.default_table)
    except ValueError as e:
        return jsonify({"error": str(e)}), 400

    app.logger.info(
        f"Received ask_stream request for: '{user_query}' (table={table or bot.default_table})"
    )

    def generate():
        try:
            yield from bot.stream(user_query, table=table)
        except Exception as e:
            app.logger.error(f"Error during stream generation: {e}", exc_info=True)

    return Response(
        stream_with_context(generate()),
        mimetype="text/plain",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


@app.route("/raw_query", methods=["POST"])
def handle_raw_query():
    """
    Testing endpoint: stream the model's raw output for a query and sources
    provided directly in the request, with NO retrieval.

    The sources are formatted into the prompt exactly as retrieved sources are,
    so this is a way to probe the model on controlled inputs.

    Expects JSON payload:
        - "query" (required): The user's question
        - "sources" (optional): List of sources, each either a string or an
          object with a "text" field. Defaults to an empty list (no sources).
        - "show_prompt" (optional): If true, the exact formatted prompt sent to
          the model is streamed first, delimited, before the output.

    Example:
        {"query": "What is CRSV?", "sources": ["CRSV means ...", "It is ..."]}

    Returns:
        A text/plain stream of the raw model output.
    """
    if not request.is_json:
        return jsonify({"error": "Request must be JSON"}), 400

    data = request.get_json()
    user_query = data.get("query")
    raw_sources = data.get("sources", [])
    show_prompt = data.get("show_prompt", False)
    if isinstance(show_prompt, str):
        show_prompt = show_prompt.strip().lower() in ("1", "true", "yes")

    if not user_query:
        return jsonify({"error": "Missing 'query' key in JSON payload"}), 400
    if not isinstance(raw_sources, list):
        return jsonify({"error": "'sources' must be a list"}), 400

    # Normalize each source to the {"text": ...} shape the formatter expects.
    # Accept plain strings or objects with a "text" field.
    sources = []
    for item in raw_sources:
        if isinstance(item, str):
            sources.append({"text": item})
        elif isinstance(item, dict) and "text" in item:
            sources.append({"text": item["text"]})
        else:
            return jsonify(
                {"error": "each source must be a string or an object with a 'text' field"}
            ), 400

    app.logger.info(
        f"Received raw_query request for: '{user_query}' ({len(sources)} sources, no retrieval)"
    )

    def generate():
        try:
            if show_prompt:
                yield "========== PROMPT ==========\n"
                yield bot.generation_engine.format_prompt(user_query, sources)
                yield "\n========== OUTPUT ==========\n"
            yield from bot.stream_raw(user_query, sources)
        except Exception as e:
            app.logger.error(f"Error during raw stream generation: {e}", exc_info=True)

    return Response(
        stream_with_context(generate()),
        mimetype="text/plain",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


def main():
    """
    Entry point: parse arguments, configure logging, load model, and start server.
    """
    global bot

    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--dataset", dest="dataset", default="data",
                        help="Dataset directory holding the table folders (default: 'data'). "
                             "Point this at your own dataset to serve your own tables.")
    parser.add_argument("-t", "--table-name", dest="table_name", default="both",
                        help="Default table (a folder inside the dataset) to search when a "
                             "request omits one: e.g. 'en', 'fr', 'both', or your own")
    parser.add_argument("--debug", action="store_true",
                        help="Enable debug logging")
    parser.add_argument("--host", default="0.0.0.0",
                        help="Host to bind the server to")
    parser.add_argument("-p", "--port", type=int, dest="port", default=8081,
                        help="Port to run the server on")
    args = parser.parse_args()

    configure_logging(args.debug)
    app.logger.info("Starting up Pleias-RAG API server...")
    app.logger.info(f"Loading model | dataset: '{args.dataset}' | default table: '{args.table_name}'...")
    bot = inference.PleiasBot(args.table_name, data_dir=args.dataset)
    app.logger.info("Model loaded successfully. Ready for requests.")
    app.logger.info("Endpoints:")
    app.logger.info("  POST /ask_stream - retrieve from a table, then stream the raw output")
    app.logger.info("  POST /raw_query  - stream the raw output for query + sources you supply")
    app.run(host=args.host, port=args.port, debug=args.debug)


if __name__ == "__main__":
    main()