32 lines
915 B
Python
32 lines
915 B
Python
from fastapi import FastAPI, WebSocket
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi import Request
|
|
from fastapi.staticfiles import StaticFiles
|
|
|
|
from chain import full_chain
|
|
|
|
app = FastAPI()
|
|
|
|
templates = Jinja2Templates(directory="templates")
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
|
|
|
@app.get("/")
|
|
def home(request: Request):
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
|
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
await websocket.accept()
|
|
|
|
try:
|
|
while True:
|
|
message = await websocket.receive_text()
|
|
await websocket.send_json({"type": "start"})
|
|
async for chunk in full_chain.astream(message):
|
|
await websocket.send_json({"type": "chunk", "data": chunk})
|
|
await websocket.send_json({"type": "end"})
|
|
except Exception:
|
|
pass
|