serve.py

"""Minimal HTTP server for the Edge Impulse docs RAG assistant.

    python serve.py --host 0.0.0.0 --port 8000

POST /ask   {"question": "...", "k": 4}      -> {"answer": "..."}
GET  /health                                 -> {"ok": true}
"""
from __future__ import annotations

import argparse
from pathlib import Path

from flask import Flask, jsonify, request

from rag import DEFAULT_API_BASE, DEFAULT_INDEX_DIR, DEFAULT_MODEL, ask


def create_app(index_dir: Path, api_base: str, model: str, k: int) -> Flask:
    app = Flask(__name__)

    @app.get("/health")
    def health():
        return jsonify({"ok": True})

    @app.post("/ask")
    def ask_route():
        payload = request.get_json(silent=True) or {}
        question = str(payload.get("question", "")).strip()
        if not question:
            return jsonify({"error": "question is required"}), 400
        try:
            answer = ask(
                question,
                index_dir=index_dir,
                k=int(payload.get("k", k)),
                max_new_tokens=int(payload.get("max_new_tokens", 320)),
                api_base=api_base,
                model=model,
            )
            return jsonify({"answer": answer})
        except Exception as exc:  # noqa: BLE001 - surface the error to the client
            return jsonify({"error": str(exc)}), 500

    return app


def main() -> None:
    parser = argparse.ArgumentParser(description="Serve the Edge Impulse docs RAG assistant.")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument("--index-dir", type=Path, default=DEFAULT_INDEX_DIR)
    parser.add_argument("--api-base", default=DEFAULT_API_BASE)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--k", type=int, default=4)
    args = parser.parse_args()
    app = create_app(args.index_dir, args.api_base, args.model, args.k)
    app.run(host=args.host, port=args.port)


if __name__ == "__main__":
    main()