From 15ad1226ea2b5574b8e26d785f94584508d537b8 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 6 Dec 2023 20:00:36 -0300
Subject: [PATCH] Add asynccontextmanager to lifespan function

---
 src/backend/langflow/main.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/backend/langflow/main.py b/src/backend/langflow/main.py
index 2df93f20b..b4acdedf9 100644
--- a/src/backend/langflow/main.py
+++ b/src/backend/langflow/main.py
@@ -1,3 +1,4 @@
+from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Optional
 from urllib.parse import urlencode
@@ -13,12 +14,20 @@ from langflow.services.utils import initialize_services, teardown_services
 from langflow.utils.logger import configure
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    initialize_services()
+    setup_llm_caching()
+    LangfuseInstance.update()
+    yield
+    teardown_services()
+
 def create_app():
     """Create the FastAPI app and include the router."""
 
     configure()
 
-    app = FastAPI()
+    app = FastAPI(lifespan=lifespan)
     origins = ["*"]
 
     app.add_middleware(
@@ -43,14 +52,9 @@ def create_app():
     def health():
         return {"status": "ok"}
 
+
     app.include_router(router)
 
-    app.on_event("startup")(initialize_services)
-    app.on_event("startup")(setup_llm_caching)
-    app.on_event("startup")(LangfuseInstance.update)
-
-    app.on_event("shutdown")(teardown_services)
-    app.on_event("shutdown")(LangfuseInstance.teardown)
 
     return app