diff --git a/codetide/agents/data_layer.py b/codetide/agents/data_layer.py index e63b6e0..c7d8ed3 100644 --- a/codetide/agents/data_layer.py +++ b/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 diff --git a/codetide/agents/tide/ui/app.py b/codetide/agents/tide/ui/app.py index 2c4f7a6..b4597b8 100644 --- a/codetide/agents/tide/ui/app.py +++ b/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, diff --git a/codetide/agents/tide/ui/utils.py b/codetide/agents/tide/ui/utils.py index 8b4e4f1..0b634ba 100644 --- a/codetide/agents/tide/ui/utils.py +++ b/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 \ No newline at end of file + 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) diff --git a/requirements-agents-ui.txt b/requirements-agents-ui.txt index e3c7f85..42830eb 100644 --- a/requirements-agents-ui.txt +++ b/requirements-agents-ui.txt @@ -1,3 +1,4 @@ chainlit==2.6.3 SQLAlchemy==2.0.36 -aiosqlite==0.21.0 \ No newline at end of file +asyncpg==0.30.0 +docker==7.1.0 \ No newline at end of file