From 30e19389a5d028ff8a0ebec6d8dea6ef8cf4818c Mon Sep 17 00:00:00 2001
From: Merwane Hamadi <merwanehamadi@gmail.com>
Date: Fri, 18 Aug 2023 19:39:03 -0700
Subject: [PATCH] Fix pool size issue on DB

---
 autogpt/db.py | 44 +++++++++++++++++++++-----------------------
 1 file changed, 21 insertions(+), 23 deletions(-)

diff --git a/autogpt/db.py b/autogpt/db.py
index c7bed99..7899fe3 100644
--- a/autogpt/db.py
+++ b/autogpt/db.py
@@ -101,18 +101,16 @@ def __init__(self, database_string) -> None:
         self.Session = sessionmaker(bind=self.engine)
         print("Databases Created")
 
-    async def create_task(
-        self, input: Optional[str], additional_input: Optional[TaskInput] = None
-    ) -> Task:
-        session = self.Session()
-        new_task = TaskModel(
-            input=input,
-            additional_input=additional_input.__root__ if additional_input else None,
-        )
-        session.add(new_task)
-        session.commit()
-        session.refresh(new_task)
-        return convert_to_task(new_task)
+    async def create_task(self, input: Optional[str], additional_input: Optional[TaskInput] = None) -> Task:
+        with self.Session() as session:
+            new_task = TaskModel(
+                input=input,
+                additional_input=additional_input.__root__ if additional_input else None,
+            )
+            session.add(new_task)
+            session.commit()
+            session.refresh(new_task)
+            return convert_to_task(new_task)
 
     async def create_step(
         self,
@@ -265,14 +263,14 @@ async def list_steps(self, task_id: str) -> List[Step]:
         ]
 
     async def list_artifacts(self, task_id: str) -> List[Artifact]:
-        session = self.Session()
-        artifacts = session.query(ArtifactModel).filter_by(task_id=task_id).all()
-        return [
-            Artifact(
-                artifact_id=str(artifact.artifact_id),
-                file_name=artifact.file_name,
-                agent_created=artifact.agent_created,
-                uri=artifact.uri,
-            )
-            for artifact in artifacts
-        ]
+        with self.Session() as session:
+            artifacts = session.query(ArtifactModel).filter_by(task_id=task_id).all()
+            return [
+                Artifact(
+                    artifact_id=str(artifact.artifact_id),
+                    file_name=artifact.file_name,
+                    agent_created=artifact.agent_created,
+                    uri=artifact.uri,
+                )
+                for artifact in artifacts
+            ]