import schedule import time import logging import threading from typing import Dict, List, Callable, Optional import datetime import psycopg2 import redis import qdrant_client import requests # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Constants for database, redis, qdrant, and ollama POSTGRES_HOST = "postgresql-genesis-u50607.vm.elestio.app" POSTGRES_PORT = 25432 POSTGRES_DB = "your_db_name" # Replace with your actual database name POSTGRES_USER = "your_db_user" # Replace with your actual database user POSTGRES_PASSWORD = "your_db_password" # Replace with your actual database password REDIS_HOST = "redis-genesis-u50607.vm.elestio.app" REDIS_PORT = 26379 QDRANT_HOST = "qdrant-b3knu-u50607.vm.elestio.app" QDRANT_PORT = 6333 OLLAMA_HOST = "localhost" OLLAMA_PORT = 23405 class HealthCheckResult: """Represents the result of a health check.""" def __init__(self, component: str, status: bool, message: Optional[str] = None): self.component = component self.status = status self.message = message self.timestamp = datetime.datetime.now() def __repr__(self): return f"HealthCheckResult(component='{self.component}', status={self.status}, message='{self.message}', timestamp={self.timestamp})" class HealthChecker: """ A class to perform health checks on various components of the Genesis system. """ def __init__(self): self.component_checks: Dict[str, Callable[[], HealthCheckResult]] = { "postgres": self.check_postgres, "redis": self.check_redis, "qdrant": self.check_qdrant, "ollama": self.check_ollama, } def check_postgres(self) -> HealthCheckResult: """Checks the health of the PostgreSQL database.""" try: conn = psycopg2.connect( host=POSTGRES_HOST, port=POSTGRES_PORT, database=POSTGRES_DB, user=POSTGRES_USER, password=POSTGRES_PASSWORD, connect_timeout=5 # Timeout in seconds ) conn.close() return HealthCheckResult("postgres", True) except Exception as e: logging.error(f"Postgres health check failed: {e}") return HealthCheckResult("postgres", False, str(e)) def check_redis(self) -> HealthCheckResult: """Checks the health of the Redis server.""" try: r = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=0, socket_connect_timeout=5) if r.ping(): return HealthCheckResult("redis", True) else: return HealthCheckResult("redis", False, "Ping failed") except Exception as e: logging.error(f"Redis health check failed: {e}") return HealthCheckResult("redis", False, str(e)) def check_qdrant(self) -> HealthCheckResult: """Checks the health of the Qdrant vector database.""" try: client = qdrant_client.QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT, timeout=5) client.get_telemetry_data() # Attempt to get telemetry data as a health check return HealthCheckResult("qdrant", True) except Exception as e: logging.error(f"Qdrant health check failed: {e}") return HealthCheckResult("qdrant", False, str(e)) def check_ollama(self) -> HealthCheckResult: """Checks the health of the AIVA Ollama server.""" try: response = requests.get(f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/", timeout=5) if response.status_code == 200: return HealthCheckResult("ollama", True) else: return HealthCheckResult("ollama", False, f"HTTP Status: {response.status_code}") except requests.exceptions.RequestException as e: logging.error(f"Ollama health check failed: {e}") return HealthCheckResult("ollama", False, str(e)) def run_health_check(self, component: str) -> HealthCheckResult: """Runs the health check for a specific component.""" if component in self.component_checks: return self.component_checks[component]() else: logging.warning(f"No health check defined for component: {component}") return HealthCheckResult(component, False, "No health check defined") class HealthScheduler: """ A class to schedule and run health checks periodically, aggregate results, and trigger alerts on degradation. """ def __init__(self, config: Dict[str, int], alert_threshold: float = 0.7): """ Initializes the HealthScheduler with a configuration and an alert threshold. Args: config (Dict[str, int]): A dictionary where keys are component names and values are the interval in seconds. alert_threshold (float): A float between 0 and 1 representing the minimum health score to avoid triggering an alert. """ self.config = config self.health_checker = HealthChecker() self.health_results: List[HealthCheckResult] = [] self.alert_threshold = alert_threshold self.scheduler = schedule.Scheduler() self.setup_schedule() def setup_schedule(self): """Sets up the schedule for health checks based on the configuration.""" for component, interval in self.config.items(): self.scheduler.every(interval).seconds.do(self.run_and_store_check, component=component) def run_and_store_check(self, component: str): """Runs a health check for a component and stores the result.""" result = self.health_checker.run_health_check(component) self.health_results.append(result) logging.info(f"Health check for {component}: {result}") self.check_and_trigger_alerts() def aggregate_health_score(self) -> float: """Aggregates the health results into a unified health score.""" if not self.health_results: return 1.0 # Assume healthy if no results yet healthy_count = sum(1 for result in self.health_results if result.status) return healthy_count / len(self.health_results) def check_and_trigger_alerts(self): """Checks the aggregated health score and triggers alerts if it falls below the threshold.""" health_score = self.aggregate_health_score() logging.info(f"Aggregated health score: {health_score}") if health_score < self.alert_threshold: self.trigger_alert(health_score) def trigger_alert(self, health_score: float): """Triggers an alert (e.g., sends a notification) when the health score is low.""" logging.warning(f"Health score is below threshold! Current score: {health_score}") # TODO: Implement alert mechanism (e.g., send email, push notification) print(f"ALERT: Health score is below threshold! Current score: {health_score}") def run_scheduler(self): """Runs the scheduler in a loop.""" while True: self.scheduler.run_pending() time.sleep(1) def main(): """Main function to start the health scheduler.""" config = { "postgres": 60, # Check every 60 seconds "redis": 30, # Check every 30 seconds "qdrant": 120, # Check every 120 seconds "ollama": 45, # Check every 45 seconds } scheduler = HealthScheduler(config) # Run the scheduler in a separate thread to avoid blocking the main thread scheduler_thread = threading.Thread(target=scheduler.run_scheduler) scheduler_thread.daemon = True # Allow the main thread to exit even if the scheduler thread is running scheduler_thread.start() # Keep the main thread alive (e.g., for other tasks) try: while True: time.sleep(60) # Check every minute if the program should exit except KeyboardInterrupt: logging.info("Exiting health scheduler.") if __name__ == "__main__": main()