import os import sys import time import logging import signal import threading import requests from urllib.parse import urlparse import redis import psycopg2 # For PostgreSQL # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.StreamHandler(sys.stdout) # Output to console ] ) class Config: """Configuration class to hold environment variables.""" def __init__(self): self.ollama_url = os.environ.get("OLLAMA_URL", "http://localhost:11434") self.vector_store_url = os.environ.get("VECTOR_STORE_URL", "http://localhost:8080") self.api_url = os.environ.get("API_URL", "http://localhost:5000") self.redis_url = os.environ.get("REDIS_URL", "redis://localhost:6379") self.postgres_url = os.environ.get("POSTGRES_URL", "postgresql://user:password@localhost:5432/database") self.startup_timeout = int(os.environ.get("STARTUP_TIMEOUT", "60")) # seconds self.health_check_interval = int(os.environ.get("HEALTH_CHECK_INTERVAL", "10")) # seconds self.metrics_port = int(os.environ.get("METRICS_PORT", "9090")) # Port for metrics endpoint self.status_dashboard_url = os.environ.get("STATUS_DASHBOARD_URL", "http://localhost:3000") # Example dashboard URL self.validate_config() def validate_config(self): """Validates the configuration values.""" try: urlparse(self.ollama_url) urlparse(self.vector_store_url) urlparse(self.api_url) urlparse(self.redis_url) urlparse(self.postgres_url) except Exception as e: logging.error(f"Invalid URL in configuration: {e}") sys.exit(1) if not (0 < self.startup_timeout <= 300): logging.warning(f"STARTUP_TIMEOUT value {self.startup_timeout} is outside reasonable bounds. Consider adjusting.") if not (0 < self.health_check_interval <= 60): logging.warning(f"HEALTH_CHECK_INTERVAL value {self.health_check_interval} is outside reasonable bounds. Consider adjusting.") if not (1024 < self.metrics_port < 65535): logging.warning(f"METRICS_PORT value {self.metrics_port} is outside reasonable bounds. Consider adjusting.") # Global shutdown event shutdown_event = threading.Event() def signal_handler(signum, frame): """Handles signals for graceful shutdown.""" logging.info(f"Received signal {signum}. Shutting down...") shutdown_event.set() def check_redis_availability(redis_url, timeout=10): """Checks if Redis is available.""" try: r = redis.Redis.from_url(redis_url, socket_connect_timeout=timeout) r.ping() logging.info("Redis is available.") return True except redis.exceptions.ConnectionError as e: logging.error(f"Redis connection error: {e}") return False except Exception as e: logging.error(f"Unexpected error checking Redis: {e}") return False def check_postgres_availability(postgres_url, timeout=10): """Checks if PostgreSQL is available.""" try: conn = psycopg2.connect(postgres_url, connect_timeout=timeout) conn.close() logging.info("PostgreSQL is available.") return True except psycopg2.OperationalError as e: logging.error(f"PostgreSQL connection error: {e}") return False except Exception as e: logging.error(f"Unexpected error checking PostgreSQL: {e}") return False def check_ollama_availability(ollama_url, timeout=10): """Checks if Ollama is available.""" try: response = requests.get(ollama_url, timeout=timeout) response.raise_for_status() # Raise HTTPError for bad responses (4xx or 5xx) logging.info("Ollama is available.") return True except requests.exceptions.RequestException as e: logging.error(f"Ollama connection error: {e}") return False except Exception as e: logging.error(f"Unexpected error checking Ollama: {e}") return False def check_vector_store_readiness(vector_store_url, timeout=10): """Checks if the vector store is ready using a /ready endpoint.""" try: ready_url = f"{vector_store_url}/ready" response = requests.get(ready_url, timeout=timeout) response.raise_for_status() if response.status_code == 200: logging.info("Vector store is ready.") return True else: logging.warning(f"Vector store /ready endpoint returned status code: {response.status_code}") return False except requests.exceptions.RequestException as e: logging.error(f"Vector store readiness check failed: {e}") return False except Exception as e: logging.error(f"Unexpected error checking vector store readiness: {e}") return False def load_memory(): """Loads memory into the system. Replace with actual implementation.""" logging.info("Loading memory...") time.sleep(5) # Simulate loading memory logging.info("Memory loaded.") def register_skills(): """Registers skills with the system. Replace with actual implementation.""" logging.info("Registering skills...") time.sleep(3) # Simulate registering skills logging.info("Skills registered.") def initialize_metrics(metrics_port): """Initializes the metrics endpoint.""" logging.info(f"Initializing metrics endpoint on port {metrics_port}...") # In a real application, you would initialize a Prometheus client or similar here. # For example, using the `prometheus_client` library: # from prometheus_client import start_http_server, Summary # start_http_server(metrics_port) logging.info("Metrics endpoint initialized (placeholder).") def check_memory_health(): """Placeholder for memory health check. Replace with actual implementation.""" # Implement checks for memory usage, fragmentation, etc. # Return True if healthy, False otherwise. return True def check_skill_health(): """Placeholder for skill health check. Replace with actual implementation.""" # Implement checks for skill availability, performance, etc. # Return True if healthy, False otherwise. return True def check_api_health(api_url, timeout=5): """Checks the health of the main API endpoint.""" try: response = requests.get(f"{api_url}/health", timeout=timeout) response.raise_for_status() if response.status_code == 200: logging.info("API is healthy.") return True else: logging.warning(f"API /health endpoint returned status code: {response.status_code}") return False except requests.exceptions.RequestException as e: logging.error(f"API health check failed: {e}") return False def get_status_dashboard_data(dashboard_url, timeout=5): """Retrieves data from the status dashboard.""" try: response = requests.get(dashboard_url, timeout=timeout) response.raise_for_status() return response.json() # Assuming the dashboard returns JSON data except requests.exceptions.RequestException as e: logging.error(f"Failed to retrieve data from status dashboard: {e}") return None except Exception as e: logging.error(f"Unexpected error getting status dashboard data: {e}") return None def health_monitoring(config): """Monitors the health of the system components.""" while not shutdown_event.is_set(): try: memory_healthy = check_memory_health() skill_healthy = check_skill_health() api_healthy = check_api_health(config.api_url) if not memory_healthy: logging.error("Memory health check failed.") if not skill_healthy: logging.error("Skill health check failed.") if not api_healthy: logging.error("API health check failed.") if memory_healthy and skill_healthy and api_healthy: logging.info("System is healthy.") else: logging.warning("System is NOT fully healthy.") # Fetch and log status dashboard data dashboard_data = get_status_dashboard_data(config.status_dashboard_url) if dashboard_data: logging.info(f"Status Dashboard Data: {dashboard_data}") else: logging.warning("Failed to retrieve status dashboard data.") except Exception as e: logging.error(f"Error during health monitoring: {e}") time.sleep(config.health_check_interval) logging.info("Health monitoring stopped.") def startup(config): """Performs the startup sequence.""" logging.info("Starting AIVA...") start_time = time.time() while time.time() - start_time < config.startup_timeout: if check_redis_availability(config.redis_url): break logging.warning("Redis not available, retrying...") time.sleep(2) else: logging.error("Timeout waiting for Redis to become available.") sys.exit(1) start_time = time.time() while time.time() - start_time < config.startup_timeout: if check_postgres_availability(config.postgres_url): break logging.warning("PostgreSQL not available, retrying...") time.sleep(2) else: logging.error("Timeout waiting for PostgreSQL to become available.") sys.exit(1) start_time = time.time() while time.time() - start_time < config.startup_timeout: if check_ollama_availability(config.ollama_url): break logging.warning("Ollama not available, retrying...") time.sleep(2) else: logging.error("Timeout waiting for Ollama to become available.") sys.exit(1) start_time = time.time() while time.time() - start_time < config.startup_timeout: if check_vector_store_readiness(config.vector_store_url): break logging.warning("Vector store not ready, retrying...") time.sleep(2) else: logging.error("Timeout waiting for vector store to become ready.") sys.exit(1) load_memory() register_skills() initialize_metrics(config.metrics_port) # Initialize metrics after dependencies are ready logging.info("AIVA started successfully.") def main(): """Main function to start the system.""" config = Config() # Register signal handlers for graceful shutdown signal.signal(signal.SIGINT, signal_handler) signal.signal(signal.SIGTERM, signal_handler) try: startup(config) # Start health monitoring in a separate thread health_thread = threading.Thread(target=health_monitoring, args=(config,), daemon=True) health_thread.start() # Keep the main thread alive until a shutdown signal is received shutdown_event.wait() except Exception as e: logging.error(f"Startup failed: {e}") sys.exit(1) finally: logging.info("Shutting down AIVA...") # Perform any cleanup tasks here (e.g., closing connections) logging.info("AIVA shutdown complete.") sys.exit(0) if __name__ == "__main__": main()