118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
import os
|
|
import sys
|
|
import uvicorn
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
|
|
from fastapi.responses import HTMLResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.templating import Jinja2Templates
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from dotenv import load_dotenv
|
|
|
|
from src.websocket.websocket_manager import ConnectionManager
|
|
from src.users.router import router as user_router
|
|
from src.files.router import router as file_router
|
|
from src.bot.router import router as bot_router
|
|
from src.etl.router import router as etl_router
|
|
from src.login.router import router as login_router
|
|
from src.conversation.router import router as conversation_router
|
|
|
|
from src.chat.index import chat_manager
|
|
from src.chat.callback import StreamingLLMCallbackHandler
|
|
from src.chat.schemas import ChatResponse
|
|
|
|
load_dotenv()
|
|
|
|
tags_metadata = [
|
|
{
|
|
"name": "users",
|
|
"description": "Operations with users.",
|
|
},
|
|
{
|
|
"name": "files",
|
|
"description": "Operations with file handling.",
|
|
},
|
|
{
|
|
"name": "bot",
|
|
"description": "Operations with bot.",
|
|
},
|
|
{
|
|
"name": "etl",
|
|
"description": "Operations with ETL process.",
|
|
},
|
|
{
|
|
"name": "conversations",
|
|
"description": "Operations with conversations process.",
|
|
},
|
|
|
|
]
|
|
|
|
|
|
app = FastAPI(openapi_tags=tags_metadata)
|
|
|
|
|
|
app.mount("/static", StaticFiles(directory="templates"), name="static")
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=False,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
app.include_router(user_router)
|
|
app.include_router(file_router)
|
|
app.include_router(bot_router)
|
|
app.include_router(etl_router)
|
|
app.include_router(login_router)
|
|
app.include_router(conversation_router)
|
|
|
|
|
|
@app.get("/",response_class=HTMLResponse)
|
|
async def read_index(request: Request):
|
|
return templates.TemplateResponse("index.html", {"request": request})
|
|
|
|
|
|
@app.get("/knowledge-base",response_class=HTMLResponse, tags=["ui"])
|
|
async def read_index(request: Request):
|
|
return templates.TemplateResponse("knowledge-base.html", {"request": request})
|
|
|
|
|
|
@app.get("/websocket-tester",response_class=HTMLResponse, tags=["tester"], summary="UI for testing websocket")
|
|
async def read_index(request: Request):
|
|
return templates.TemplateResponse("websocket-tester.html", {"request": request})
|
|
|
|
@app.get("/upload-tester",response_class=HTMLResponse, tags=["tester"], summary="UI for testing file upload")
|
|
async def read_index(request: Request):
|
|
return templates.TemplateResponse("upload-tester.html", {"request": request})
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
@app.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket):
|
|
await manager.connect(websocket)
|
|
|
|
while True:
|
|
try:
|
|
data = await websocket.receive_json()
|
|
q = data["message"]
|
|
|
|
stream_handler = StreamingLLMCallbackHandler(websocket)
|
|
await websocket.send_text(".-botResponseStartStreaming")
|
|
response = await chat_manager.chat_response({
|
|
"question": q,
|
|
"chat_session": data["chat_session"],
|
|
"bot_id": data["bot_id"],
|
|
"stream_handler": stream_handler,
|
|
})
|
|
|
|
await websocket.send_text(".-botResponseEndStreaming")
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket)
|
|
# await websocket.send_text(f"Connection lost")
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host='0.0.0.0', port=int(os.environ.get('PORT', 8000)), reload=True) |