jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor main.py to include socketio server

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-31 23:50:38 -03:00
commit 4c5c95af48
1 changed files with 20 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

30
src/backend/langflow/main.py
View file

@ -3,11 +3,11 @@ from pathlib import Path
from typing import Optional from typing import Optional
from urllib.parse import urlencode from urllib.parse import urlencode
import socketio
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles from fastapi.staticfiles import StaticFiles
from langflow.api import router from langflow.api import router
from langflow.interface.utils import setup_llm_caching from langflow.interface.utils import setup_llm_caching
from langflow.services.plugins.langfuse_plugin import LangfuseInstance from langflow.services.plugins.langfuse_plugin import LangfuseInstance
@ -15,20 +15,24 @@ from langflow.services.utils import initialize_services, teardown_services
from langflow.utils.logger import configure from langflow.utils.logger import configure
@asynccontextmanager def get_lifespan(fix_migration=False, socketio_server=None):
async def lifespan(app: FastAPI): @asynccontextmanager
initialize_services() async def lifespan(app: FastAPI):
setup_llm_caching() initialize_services(fix_migration=fix_migration, socketio_server=socketio_server)
LangfuseInstance.update() setup_llm_caching()
yield LangfuseInstance.update()
teardown_services() yield
teardown_services()
return lifespan
def create_app(): def create_app():
"""Create the FastAPI app and include the router.""" """Create the FastAPI app and include the router."""
configure() configure()
socketio_server = socketio.AsyncServer(async_mode="asgi", cors_allowed_origins="*", logger=True)
lifespan = get_lifespan(socketio_server=socketio_server)
app = FastAPI(lifespan=lifespan) app = FastAPI(lifespan=lifespan)
origins = ["*"] origins = ["*"]
@ -56,6 +60,13 @@ def create_app():
app.include_router(router) app.include_router(router)
app = mount_socketio(app, socketio_server)
return app
def mount_socketio(app: FastAPI, socketio_server: socketio.AsyncServer):
app.mount("/ws", socketio.ASGIApp(socketio_server))
return app return app
@ -103,7 +114,6 @@ def setup_app(static_files_dir: Optional[Path] = None, backend_only: bool = Fals
if __name__ == "__main__": if __name__ == "__main__":
import uvicorn import uvicorn
from langflow.__main__ import get_number_of_workers from langflow.__main__ import get_number_of_workers
configure() configure()