From 8f3ed733b9c6ef2d04cd91d941d1aa37640a25f1 Mon Sep 17 00:00:00 2001 From: Nicholas Tindle Date: Fri, 14 Jun 2024 06:56:53 -0500 Subject: [PATCH] chore(rnd): Fix warning during server tests (#7213) --- .../autogpt_server/server/server.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/rnd/autogpt_server/autogpt_server/server/server.py b/rnd/autogpt_server/autogpt_server/server/server.py index e877a9b38c..53a24db384 100644 --- a/rnd/autogpt_server/autogpt_server/server/server.py +++ b/rnd/autogpt_server/autogpt_server/server/server.py @@ -1,4 +1,6 @@ import asyncio +from contextlib import asynccontextmanager + import uuid import uvicorn @@ -8,7 +10,15 @@ from autogpt_server.data import db, execution, graph class AgentServer: + def __init__(self, queue: execution.ExecutionQueue): + + @asynccontextmanager + async def lifespan(app: FastAPI): + await db.connect() + yield + await db.disconnect() + self.app = FastAPI( title="AutoGPT Agent Server", description=( @@ -17,6 +27,7 @@ class AgentServer: ), summary="AutoGPT Agent Server", version="0.1", + lifespan=lifespan, ) self.execution_queue = queue @@ -28,8 +39,6 @@ class AgentServer: methods=["POST"], ) self.app.include_router(self.router) - self.app.on_event("startup")(db.connect) - self.app.on_event("shutdown")(db.disconnect) async def execute_agent(self, agent_id: str, node_input: dict): agent = await graph.get_graph(agent_id) @@ -49,9 +58,7 @@ class AgentServer: ) task = execution.add_execution( - execution.Execution( - run_id=run_id, node_id=node.id, data=node_input - ), + execution.Execution(run_id=run_id, node_id=node.id, data=node_input), self.execution_queue, )