FastAPI-based working frontend (#10)

playground/streaming_fastapi_worker.py

@@ -0,0 +1,40 @@
+import argparse
+import asyncio
+import time
+from typing import Union
+import json
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import uvicorn
+
+
+app = FastAPI()
+
+
+async def text_streamer(args):
+    context = args["prompt"]
+    words = context.split(" ")
+    for word in words:
+        await asyncio.sleep(1)
+        print("word:", word)
+        ret = {
+            "text": word + " ",
+            "error": 0,
+        }
+        yield (json.dumps(ret) + "\0").encode("utf-8")
+
+
+@app.post("/")
+async def read_root(request: Request):
+    args = await request.json()
+    return StreamingResponse(text_streamer(args))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=10002)
+    args = parser.parse_args()
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")