FastAPI-based working frontend (#10)

playground/http_client.py

@@ -0,0 +1,20 @@
+import requests
+import json
+
+def http_bot():
+    prompt = "How are you? I'm fine."
+
+    headers = {"User-Agent": "Test Client"}
+    pload = {
+        "prompt": prompt,
+    }
+    response = requests.post("http://localhost:10002", headers=headers, json=pload, stream=True)
+
+    for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"]
+            yield output
+
+for h in http_bot():
+    print(h, end="", flush=True)

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