Merge pull request #35 from BrunoV21/pg-data-layer

Pg data layer

Author: BrunoV21

codetide/agents/data_layer.py

@@ -3,16 +3,16 @@
     from sqlalchemy import String, Text, ForeignKey, Boolean, Integer
     from sqlalchemy.ext.asyncio import create_async_engine
     from sqlalchemy.types import TypeDecorator
+    from sqlalchemy.exc import OperationalError
 except ImportError as e:
     raise ImportError(
         "This module requires 'sqlalchemy' and 'ulid-py'. "
         "Install them with: pip install codetide[agents-ui]"
-    ) from e
+    ) from e   
 
+import asyncio
 from datetime import datetime
-from sqlalchemy import Select
 from ulid import ulid
-import asyncio
 import json
 
 # SQLite-compatible JSON and UUID types
@@ -169,11 +169,27 @@ class Feedback(Base):
 #     chats = await db.list_chats()
 #     for c in chats:
 #         print(f"{c.id} — {c.name}")
-async def init_db(path: str):
-    from sqlalchemy.ext.asyncio import async_sessionmaker, AsyncSession
-    engine = create_async_engine(f"sqlite+aiosqlite:///{path}")
-    async with engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
-
-if __name__ == "__main__":
-    asyncio.run(init_db("database.db"))
+
+async def init_db(conn_str: str, max_retries: int = 5, retry_delay: int = 2):
+    """
+    Initialize database with retry logic for connection issues.
+    """
+    engine = create_async_engine(conn_str)
+    
+    for attempt in range(max_retries):
+        try:
+            async with engine.begin() as conn:
+                await conn.run_sync(Base.metadata.create_all)
+            print("Database initialized successfully!")
+            return
+        except OperationalError as e:
+            if attempt == max_retries - 1:
+                print(f"Failed to initialize database after {max_retries} attempts: {e}")
+                raise
+            else:
+                print(f"Database connection failed (attempt {attempt + 1}/{max_retries}): {e}")
+                print(f"Retrying in {retry_delay} seconds...")
+                await asyncio.sleep(retry_delay)
+        except Exception as e:
+            print(f"Unexpected error initializing database: {e}")
+        raise

codetide/agents/tide/ui/app.py

@@ -9,8 +9,8 @@
     from aicore.config import Config
     from aicore.llm import Llm, LlmConfig
     from aicore.models import AuthenticationError, ModelError
-    from aicore.const import STREAM_END_TOKEN, STREAM_START_TOKEN#, REASONING_START_TOKEN, REASONING_STOP_TOKEN    
-    from codetide.agents.tide.ui.utils import process_thread, run_concurrent_tasks, send_reasoning_msg
+    from aicore.const import STREAM_END_TOKEN, STREAM_START_TOKEN#, REASONING_START_TOKEN, REASONING_STOP_TOKEN
+    from codetide.agents.tide.ui.utils import process_thread, run_concurrent_tasks, send_reasoning_msg, check_docker, launch_postgres
     from codetide.agents.tide.ui.stream_processor import StreamProcessor, MarkerConfig
     from codetide.agents.tide.ui.defaults import AGENT_TIDE_PORT, STARTERS
     from codetide.agents.tide.ui.agent_tide_ui import AgentTideUi
@@ -24,29 +24,33 @@
 except ImportError as e:
     raise ImportError(
         "The 'codetide.agents' module requires the 'aicore' and 'chainlit' packages. "
-        "Install it with: pip install codetide[agents-ui]"
+        "Install it with: pip install codetide[aasygents-ui]"
     ) from e
 
 from codetide.agents.tide.ui.defaults import AICORE_CONFIG_EXAMPLE, EXCEPTION_MESSAGE, MISSING_CONFIG_MESSAGE
 from codetide.agents.tide.defaults import DEFAULT_AGENT_TIDE_LLM_CONFIG_PATH
 from codetide.core.defaults import DEFAULT_ENCODING
+from dotenv import get_key, load_dotenv, set_key
 from codetide.agents.data_layer import init_db
 from ulid import ulid
 import argparse
 import getpass
 import asyncio
+import secrets
+import string
 import json
 import yaml
 import time
 
-@cl.password_auth_callback
-def auth():
-    username = getpass.getuser()
-    return cl.User(identifier=username, display_name=username)
+if check_docker and os.getenv("AGENTTIDE_PG_CONN_STR") is not None:
+    @cl.password_auth_callback
+    def auth():
+        username = getpass.getuser()
+        return cl.User(identifier=username, display_name=username)
 
-@cl.data_layer
-def get_data_layer():
-    return SQLAlchemyDataLayer(conninfo=f"sqlite+aiosqlite:///{os.environ['CHAINLIT_APP_ROOT']}/database.db")
+    @cl.data_layer
+    def get_data_layer():
+        return SQLAlchemyDataLayer(conninfo=os.getenv("AGENTTIDE_PG_CONN_STR"))
 
 @cl.on_settings_update
 async def setup_llm_config(settings):
@@ -442,9 +446,18 @@ async def agent_loop(message: Optional[cl.Message]=None, codeIdentifiers: Option
                 chat_history.append({"role": "user", "content": feedback})
                 await agent_loop(agent_tide_ui=agent_tide_ui)
 
-# def generate_temp_password(length=16):
-#     characters = string.ascii_letters + string.digits + string.punctuation
-#     return ''.join(secrets.choice(characters) for _ in range(length))
+def generate_password(length: int = 16) -> str:
+    """
+    Generate a secure random password.
+    Works on Linux, macOS, and Windows.
+    """
+    if password  := get_key(Path(os.environ['CHAINLIT_APP_ROOT']) / ".env", "AGENTTDE_PG_PASSWORD"):
+        return password
+    
+    safe_chars = string.ascii_letters + string.digits + '-_@#$%^&*+=[]{}|:;<>?'
+    password = ''.join(secrets.choice(safe_chars) for _ in range(length))
+    set_key(Path(os.environ['CHAINLIT_APP_ROOT']) / ".env","AGENTTDE_PG_PASSWORD", password)
+    return password
 
 def serve(
     host=None,
@@ -454,14 +467,7 @@ def serve(
     ssl_keyfile=None,
     ws_per_message_deflate="true",
     ws_protocol="auto"
-):    
-    username = getpass.getuser()    
-    GREEN = "\033[92m"
-    RESET = "\033[0m"
-
-    print(f"\n{GREEN}Your chainlit username is `{username}`{RESET}\n")
-
-
+):
     # if not os.getenv("_PASSWORD"):
     #     temp_password = generate_temp_password()
     #     os.environ["_PASSWORD"] = temp_password
@@ -513,10 +519,33 @@ def main():
     parser.add_argument("--config-path", type=str, default=DEFAULT_AGENT_TIDE_LLM_CONFIG_PATH, help="Path to the config file")
     args = parser.parse_args()
 
+    load_dotenv()
     os.environ["AGENT_TIDE_PROJECT_PATH"] = str(Path(args.project_path))
     os.environ["AGENT_TIDE_CONFIG_PATH"] = str(Path(args.project_path) / args.config_path)
+    
+    load_dotenv()
+    username = getpass.getuser()
+    GREEN = "\033[92m"
+    RED = "\033[91m"
+    RESET = "\033[0m"
+
+    print(f"\n{GREEN}Your chainlit username is `{username}`{RESET}\n")
+
+    if check_docker():
+        password = generate_password()
+        launch_postgres(username, password, f"{os.environ['CHAINLIT_APP_ROOT']}/pgdata")
+    
+        conn_string = f"postgresql+asyncpg://{username}:{password}@localhost:{os.getenv('AGENTTIDE_PG_PORT', 5437)}/agenttidedb"
+        os.environ["AGENTTIDE_PG_CONN_STR"] = conn_string
+        asyncio.run(init_db(os.environ["AGENTTIDE_PG_CONN_STR"]))
 
-    asyncio.run(init_db(f"{os.environ['CHAINLIT_APP_ROOT']}/database.db"))
+        print(f"{GREEN} PostgreSQL launched on port {os.getenv('AGENTTIDE_PG_PORT', 5437)}{RESET}")
+        print(f"{GREEN} Connection string stored in env var: AGENTTIDE_PG_CONN_STR{RESET}\n")
+    else:
+        print(f"{RED} Could not find Docker on this system.{RESET}")
+        print("   PostgreSQL could not be launched for persistent data storage.")
+        print("   You won't have access to multiple conversations or history beyond each session.")
+        print("   Consider installing Docker and ensuring it is running.\n")
 
     serve(
         host=args.host,

codetide/agents/tide/ui/utils.py

@@ -2,12 +2,16 @@
 
 from typing import List, Optional, Tuple
 from chainlit.types import ThreadDict
+from rich.progress import Progress
 from aicore.logger import _logger
 from aicore.llm import LlmConfig
 import chainlit as cl
 import asyncio
 import orjson
+import docker
 import time
+import os
+
 
 def process_thread(thread :ThreadDict)->Tuple[List[dict], Optional[LlmConfig], str]:
     ### type: tool
@@ -93,4 +97,110 @@ async def send_reasoning_msg(loading_msg :cl.message, context_msg :cl.Message, a
         )
     )
     await context_msg.send()
-    return True
+    return True
+
+def check_docker():
+    try:
+        client = docker.from_env()
+        client.ping()  # Simple API check
+        return True
+    except Exception:
+        return False
+    
+tasks = {}
+
+# Show task progress (red for download, green for extract)
+def show_progress(line, progress):
+    if line['status'] == 'Downloading':
+        id = f'[red][Download {line["id"]}]'
+    elif line['status'] == 'Extracting':
+        id = f'[green][Extract  {line["id"]}]'
+    else:
+        # skip other statuses
+        return
+
+    if id not in tasks.keys():
+        tasks[id] = progress.add_task(f"{id}", total=line['progressDetail']['total'])
+    else:
+        progress.update(tasks[id], completed=line['progressDetail']['current'])
+
+def image_pull(client :docker.DockerClient, image_name):
+    print(f'Pulling image: {image_name}')
+    with Progress() as progress:
+        resp = client.api.pull(image_name, stream=True, decode=True)
+        for line in resp:
+            show_progress(line, progress)
+
+def wait_for_postgres_ready(container, username: str, password: str, max_attempts: int = 30, delay: int = 2) -> bool:
+    """
+    Wait for PostgreSQL to be ready by checking container logs and attempting connections.
+    """
+    print("Waiting for PostgreSQL to be ready...")
+    
+    for attempt in range(max_attempts):
+        try:
+            # First, check if container is still running
+            container.reload()
+            if container.status != "running":
+                print(f"Container stopped unexpectedly. Status: {container.status}")
+                return False
+            
+            # Check logs for readiness indicator
+            logs = container.logs().decode('utf-8')
+            if "database system is ready to accept connections" in logs:
+                print("PostgreSQL is ready to accept connections!")
+                # Give it one more second to be completely ready
+                time.sleep(5)
+                return True
+                
+            print(f"Attempt {attempt + 1}/{max_attempts}: PostgreSQL not ready yet...")
+            time.sleep(delay)
+            
+        except Exception as e:
+            print(f"Error checking PostgreSQL readiness: {e}")
+            time.sleep(delay)
+    
+    print("Timeout waiting for PostgreSQL to be ready")
+    return False
+
+def launch_postgres(POSTGRES_USER: str, POSTGRES_PASSWORD: str, volume_path: str):
+    client = docker.from_env()
+    container_name = "agent-tide-postgres"
+
+    # Check if the container already exists
+    try:
+        container = client.containers.get(container_name)
+        status = container.status
+        print(f"Container '{container_name}' status: {status}")
+        if status == "running":
+            print("Container is already running. No need to relaunch.")
+            return
+        else:
+            print("Container exists but is not running. Starting container...")
+            container.start()
+            return
+    except docker.errors.NotFound:
+        # Container does not exist, we need to create it
+        print("Container does not exist. Launching a new one...")
+
+
+    image_pull(client, "postgres:alpine")
+    print("Image pulled successfully")
+    # Launch a new container
+    container = client.containers.run(
+        "postgres:alpine",
+        name=container_name,
+        environment={
+            "POSTGRES_USER": POSTGRES_USER,
+            "POSTGRES_PASSWORD": POSTGRES_PASSWORD,
+            "POSTGRES_DB": "agenttidedb"
+        },
+        ports={"5432/tcp": os.getenv('AGENTTIDE_PG_PORT', 5437)},
+        volumes={volume_path: {"bind": "/var/lib/postgresql/data", "mode": "rw"}},
+        detach=True,
+        restart_policy={"Name": "always"}
+    )
+
+    print(f"Container '{container_name}' launched successfully with status: {container.status}")
+    # Wait for PostgreSQL to be ready
+    return wait_for_postgres_ready(container, POSTGRES_USER, POSTGRES_PASSWORD)

requirements-agents-ui.txt

@@ -1,3 +1,4 @@
 chainlit==2.6.3
 SQLAlchemy==2.0.36
-aiosqlite==0.21.0
+asyncpg==0.30.0
+docker==7.1.0