feat(rnd,infra): Pull out websockets away from server api (#7899)

* standalone websocket server

* add websocket url

* wip: talk to ws directly

* rename to api server

* dockerfile and queue

* fix paths

* update poetry lock

* helm charts for websockets

* create seperate deployments for websockets and rest server with redis queue for async comms

* delete duplicate queue

* add depends in ws_api

* singleton for conn manager

* update from review

* fix CI

* address feedback

* update readme

* update docker file and add migration step in readm

* ad watch

* add step to copy example env file

* put connect back in

rnd/autogpt_builder/.env.example

@@ -1,4 +1,5 @@
 NEXT_PUBLIC_AGPT_SERVER_URL=http://localhost:8000/api
+NEXT_PUBLIC_AGPT_WS_SERVER_URL=ws://localhost:8001/ws
 NEXT_PUBLIC_AGPT_MARKETPLACE_URL=http://localhost:8001/api/v1/market
 
 ## Supabase credentials

rnd/autogpt_builder/src/lib/autogpt-server-api/client.ts

@@ -21,9 +21,11 @@ export default class AutoGPTServerAPI {
   constructor(
     baseUrl: string = process.env.NEXT_PUBLIC_AGPT_SERVER_URL ||
       "http://localhost:8000/api",
+    wsUrl: string = process.env.NEXT_PUBLIC_AGPT_WS_SERVER_URL ||
+      "ws://localhost:8001/ws",
   ) {
     this.baseUrl = baseUrl;
-    this.wsUrl = `ws://${new URL(this.baseUrl).host}/ws`;
+    this.wsUrl = wsUrl;
   }
 
   async createUser(): Promise<User> {

rnd/autogpt_server/Dockerfile

@@ -25,16 +25,23 @@ ENV POETRY_VERSION=1.8.3 \
     PATH="$POETRY_HOME/bin:$PATH"
 RUN pip3 install poetry
 
-COPY rnd/autogpt_server /app/rnd/autogpt_server
-COPY rnd/autogpt_libs /app/rnd/autogpt_libs
 COPY autogpt /app/autogpt
 COPY forge /app/forge
+COPY rnd/autogpt_libs /app/rnd/autogpt_libs
 
 WORKDIR /app/rnd/autogpt_server
 
-# Install dependencies
+COPY rnd/autogpt_server/pyproject.toml rnd/autogpt_server/poetry.lock ./
+
 RUN poetry install --no-interaction --no-ansi
 
+COPY rnd/autogpt_server/postgres/schema.prisma app/rnd/autogpt_server/schema.prisma
+RUN poetry run prisma generate
+
+COPY rnd/autogpt_server /app/rnd/autogpt_server
+
+WORKDIR /app/rnd/autogpt_server
+
 RUN poetry run prisma generate
 
 FROM server_base as server

rnd/autogpt_server/Dockerfile.ws

@@ -0,0 +1,52 @@
+FROM python:3.11-slim-buster as server_base
+
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+
+WORKDIR /app
+
+RUN apt-get update \
+    && apt-get install -y build-essential curl ffmpeg wget libcurl4-gnutls-dev libexpat1-dev gettext libz-dev libssl-dev \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/* \
+    && wget https://github.com/git/git/archive/v2.28.0.tar.gz -O git.tar.gz \
+    && tar -zxf git.tar.gz \
+    && cd git-* \
+    && make prefix=/usr all \
+    && make prefix=/usr install
+
+
+ENV POETRY_VERSION=1.8.3 \
+    POETRY_HOME="/opt/poetry" \
+    POETRY_NO_INTERACTION=1 \
+    POETRY_VIRTUALENVS_CREATE=false \
+    PATH="$POETRY_HOME/bin:$PATH"
+RUN pip3 install poetry
+
+COPY autogpt /app/autogpt
+COPY forge /app/forge
+COPY rnd/autogpt_libs /app/rnd/autogpt_libs
+
+WORKDIR /app/rnd/autogpt_server
+
+COPY rnd/autogpt_server/pyproject.toml rnd/autogpt_server/poetry.lock ./
+
+RUN poetry install --no-interaction --no-ansi
+
+COPY rnd/autogpt_server/postgres/schema.prisma app/rnd/autogpt_server/schema.prisma
+RUN poetry run prisma generate
+
+COPY rnd/autogpt_server /app/rnd/autogpt_server
+
+WORKDIR /app/rnd/autogpt_server
+
+RUN poetry run prisma generate
+
+FROM server_base as server
+
+ENV PORT=8001
+ENV DATABASE_URL=""
+
+CMD ["poetry", "run", "ws"]

rnd/autogpt_server/README.md

@@ -32,8 +32,14 @@ We use the Poetry to manage the dependencies. To set up the project, follow thes
    ```sh
    poetry install
    ```
+
+4. Copy .env.example to .env
+
+   ```sh
+   cp .env.example .env
+   ```
    
-4. Generate the Prisma client
+5. Generate the Prisma client
 
    ```sh
    poetry run prisma generate
@@ -49,20 +55,45 @@ We use the Poetry to manage the dependencies. To set up the project, follow thes
    > Then run the generation again. The path *should* look something like this:  
    > `<some path>/pypoetry/virtualenvs/autogpt-server-TQIRSwR6-py3.12/bin/prisma`
 
-5. Migrate the database. Be careful because this deletes current data in the database.
+6. Migrate the database. Be careful because this deletes current data in the database.
 
    ```sh
-   poetry run prisma migrate dev
+   docker compose up postgres -d
+   poetry run prisma migrate dev --schema postgres/schema.prisma
    ```
 
 ## Running The Server
 
 ### Starting the server directly
 
-Run the following command:
+Run the following command to build the dockerfiles:
 
 ```sh
-poetry run app
+docker compose build
+```
+
+Run the following command to run the app:
+
+```sh
+docker compose up
+```
+
+Run the following to automatically rebuild when code changes, in another terminal:
+
+```sh
+docker compose watch
+```
+
+Run the following command to shut down:
+
+```sh
+docker compose down
+```
+
+If you run into issues with dangling orphans, try:
+
+```sh
+docker-compose down --volumes --remove-orphans && docker-compose up --force-recreate --renew-anon-volumes --remove-orphans  
 ```
 
 ## Testing

rnd/autogpt_server/autogpt_server/data/queue.py

@@ -0,0 +1,78 @@
+import json
+import logging
+import os
+from abc import ABC, abstractmethod
+from datetime import datetime
+
+from redis.asyncio import Redis
+
+from autogpt_server.data.execution import ExecutionResult
+
+logger = logging.getLogger(__name__)
+
+
+class DateTimeEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, datetime):
+            return o.isoformat()
+        return super().default(o)
+
+
+class AsyncEventQueue(ABC):
+    @abstractmethod
+    async def connect(self):
+        pass
+
+    @abstractmethod
+    async def close(self):
+        pass
+
+    @abstractmethod
+    async def put(self, execution_result: ExecutionResult):
+        pass
+
+    @abstractmethod
+    async def get(self) -> ExecutionResult | None:
+        pass
+
+
+class AsyncRedisEventQueue(AsyncEventQueue):
+
+    def __init__(self):
+        self.host = os.getenv("REDIS_HOST", "localhost")
+        self.port = int(os.getenv("REDIS_PORT", "6379"))
+        self.password = os.getenv("REDIS_PASSWORD", None)
+        self.queue_name = os.getenv("REDIS_QUEUE", "execution_events")
+        self.connection = None
+
+    async def connect(self):
+        if not self.connection:
+            self.connection = Redis(
+                host=self.host,
+                port=self.port,
+                password=self.password,
+                decode_responses=True,
+            )
+            await self.connection.ping()
+            logger.info(f"Connected to Redis on {self.host}:{self.port}")
+
+    async def put(self, execution_result: ExecutionResult):
+        if self.connection:
+            message = json.dumps(execution_result.model_dump(), cls=DateTimeEncoder)
+            logger.info(f"Put {message}")
+            await self.connection.lpush(self.queue_name, message)  # type: ignore
+
+    async def get(self) -> ExecutionResult | None:
+        if self.connection:
+            message = await self.connection.rpop(self.queue_name)  # type: ignore
+            if message is not None and isinstance(message, (str, bytes, bytearray)):
+                data = json.loads(message)
+                logger.info(f"Get {data}")
+                return ExecutionResult(**data)
+        return None
+
+    async def close(self):
+        if self.connection:
+            await self.connection.close()
+            self.connection = None
+            logger.info("Closed connection to Redis")

rnd/autogpt_server/autogpt_server/executor/manager.py

@@ -5,7 +5,7 @@ from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Coroutine, Generator, TypeVar
 
 if TYPE_CHECKING:
-    from autogpt_server.server.server import AgentServer
+    from autogpt_server.server.rest_api import AgentServer
 
 from autogpt_server.blocks.basic import InputBlock
 from autogpt_server.data import db
@@ -300,7 +300,7 @@ def validate_exec(
 
 
 def get_agent_server_client() -> "AgentServer":
-    from autogpt_server.server.server import AgentServer
+    from autogpt_server.server.rest_api import AgentServer
 
     return get_service_client(AgentServer)
 
@@ -409,6 +409,7 @@ class ExecutionManager(AppService):
     def __init__(self):
         self.pool_size = Config().num_graph_workers
         self.queue = ExecutionQueue[GraphExecution]()
+        self.use_redis = False
 
     def run_service(self):
         with ProcessPoolExecutor(
@@ -433,7 +434,6 @@ class ExecutionManager(AppService):
         if not graph:
             raise Exception(f"Graph #{graph_id} not found.")
         graph.validate_graph(for_run=True)
-
         nodes_input = []
         for node in graph.starting_nodes:
             input_data = {}

rnd/autogpt_server/autogpt_server/executor/scheduler.py

@@ -22,6 +22,7 @@ class ExecutionScheduler(AppService):
     def __init__(self, refresh_interval=10):
         self.last_check = datetime.min
         self.refresh_interval = refresh_interval
+        self.use_redis = False
 
     @property
     def execution_manager_client(self) -> ExecutionManager:

rnd/autogpt_server/autogpt_server/server/__init__.py

@@ -1,3 +1,3 @@
-from .server import AgentServer
+from .rest_api import AgentServer
 
 __all__ = ["AgentServer"]

rnd/autogpt_server/autogpt_server/server/rest_api.py

@@ -1,4 +1,3 @@
-import asyncio
 import inspect
 from collections import defaultdict
 from contextlib import asynccontextmanager
@@ -6,21 +5,11 @@ from functools import wraps
 from typing import Annotated, Any, Dict
 
 import uvicorn
-from autogpt_libs.auth.jwt_utils import parse_jwt_token
 from autogpt_libs.auth.middleware import auth_middleware
-from fastapi import (
-    APIRouter,
-    Body,
-    Depends,
-    FastAPI,
-    HTTPException,
-    WebSocket,
-    WebSocketDisconnect,
-)
+from fastapi import APIRouter, Body, Depends, FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
-import autogpt_server.server.ws_api
 from autogpt_server.data import block, db
 from autogpt_server.data import graph as graph_db
 from autogpt_server.data import user as user_db
@@ -30,15 +19,11 @@ from autogpt_server.data.execution import (
     get_execution_results,
     list_executions,
 )
-from autogpt_server.data.user import DEFAULT_USER_ID, get_or_create_user
+from autogpt_server.data.queue import AsyncEventQueue, AsyncRedisEventQueue
+from autogpt_server.data.user import get_or_create_user
 from autogpt_server.executor import ExecutionManager, ExecutionScheduler
-from autogpt_server.server.conn_manager import ConnectionManager
-from autogpt_server.server.model import (
-    CreateGraph,
-    Methods,
-    SetGraphActiveVersion,
-    WsMessage,
-)
+from autogpt_server.server.model import CreateGraph, SetGraphActiveVersion
+from autogpt_server.util.auth import get_user_id
 from autogpt_server.util.lock import KeyedMutex
 from autogpt_server.util.service import AppService, expose, get_service_client
 from autogpt_server.util.settings import Settings
@@ -46,37 +31,24 @@ from autogpt_server.util.settings import Settings
 settings = Settings()
 
 
-def get_user_id(payload: dict = Depends(auth_middleware)) -> str:
-    if not payload:
-        # This handles the case when authentication is disabled
-        return DEFAULT_USER_ID
-
-    user_id = payload.get("sub")
-    if not user_id:
-        raise HTTPException(status_code=401, detail="User ID not found in token")
-    return user_id
-
-
 class AgentServer(AppService):
-    event_queue: asyncio.Queue[ExecutionResult] = asyncio.Queue()
-    manager = ConnectionManager()
     mutex = KeyedMutex()
     use_db = False
+    use_redis = True
     _test_dependency_overrides = {}
 
-    async def event_broadcaster(self):
-        while True:
-            event: ExecutionResult = await self.event_queue.get()
-            await self.manager.send_execution_result(event)
+    def __init__(self, event_queue: AsyncEventQueue | None = None):
+        self.event_queue = event_queue or AsyncRedisEventQueue()
 
     @asynccontextmanager
     async def lifespan(self, _: FastAPI):
         await db.connect()
+        self.run_and_wait(self.event_queue.connect())
         await block.initialize_blocks()
         if await user_db.create_default_user(settings.config.enable_auth):
             await graph_db.import_packaged_templates()
-        asyncio.create_task(self.event_broadcaster())
         yield
+        await self.event_queue.close()
         await db.disconnect()
 
     def run_service(self):
@@ -223,10 +195,6 @@ class AgentServer(AppService):
 
         app.include_router(router)
 
-        @app.websocket("/ws")
-        async def websocket_endpoint(websocket: WebSocket):  # type: ignore
-            await self.websocket_router(websocket)
-
         uvicorn.run(app, host="0.0.0.0", port=8000)
 
     def set_test_dependency_overrides(self, overrides: dict):
@@ -276,64 +244,6 @@ class AgentServer(AppService):
             status_code=500,
         )
 
-    async def authenticate_websocket(self, websocket: WebSocket) -> str:
-        if settings.config.enable_auth.lower() == "true":
-            token = websocket.query_params.get("token")
-            if not token:
-                await websocket.close(code=4001, reason="Missing authentication token")
-                return ""
-
-            try:
-                payload = parse_jwt_token(token)
-                user_id = payload.get("sub")
-                if not user_id:
-                    await websocket.close(code=4002, reason="Invalid token")
-                    return ""
-                return user_id
-            except ValueError:
-                await websocket.close(code=4003, reason="Invalid token")
-                return ""
-        else:
-            return user_db.DEFAULT_USER_ID
-
-    async def websocket_router(self, websocket: WebSocket):
-        user_id = await self.authenticate_websocket(websocket)
-        if not user_id:
-            return
-        await self.manager.connect(websocket)
-        try:
-            while True:
-                data = await websocket.receive_text()
-                message = WsMessage.model_validate_json(data)
-                if message.method == Methods.SUBSCRIBE:
-                    await autogpt_server.server.ws_api.handle_subscribe(
-                        websocket, self.manager, message
-                    )
-
-                elif message.method == Methods.UNSUBSCRIBE:
-                    await autogpt_server.server.ws_api.handle_unsubscribe(
-                        websocket, self.manager, message
-                    )
-
-                elif message.method == Methods.ERROR:
-                    print("WebSocket Error message received:", message.data)
-
-                else:
-                    print(
-                        f"Message type {message.method} is not processed by the server"
-                    )
-                    await websocket.send_text(
-                        WsMessage(
-                            method=Methods.ERROR,
-                            success=False,
-                            error="Message type is not processed by the server",
-                        ).model_dump_json()
-                    )
-
-        except WebSocketDisconnect:
-            self.manager.disconnect(websocket)
-            print("Client Disconnected")
-
     @classmethod
     async def get_or_create_user_route(cls, user_data: dict = Depends(auth_middleware)):
         user = await get_or_create_user(user_data)

rnd/autogpt_server/autogpt_server/server/ws_api.py

@@ -1,33 +1,80 @@
-from fastapi import WebSocket, WebSocketDisconnect
+import asyncio
+import logging
 
+from autogpt_libs.auth import parse_jwt_token
+from fastapi import Depends, FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+
+from autogpt_server.data.queue import AsyncRedisEventQueue
+from autogpt_server.data.user import DEFAULT_USER_ID
 from autogpt_server.server.conn_manager import ConnectionManager
 from autogpt_server.server.model import ExecutionSubscription, Methods, WsMessage
+from autogpt_server.util.settings import Settings
+
+settings = Settings()
+
+app = FastAPI()
+event_queue = AsyncRedisEventQueue()
+_connection_manager = None
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=[
+        "http://localhost:3000",
+        "http://127.0.0.1:3000",
+        "https://dev-builder.agpt.co",
+    ],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 
-async def websocket_router(websocket: WebSocket, manager: ConnectionManager):
-    await manager.connect(websocket)
-    try:
-        while True:
-            data = await websocket.receive_text()
-            message = WsMessage.model_validate_json(data)
-            if message.method == Methods.SUBSCRIBE:
-                await handle_subscribe(websocket, manager, message)
+def get_connection_manager():
+    global _connection_manager
+    if _connection_manager is None:
+        _connection_manager = ConnectionManager()
+    return _connection_manager
 
-            elif message.method == Methods.UNSUBSCRIBE:
-                await handle_unsubscribe(websocket, manager, message)
-            else:
-                print("Message type is not processed by the server")
-                await websocket.send_text(
-                    WsMessage(
-                        method=Methods.ERROR,
-                        success=False,
-                        error="Message type is not processed by the server",
-                    ).model_dump_json()
-                )
 
-    except WebSocketDisconnect:
-        manager.disconnect(websocket)
-        print("Client Disconnected")
+@app.on_event("startup")
+async def startup_event():
+    await event_queue.connect()
+    manager = get_connection_manager()
+    asyncio.create_task(event_broadcaster(manager))
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    await event_queue.close()
+
+
+async def event_broadcaster(manager: ConnectionManager):
+    while True:
+        event = await event_queue.get()
+        if event is not None:
+            await manager.send_execution_result(event)
+
+
+async def authenticate_websocket(websocket: WebSocket) -> str:
+    if settings.config.enable_auth.lower() == "true":
+        token = websocket.query_params.get("token")
+        if not token:
+            await websocket.close(code=4001, reason="Missing authentication token")
+            return ""
+
+        try:
+            payload = parse_jwt_token(token)
+            user_id = payload.get("sub")
+            if not user_id:
+                await websocket.close(code=4002, reason="Invalid token")
+                return ""
+            return user_id
+        except ValueError:
+            await websocket.close(code=4003, reason="Invalid token")
+            return ""
+    else:
+        return DEFAULT_USER_ID
 
 
 async def handle_subscribe(
@@ -76,3 +123,46 @@ async def handle_unsubscribe(
                 channel=ex_sub.graph_id,
             ).model_dump_json()
         )
+
+
+@app.get("/")
+async def health():
+    return {"status": "healthy"}
+
+
+@app.websocket("/ws")
+async def websocket_router(
+    websocket: WebSocket, manager: ConnectionManager = Depends(get_connection_manager)
+):
+    user_id = await authenticate_websocket(websocket)
+    if not user_id:
+        return
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_text()
+            message = WsMessage.model_validate_json(data)
+            if message.method == Methods.SUBSCRIBE:
+                await handle_subscribe(websocket, manager, message)
+
+            elif message.method == Methods.UNSUBSCRIBE:
+                await handle_unsubscribe(websocket, manager, message)
+
+            elif message.method == Methods.ERROR:
+                logging.error("WebSocket Error message received:", message.data)
+
+            else:
+                logging.info(
+                    f"Message type {message.method} is not processed by the server"
+                )
+                await websocket.send_text(
+                    WsMessage(
+                        method=Methods.ERROR,
+                        success=False,
+                        error="Message type is not processed by the server",
+                    ).model_dump_json()
+                )
+
+    except WebSocketDisconnect:
+        manager.disconnect(websocket)
+        logging.info("Client Disconnected")

rnd/autogpt_server/autogpt_server/util/auth.py

@@ -0,0 +1,15 @@
+from autogpt_libs.auth import auth_middleware
+from fastapi import Depends, HTTPException
+
+from autogpt_server.data.user import DEFAULT_USER_ID
+
+
+def get_user_id(payload: dict = Depends(auth_middleware)) -> str:
+    if not payload:
+        # This handles the case when authentication is disabled
+        return DEFAULT_USER_ID
+
+    user_id = payload.get("sub")
+    if not user_id:
+        raise HTTPException(status_code=401, detail="User ID not found in token")
+    return user_id

rnd/autogpt_server/autogpt_server/util/service.py

@@ -7,14 +7,17 @@ from typing import Any, Callable, Coroutine, Type, TypeVar, cast
 
 from Pyro5 import api as pyro
 from Pyro5 import nameserver
-from tenacity import retry, stop_after_delay, wait_exponential
+from tenacity import retry, stop_after_attempt, wait_exponential
 
 from autogpt_server.data import db
+from autogpt_server.data.queue import AsyncEventQueue, AsyncRedisEventQueue
 from autogpt_server.util.process import AppProcess
 from autogpt_server.util.settings import Config
 
 logger = logging.getLogger(__name__)
-conn_retry = retry(stop=stop_after_delay(5), wait=wait_exponential(multiplier=0.1))
+conn_retry = retry(
+    stop=stop_after_attempt(30), wait=wait_exponential(multiplier=1, min=1, max=30)
+)
 T = TypeVar("T")
 
 pyro_host = Config().pyro_host
@@ -43,7 +46,9 @@ class PyroNameServer(AppProcess):
 
 class AppService(AppProcess):
     shared_event_loop: asyncio.AbstractEventLoop
+    event_queue: AsyncEventQueue = AsyncRedisEventQueue()
     use_db: bool = True
+    use_redis: bool = False
 
     @classmethod
     @property
@@ -66,6 +71,8 @@ class AppService(AppProcess):
         self.shared_event_loop = asyncio.get_event_loop()
         if self.use_db:
             self.shared_event_loop.run_until_complete(db.connect())
+        if self.use_redis:
+            self.shared_event_loop.run_until_complete(self.event_queue.connect())
 
         # Initialize the async loop.
         async_thread = threading.Thread(target=self.__start_async_loop)

rnd/autogpt_server/autogpt_server/util/test.py

@@ -1,21 +1,57 @@
+import asyncio
 import time
 
 from autogpt_server.data import db
 from autogpt_server.data.block import Block, initialize_blocks
-from autogpt_server.data.execution import ExecutionStatus
+from autogpt_server.data.execution import ExecutionResult, ExecutionStatus
+from autogpt_server.data.queue import AsyncEventQueue
 from autogpt_server.executor import ExecutionManager, ExecutionScheduler
 from autogpt_server.server import AgentServer
-from autogpt_server.server.server import get_user_id
+from autogpt_server.server.rest_api import get_user_id
 from autogpt_server.util.service import PyroNameServer
 
 log = print
 
 
+class InMemoryAsyncEventQueue(AsyncEventQueue):
+    def __init__(self):
+        self.queue = asyncio.Queue()
+        self.connected = False
+        self.closed = False
+
+    async def connect(self):
+        if not self.connected:
+            self.connected = True
+        return
+
+    async def close(self):
+        self.closed = True
+        self.connected = False
+        return
+
+    async def put(self, execution_result: ExecutionResult):
+        if not self.connected:
+            raise RuntimeError("Queue is not connected")
+        await self.queue.put(execution_result)
+
+    async def get(self):
+        if self.closed:
+            return None
+        if not self.connected:
+            raise RuntimeError("Queue is not connected")
+        try:
+            item = await asyncio.wait_for(self.queue.get(), timeout=0.1)
+            return item
+        except asyncio.TimeoutError:
+            return None
+
+
 class SpinTestServer:
     def __init__(self):
         self.name_server = PyroNameServer()
         self.exec_manager = ExecutionManager()
-        self.agent_server = AgentServer()
+        self.in_memory_queue = InMemoryAsyncEventQueue()
+        self.agent_server = AgentServer(event_queue=self.in_memory_queue)
         self.scheduler = ExecutionScheduler()
 
     @staticmethod

rnd/autogpt_server/autogpt_server/ws_app.py

@@ -0,0 +1,11 @@
+import uvicorn
+
+from autogpt_server.server.ws_api import app
+
+
+def main():
+    uvicorn.run(app, host="0.0.0.0", port=8001)
+
+
+if __name__ == "__main__":
+    main()

rnd/autogpt_server/pyproject.toml

@@ -44,6 +44,9 @@ uvicorn = { extras = ["standard"], version = "^0.30.1" }
 websockets = "^12.0"
 youtube-transcript-api = "^0.6.2"
 
+aio-pika = "^9.4.3"
+redis = "^5.0.8"
+
 [tool.poetry.group.dev.dependencies]
 cx-freeze = { git = "https://github.com/ntindle/cx_Freeze.git", rev = "main", develop = true }
 poethepoet = "^0.26.1"
@@ -61,6 +64,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.scripts]
 app = "autogpt_server.app:main"
+ws = "autogpt_server.ws_app:main"
 cli = "autogpt_server.cli:main"
 format = "linter:format"
 lint = "linter:lint"

rnd/autogpt_server/test/util/test_service.py

@@ -7,6 +7,8 @@ from autogpt_server.util.service import (
 
 
 class TestService(AppService):
+    def __init__(self):
+        self.use_redis = False
 
     def run_service(self):
         super().run_service()

rnd/docker-compose.yml

@@ -6,7 +6,6 @@ services:
       POSTGRES_USER: agpt_user
       POSTGRES_PASSWORD: pass123
       POSTGRES_DB: agpt_local
-      PGUSER: 5432
     healthcheck:
       test: pg_isready -U $$POSTGRES_USER -d $$POSTGRES_DB
       interval: 10s
@@ -14,3 +13,51 @@ services:
       retries: 5
     ports:
       - "5432:5432"
+
+  redis:
+    image: redis:latest
+    command: redis-server --requirepass password
+    ports:
+      - "6379:6379"
+
+  rest_server:
+    build:
+      context: ../
+      dockerfile: rnd/autogpt_server/Dockerfile
+    develop:
+      watch:
+        - path: ./
+          target: rnd/autogpt_server/
+          action: rebuild
+    depends_on:
+      - postgres
+      - redis
+    environment:
+      - DATABASE_URL=postgresql://agpt_user:pass123@postgres:5432/agpt_local
+      - REDIS_HOST=redis
+      - REDIS_PORT=6379
+      - REDIS_PASSWORD=password
+      - AUTH_ENABLED=false
+    ports:
+      - "8000:8000"
+
+  ws_server:
+    build:
+      context: ../
+      dockerfile: rnd/autogpt_server/Dockerfile.ws
+    develop:
+      watch:
+        - path: ./
+          target: rnd/autogpt_server/
+          action: rebuild
+    depends_on:
+      - postgres
+      - redis
+    environment:
+      - DATABASE_URL=postgresql://agpt_user:pass123@postgres:5432/agpt_local
+      - REDIS_HOST=redis
+      - REDIS_PORT=6379
+      - REDIS_PASSWORD=password
+      - AUTH_ENABLED=false
+    ports:
+      - "8001:8001"

rnd/infra/helm/autogpt-server/values.dev.yaml

@@ -42,11 +42,11 @@ ingress:
 
 resources:
   requests:
-    cpu: 100m
-    memory: 128Mi
+    cpu: 200m
+    memory: 256Mi
   limits:
-    cpu: 500m
-    memory: 512Mi
+    cpu: 2
+    memory: 2Gi
 
 livenessProbe:
   httpGet:
@@ -82,4 +82,6 @@ env:
   APP_ENV: "dev"
   PYRO_HOST: "0.0.0.0"
   NUM_GRAPH_WORKERS: 100
-  NUM_NODE_WORKERS: 100
+  NUM_NODE_WORKERS: 100
+  REDIS_HOST: "redis-dev-master.redis-dev.svc.cluster.local"
+  REDIS_PORT: "6379"

rnd/infra/helm/autogpt-websocket-server/.helmignore

@@ -0,0 +1,23 @@
+# Patterns to ignore when building packages.
+# This supports shell glob matching, relative path matching, and
+# negation (prefixed with !). Only one pattern per line.
+.DS_Store
+# Common VCS dirs
+.git/
+.gitignore
+.bzr/
+.bzrignore
+.hg/
+.hgignore
+.svn/
+# Common backup files
+*.swp
+*.bak
+*.tmp
+*.orig
+*~
+# Various IDEs
+.project
+.idea/
+*.tmproj
+.vscode/

rnd/infra/helm/autogpt-websocket-server/Chart.yaml

@@ -0,0 +1,10 @@
+apiVersion: v2
+name: autogpt-websocket-server
+description: A Helm chart for Websocket Server
+
+type: application
+
+
+version: 0.1.0
+
+appVersion: "1.0.0"

rnd/infra/helm/autogpt-websocket-server/templates/NOTES.txt

@@ -0,0 +1,22 @@
+1. Get the application URL by running these commands:
+{{- if .Values.ingress.enabled }}
+{{- range $host := .Values.ingress.hosts }}
+  {{- range .paths }}
+  http{{ if $.Values.ingress.tls }}s{{ end }}://{{ $host.host }}{{ .path }}
+  {{- end }}
+{{- end }}
+{{- else if contains "NodePort" .Values.service.type }}
+  export NODE_PORT=$(kubectl get --namespace {{ .Release.Namespace }} -o jsonpath="{.spec.ports[0].nodePort}" services {{ include "autogpt-websocket-server.fullname" . }})
+  export NODE_IP=$(kubectl get nodes --namespace {{ .Release.Namespace }} -o jsonpath="{.items[0].status.addresses[0].address}")
+  echo http://$NODE_IP:$NODE_PORT
+{{- else if contains "LoadBalancer" .Values.service.type }}
+     NOTE: It may take a few minutes for the LoadBalancer IP to be available.
+           You can watch its status by running 'kubectl get --namespace {{ .Release.Namespace }} svc -w {{ include "autogpt-websocket-server.fullname" . }}'
+  export SERVICE_IP=$(kubectl get svc --namespace {{ .Release.Namespace }} {{ include "autogpt-websocket-server.fullname" . }} --template "{{"{{ range (index .status.loadBalancer.ingress 0) }}{{.}}{{ end }}"}}")
+  echo http://$SERVICE_IP:{{ .Values.service.port }}
+{{- else if contains "ClusterIP" .Values.service.type }}
+  export POD_NAME=$(kubectl get pods --namespace {{ .Release.Namespace }} -l "app.kubernetes.io/name={{ include "autogpt-websocket-server.name" . }},app.kubernetes.io/instance={{ .Release.Name }}" -o jsonpath="{.items[0].metadata.name}")
+  export CONTAINER_PORT=$(kubectl get pod --namespace {{ .Release.Namespace }} $POD_NAME -o jsonpath="{.spec.containers[0].ports[0].containerPort}")
+  echo "Visit http://127.0.0.1:8080 to use your application"
+  kubectl --namespace {{ .Release.Namespace }} port-forward $POD_NAME 8080:$CONTAINER_PORT
+{{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/_helpers.tpl

@@ -0,0 +1,62 @@
+{{/*
+Expand the name of the chart.
+*/}}
+{{- define "autogpt-websocket-server.name" -}}
+{{- default .Chart.Name .Values.nameOverride | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Create a default fully qualified app name.
+We truncate at 63 chars because some Kubernetes name fields are limited to this (by the DNS naming spec).
+If release name contains chart name it will be used as a full name.
+*/}}
+{{- define "autogpt-websocket-server.fullname" -}}
+{{- if .Values.fullnameOverride }}
+{{- .Values.fullnameOverride | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- $name := default .Chart.Name .Values.nameOverride }}
+{{- if contains $name .Release.Name }}
+{{- .Release.Name | trunc 63 | trimSuffix "-" }}
+{{- else }}
+{{- printf "%s-%s" .Release.Name $name | trunc 63 | trimSuffix "-" }}
+{{- end }}
+{{- end }}
+{{- end }}
+
+{{/*
+Create chart name and version as used by the chart label.
+*/}}
+{{- define "autogpt-websocket-server.chart" -}}
+{{- printf "%s-%s" .Chart.Name .Chart.Version | replace "+" "_" | trunc 63 | trimSuffix "-" }}
+{{- end }}
+
+{{/*
+Common labels
+*/}}
+{{- define "autogpt-websocket-server.labels" -}}
+helm.sh/chart: {{ include "autogpt-websocket-server.chart" . }}
+{{ include "autogpt-websocket-server.selectorLabels" . }}
+{{- if .Chart.AppVersion }}
+app.kubernetes.io/version: {{ .Chart.AppVersion | quote }}
+{{- end }}
+app.kubernetes.io/managed-by: {{ .Release.Service }}
+{{- end }}
+
+{{/*
+Selector labels
+*/}}
+{{- define "autogpt-websocket-server.selectorLabels" -}}
+app.kubernetes.io/name: {{ include "autogpt-websocket-server.name" . }}
+app.kubernetes.io/instance: {{ .Release.Name }}
+{{- end }}
+
+{{/*
+Create the name of the service account to use
+*/}}
+{{- define "autogpt-websocket-server.serviceAccountName" -}}
+{{- if .Values.serviceAccount.create }}
+{{- default (include "autogpt-websocket-server.fullname" .) .Values.serviceAccount.name }}
+{{- else }}
+{{- default "default" .Values.serviceAccount.name }}
+{{- end }}
+{{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/configmap.yaml

@@ -0,0 +1,10 @@
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: {{ include "autogpt-websocket-server.fullname" . }}-config
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+data:
+  {{- range $key, $value := .Values.env }}
+  {{ $key }}: {{ $value | quote }}
+  {{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/deployment.yaml

@@ -0,0 +1,75 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ include "autogpt-websocket-server.fullname" . }}
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+spec:
+  {{- if not .Values.autoscaling.enabled }}
+  replicas: {{ .Values.replicaCount }}
+  {{- end }}
+  selector:
+    matchLabels:
+      {{- include "autogpt-websocket-server.selectorLabels" . | nindent 6 }}
+  template:
+    metadata:
+      {{- with .Values.podAnnotations }}
+      annotations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      labels:
+        {{- include "autogpt-websocket-server.labels" . | nindent 8 }}
+        {{- with .Values.podLabels }}
+        {{- toYaml . | nindent 8 }}
+        {{- end }}
+    spec:
+      {{- with .Values.imagePullSecrets }}
+      imagePullSecrets:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      serviceAccountName: {{ include "autogpt-websocket-server.serviceAccountName" . }}
+      securityContext:
+        {{- toYaml .Values.podSecurityContext | nindent 8 }}
+      containers:
+        - name: {{ .Chart.Name }}
+          envFrom:
+            - configMapRef:
+                name: {{ include "autogpt-websocket-server.fullname" . }}-config
+          securityContext:
+            {{- toYaml .Values.securityContext | nindent 12 }}
+          image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
+          imagePullPolicy: {{ .Values.image.pullPolicy }}
+          ports:
+            - name: ws
+              containerPort: {{ .Values.service.port }}
+              protocol: TCP
+          {{- if .Values.livenessProbe.enabled }}
+          livenessProbe:
+            {{- toYaml .Values.livenessProbe | nindent 12 }}
+          {{- end }}
+          {{- if .Values.readinessProbe.enabled }}
+          readinessProbe:
+            {{- toYaml .Values.readinessProbe | nindent 12 }}
+          {{- end }}
+          resources:
+            {{- toYaml .Values.resources | nindent 12 }}
+          {{- with .Values.volumeMounts }}
+          volumeMounts:
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
+      {{- with .Values.volumes }}
+      volumes:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.nodeSelector }}
+      nodeSelector:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.affinity }}
+      affinity:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with .Values.tolerations }}
+      tolerations:
+        {{- toYaml . | nindent 8 }}
+      {{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/hpa.yaml

@@ -0,0 +1,32 @@
+{{- if .Values.autoscaling.enabled }}
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: {{ include "autogpt-websocket-server.fullname" . }}
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: Deployment
+    name: {{ include "autogpt-websocket-server.fullname" . }}
+  minReplicas: {{ .Values.autoscaling.minReplicas }}
+  maxReplicas: {{ .Values.autoscaling.maxReplicas }}
+  metrics:
+    {{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: cpu
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
+    {{- end }}
+    {{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    - type: Resource
+      resource:
+        name: memory
+        target:
+          type: Utilization
+          averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
+    {{- end }}
+{{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/ingress.yaml

@@ -0,0 +1,61 @@
+{{- if .Values.ingress.enabled -}}
+{{- $fullName := include "autogpt-websocket-server.fullname" . -}}
+{{- $svcPort := .Values.service.port -}}
+{{- if and .Values.ingress.className (not (semverCompare ">=1.18-0" .Capabilities.KubeVersion.GitVersion)) }}
+  {{- if not (hasKey .Values.ingress.annotations "kubernetes.io/ingress.class") }}
+  {{- $_ := set .Values.ingress.annotations "kubernetes.io/ingress.class" .Values.ingress.className}}
+  {{- end }}
+{{- end }}
+{{- if semverCompare ">=1.19-0" .Capabilities.KubeVersion.GitVersion -}}
+apiVersion: networking.k8s.io/v1
+{{- else if semverCompare ">=1.14-0" .Capabilities.KubeVersion.GitVersion -}}
+apiVersion: networking.k8s.io/v1beta1
+{{- else -}}
+apiVersion: extensions/v1beta1
+{{- end }}
+kind: Ingress
+metadata:
+  name: {{ $fullName }}
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+  {{- with .Values.ingress.annotations }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+spec:
+  {{- if and .Values.ingress.className (semverCompare ">=1.18-0" .Capabilities.KubeVersion.GitVersion) }}
+  ingressClassName: {{ .Values.ingress.className }}
+  {{- end }}
+  {{- if .Values.ingress.tls }}
+  tls:
+    {{- range .Values.ingress.tls }}
+    - hosts:
+        {{- range .hosts }}
+        - {{ . | quote }}
+        {{- end }}
+      secretName: {{ .secretName }}
+    {{- end }}
+  {{- end }}
+  rules:
+    {{- range .Values.ingress.hosts }}
+    - host: {{ .host | quote }}
+      http:
+        paths:
+          {{- range .paths }}
+          - path: {{ .path }}
+            {{- if and .pathType (semverCompare ">=1.18-0" $.Capabilities.KubeVersion.GitVersion) }}
+            pathType: {{ .pathType }}
+            {{- end }}
+            backend:
+              {{- if semverCompare ">=1.19-0" $.Capabilities.KubeVersion.GitVersion }}
+              service:
+                name: {{ $fullName }}
+                port:
+                  number: {{ $svcPort }}
+              {{- else }}
+              serviceName: {{ $fullName }}
+              servicePort: {{ $svcPort }}
+              {{- end }}
+          {{- end }}
+    {{- end }}
+{{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/managedcert.yaml

@@ -0,0 +1,7 @@
+apiVersion: networking.gke.io/v1
+kind: ManagedCertificate
+metadata:
+  name: {{ include "autogpt-websocket-server.fullname" . }}-cert
+spec:
+  domains:
+    - {{ .Values.domain }}

rnd/infra/helm/autogpt-websocket-server/templates/service.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ include "autogpt-websocket-server.fullname" . }}
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+spec:
+  type: {{ .Values.service.type }}
+  ports:
+    - port: {{ .Values.service.port }}
+      targetPort: 8001
+      protocol: TCP
+      name: ws
+  selector:
+    {{- include "autogpt-websocket-server.selectorLabels" . | nindent 4 }}

rnd/infra/helm/autogpt-websocket-server/templates/serviceaccount.yaml

@@ -0,0 +1,13 @@
+{{- if .Values.serviceAccount.create -}}
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: {{ include "autogpt-websocket-server.serviceAccountName" . }}
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+  {{- with .Values.serviceAccount.annotations }}
+  annotations:
+    {{- toYaml . | nindent 4 }}
+  {{- end }}
+automountServiceAccountToken: {{ .Values.serviceAccount.automount }}
+{{- end }}

rnd/infra/helm/autogpt-websocket-server/templates/tests/test-connection.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: "{{ include "autogpt-websocket-server.fullname" . }}-test-connection"
+  labels:
+    {{- include "autogpt-websocket-server.labels" . | nindent 4 }}
+  annotations:
+    "helm.sh/hook": test
+spec:
+  containers:
+    - name: wget
+      image: busybox
+      command: ['wget']
+      args: ['{{ include "autogpt-websocket-server.fullname" . }}:{{ .Values.service.port }}']
+  restartPolicy: Never

rnd/infra/helm/autogpt-websocket-server/values.dev.yaml

@@ -0,0 +1,63 @@
+replicaCount: 1 # not scaling websocket server for now
+
+image:
+  repository: us-east1-docker.pkg.dev/agpt-dev/agpt-ws-server-dev/agpt-ws-server-dev
+  tag: latest
+  pullPolicy: Always
+
+service:
+  type: ClusterIP
+  port: 8001
+
+ingress:
+  enabled: true
+  className: "gce"
+  annotations:
+    kubernetes.io/ingress.class: gce
+    kubernetes.io/ingress.global-static-ip-name: "agpt-dev-agpt-ws-server-ip"
+    networking.gke.io/managed-certificates: "autogpt-websocket-server-cert"
+  hosts:
+    - host: dev-ws-server.agpt.co
+      paths:
+        - path: /
+          pathType: Prefix
+      backend:
+        service:
+          name: autogpt-websocket-server
+          port: 8001
+  defaultBackend:
+    service:
+      name: autogpt-websocket-server
+      port:
+        number: 8001
+
+domain: "dev-ws-server.agpt.co"
+
+resources:
+  limits:
+    cpu: 200m
+    memory: 256Mi
+  requests:
+    cpu: 100m
+    memory: 128Mi
+
+autoscaling:
+  enabled: false
+
+readinessProbe:
+  httpGet:
+    path: /
+    port: 8001
+  initialDelaySeconds: 10
+  periodSeconds: 5
+livenessProbe:
+  httpGet:
+    path: /
+    port: 8001
+  initialDelaySeconds: 15
+  periodSeconds: 10
+
+env:
+  REDIS_HOST: "redis-dev-master.redis-dev.svc.cluster.local"
+  REDIS_PORT: "6379"
+  REDIS_PASSWORD: "password"

rnd/infra/helm/autogpt-websocket-server/values.yaml

@@ -0,0 +1,98 @@
+# Default values for autogpt-websocket-server.
+# This is a YAML-formatted file.
+# Declare variables to be passed into your templates.
+
+replicaCount: 1
+
+image:
+  repository: nginx
+  pullPolicy: IfNotPresent
+  # Overrides the image tag whose default is the chart appVersion.
+  tag: ""
+
+imagePullSecrets: []
+nameOverride: ""
+fullnameOverride: ""
+
+serviceAccount:
+  # Specifies whether a service account should be created
+  create: true
+  # Automatically mount a ServiceAccount's API credentials?
+  automount: true
+  # Annotations to add to the service account
+  annotations: {}
+  # The name of the service account to use.
+  # If not set and create is true, a name is generated using the fullname template
+  name: ""
+
+podAnnotations: {}
+podLabels: {}
+
+podSecurityContext: {}
+  # fsGroup: 2000
+
+securityContext: {}
+  # capabilities:
+  #   drop:
+  #   - ALL
+  # readOnlyRootFilesystem: true
+  # runAsNonRoot: true
+  # runAsUser: 1000
+
+service:
+  type: ClusterIP
+  port: 80
+
+ingress:
+  enabled: false
+  className: ""
+  annotations: {}
+    # kubernetes.io/ingress.class: nginx
+    # kubernetes.io/tls-acme: "true"
+  hosts:
+    - host: chart-example.local
+      paths:
+        - path: /
+          pathType: ImplementationSpecific
+  tls: []
+  #  - secretName: chart-example-tls
+  #    hosts:
+  #      - chart-example.local
+
+resources: {}
+  # We usually recommend not to specify default resources and to leave this as a conscious
+  # choice for the user. This also increases chances charts run on environments with little
+  # resources, such as Minikube. If you do want to specify resources, uncomment the following
+  # lines, adjust them as necessary, and remove the curly braces after 'resources:'.
+  # limits:
+  #   cpu: 100m
+  #   memory: 128Mi
+  # requests:
+  #   cpu: 100m
+  #   memory: 128Mi
+
+autoscaling:
+  enabled: false
+  minReplicas: 1
+  maxReplicas: 100
+  targetCPUUtilizationPercentage: 80
+  # targetMemoryUtilizationPercentage: 80
+
+# Additional volumes on the output Deployment definition.
+volumes: []
+# - name: foo
+#   secret:
+#     secretName: mysecret
+#     optional: false
+
+# Additional volumeMounts on the output Deployment definition.
+volumeMounts: []
+# - name: foo
+#   mountPath: "/etc/foo"
+#   readOnly: true
+
+nodeSelector: {}
+
+tolerations: []
+
+affinity: {}

rnd/infra/helm/redis-values.dev.yaml

@@ -0,0 +1,15 @@
+architecture: standalone
+auth:
+  enabled: true
+  password: password
+master:
+  persistence:
+    enabled: true
+    size: 3Gi
+  configmap:
+    redis.conf: |
+      bind 127.0.0.1
+      protected-mode yes
+      requirepass password
+replica:
+  replicaCount: 0