scripts/serve.py

from __future__ import annotations

import argparse
from pathlib import Path

from flask import Flask, jsonify, request

from ask import ask


def create_app(index_dir: Path, base_model: str, adapter: 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,
                base_model=base_model,
                adapter=adapter,
                k=int(payload.get("k", k)),
                max_new_tokens=int(payload.get("max_new_tokens", 320)),
            )
            return jsonify({"answer": answer})
        except Exception as exc:
            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=8080)
    parser.add_argument("--index-dir", type=Path, default=Path("data/index"))
    parser.add_argument("--base-model", default="Qwen/Qwen2.5-Coder-0.5B-Instruct")
    parser.add_argument("--adapter", default="")
    parser.add_argument("--k", type=int, default=4)
    args = parser.parse_args()
    app = create_app(args.index_dir, args.base_model, args.adapter, args.k)
    app.run(host=args.host, port=args.port)


if __name__ == "__main__":
    main()