""" 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()