feat(backend): Introduce executors shared DB connection

autogpt_platform/backend/backend/app.py

@@ -25,10 +25,12 @@ def main(**kwargs):
     """
 
     from backend.executor import ExecutionManager, ExecutionScheduler
+    from backend.server.db_api import DatabaseAPI
     from backend.server.rest_api import AgentServer
     from backend.server.ws_api import WebsocketServer
 
     run_processes(
+        DatabaseAPI(),
         ExecutionManager(),
         ExecutionScheduler(),
         WebsocketServer(),

autogpt_platform/backend/backend/data/credit.py

@@ -18,7 +18,7 @@
     LlmModel,
 )
 from backend.blocks.talking_head import CreateTalkingAvatarVideoBlock
-from backend.data.block import Block, BlockInput
+from backend.data.block import Block, BlockInput, get_block
 from backend.util.settings import Config
 
 
@@ -96,7 +96,7 @@ async def spend_credits(
         self,
         user_id: str,
         user_credit: int,
-        block: Block,
+        block_id: str,
         input_data: BlockInput,
         data_size: float,
         run_time: float,
@@ -107,7 +107,7 @@ async def spend_credits(
         Args:
             user_id (str): The user ID.
             user_credit (int): The current credit for the user.
-            block (Block): The block that is being used.
+            block_id (str): The block ID.
             input_data (BlockInput): The input data for the block.
             data_size (float): The size of the data being processed.
             run_time (float): The time taken to run the block.
@@ -208,12 +208,16 @@ async def spend_credits(
         self,
         user_id: str,
         user_credit: int,
-        block: Block,
+        block_id: str,
         input_data: BlockInput,
         data_size: float,
         run_time: float,
         validate_balance: bool = True,
     ) -> int:
+        block = get_block(block_id)
+        if not block:
+            raise ValueError(f"Block not found: {block_id}")
+
         cost, matching_filter = self._block_usage_cost(
             block=block, input_data=input_data, data_size=data_size, run_time=run_time
         )

autogpt_platform/backend/backend/data/graph.py

@@ -7,7 +7,7 @@
 import prisma.types
 from prisma.models import AgentGraph, AgentGraphExecution, AgentNode, AgentNodeLink
 from prisma.types import AgentGraphInclude
-from pydantic import BaseModel, PrivateAttr
+from pydantic import BaseModel
 from pydantic_core import PydanticUndefinedType
 
 from backend.blocks.basic import AgentInputBlock, AgentOutputBlock
@@ -51,17 +51,8 @@ class Node(BaseDbModel):
     block_id: str
     input_default: BlockInput = {}  # dict[input_name, default_value]
     metadata: dict[str, Any] = {}
-
-    _input_links: list[Link] = PrivateAttr(default=[])
-    _output_links: list[Link] = PrivateAttr(default=[])
-
-    @property
-    def input_links(self) -> list[Link]:
-        return self._input_links
-
-    @property
-    def output_links(self) -> list[Link]:
-        return self._output_links
+    input_links: list[Link] = []
+    output_links: list[Link] = []
 
     @staticmethod
     def from_db(node: AgentNode):
@@ -73,8 +64,8 @@ def from_db(node: AgentNode):
             input_default=json.loads(node.constantInput),
             metadata=json.loads(node.metadata),
         )
-        obj._input_links = [Link.from_db(link) for link in node.Input or []]
-        obj._output_links = [Link.from_db(link) for link in node.Output or []]
+        obj.input_links = [Link.from_db(link) for link in node.Input or []]
+        obj.output_links = [Link.from_db(link) for link in node.Output or []]
         return obj
 
 

autogpt_platform/backend/backend/data/queue.py

@@ -2,6 +2,7 @@
 import logging
 from abc import ABC, abstractmethod
 from datetime import datetime
+from functools import cached_property
 
 from backend.data import redis
 from backend.data.execution import ExecutionResult
@@ -17,14 +18,6 @@ def default(self, o):
 
 
 class AbstractEventQueue(ABC):
-    @abstractmethod
-    def connect(self):
-        pass
-
-    @abstractmethod
-    def close(self):
-        pass
-
     @abstractmethod
     def put(self, execution_result: ExecutionResult):
         pass
@@ -36,26 +29,23 @@ def get(self) -> ExecutionResult | None:
 
 class RedisEventQueue(AbstractEventQueue):
     def __init__(self):
-        self.connection = None
         self.queue_name = redis.QUEUE_NAME
 
-    def connect(self):
-        self.connection = redis.connect()
+    @cached_property
+    def connection(self):
+        return redis.get_redis()
 
     def put(self, execution_result: ExecutionResult):
-        if self.connection:
-            message = json.dumps(execution_result.model_dump(), cls=DateTimeEncoder)
-            logger.info(f"Putting execution result to Redis {message}")
-            self.connection.lpush(self.queue_name, message)
+        message = json.dumps(execution_result.model_dump(), cls=DateTimeEncoder)
+        logger.info(f"Putting execution result to Redis {message}")
+        self.connection.lpush(self.queue_name, message)
 
     def get(self) -> ExecutionResult | None:
-        if self.connection:
-            message = self.connection.rpop(self.queue_name)
-            if message is not None and isinstance(message, (str, bytes, bytearray)):
-                data = json.loads(message)
-                logger.info(f"Getting execution result from Redis {data}")
-                return ExecutionResult(**data)
+        message = self.connection.rpop(self.queue_name)
+        if message is not None and isinstance(message, (str, bytes, bytearray)):
+            data = json.loads(message)
+            logger.info(f"Getting execution result from Redis {data}")
+            return ExecutionResult(**data)
+        elif message is not None:
+            logger.error(f"Failed to get execution result from Redis {message}")
         return None
-
-    def close(self):
-        redis.disconnect()

autogpt_platform/backend/backend/exec.py

@@ -1,12 +1,14 @@
 from backend.app import run_processes
 from backend.executor import ExecutionManager
+from backend.server.db_api import DatabaseAPI
 
 
 def main():
     """
     Run all the processes required for the AutoGPT-server REST API.
     """
     run_processes(
+        DatabaseAPI(),
         ExecutionManager(),
     )