Make API docs better

megabots/api.py

@@ -1,14 +1,19 @@
 import os
 from megabots import bot
-from lcserve import serving
+from megabots.utils import create_api
+
+# from lcserve import serving
 
 cur_dir = os.path.dirname(os.path.abspath(__file__))
 index_dir = os.path.join(cur_dir, "..", "examples", "files")
 
-print("hey")
+
+mybot = bot("qna-over-docs", index="./index.pkl")
+
+
+# @serving
+# def ask(question: str) -> str:
+#     return mybot.ask(question)
 
 
-@serving
-def ask(question: str) -> str:
-    mybot = bot("qna-over-docs", index=index_dir)
-    return mybot.ask(question)
+app = create_api(mybot)

megabots/utils.py

@@ -1,15 +1,43 @@
 import gradio as gr
 from fastapi import FastAPI
 from megabots.bot import Bot
+from fastapi.openapi.utils import get_openapi
+from pydantic import BaseModel
 
 
-def create_api(bot: Bot):
+def _custom_openapi(app: FastAPI, version: str):
+    if app.openapi_schema:
+        return app.openapi_schema
+
+    openapi_schema = get_openapi(
+        title="🤖 Megabots API",
+        version=version,
+        description="Use this API to interact with the bot.",
+        routes=app.routes,
+    )
+    return openapi_schema
+
+
+class Answer(BaseModel):
+    text: str
+
+
+def create_api(bot: Bot, version: str = "0.0.1"):
     app = FastAPI()
 
-    @app.get("/v1/ask/{question}")
-    async def ask(question: str):
+    @app.get(
+        "/v1/bot/ask/{question}",
+        tags=["Bot"],
+        summary="Ask bot",
+        description="Send question to the bot.",
+        responses={200: {"description": "Bot answer"}},
+        response_model=Answer,
+    )
+    async def ask(question: str) -> Answer:
         answer = bot.ask(question)
-        return {"answer": answer}
+        return Answer(text=answer)
+
+    app.openapi_schema = _custom_openapi(app, version)
 
     return app