import logging import random import time from typing import List, Dict, Any import psycopg2 import redis from qdrant_client import QdrantClient, models from qdrant_client.models import VectorParams, Distance, PointStruct # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Database connection details DB_HOST = "postgresql-genesis-u50607.vm.elestio.app" DB_PORT = 25432 DB_NAME = "your_db_name" # Replace with your actual database name DB_USER = "your_db_user" # Replace with your actual database user DB_PASSWORD = "your_db_password" # Replace with your actual database password # Redis connection details REDIS_HOST = "redis-genesis-u50607.vm.elestio.app" REDIS_PORT = 26379 # Qdrant connection details QDRANT_HOST = "qdrant-b3knu-u50607.vm.elestio.app" QDRANT_PORT = 6333 # Ollama details OLLAMA_HOST = "localhost" OLLAMA_PORT = 23405 class DataGenerator: """ Generates realistic test data for Genesis system components. """ def __init__(self, volume: int = 100): """ Initializes the DataGenerator with a specified data volume. Args: volume: The number of entities/records to generate. """ self.volume = volume self.db_conn = None self.redis_client = None self.qdrant_client = None def connect_db(self) -> None: """ Connects to the PostgreSQL database. """ try: self.db_conn = psycopg2.connect( host=DB_HOST, port=DB_PORT, database=DB_NAME, user=DB_USER, password=DB_PASSWORD ) self.db_conn.autocommit = True # Enable autocommit for test data generation logging.info("Connected to PostgreSQL database.") except psycopg2.Error as e: logging.error(f"Error connecting to PostgreSQL: {e}") raise def connect_redis(self) -> None: """ Connects to the Redis server. """ try: self.redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, decode_responses=True) self.redis_client.ping() # Check connection logging.info("Connected to Redis.") except redis.exceptions.ConnectionError as e: logging.error(f"Error connecting to Redis: {e}") raise def connect_qdrant(self) -> None: """ Connects to the Qdrant vector database. """ try: self.qdrant_client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) logging.info("Connected to Qdrant.") except Exception as e: logging.error(f"Error connecting to Qdrant: {e}") raise def generate_entities(self) -> List[Dict[str, Any]]: """ Generates a list of entity dictionaries. Returns: A list of entity dictionaries. """ entities = [] for i in range(self.volume): entity = { "id": i + 1, "name": f"Entity {i + 1}", "description": f"This is a test entity number {i + 1}.", "created_at": time.time() } entities.append(entity) logging.info(f"Generated {self.volume} entities.") return entities def generate_relations(self, entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Generates a list of relation dictionaries between entities. Args: entities: A list of entity dictionaries to create relations between. Returns: A list of relation dictionaries. """ relations = [] for i in range(self.volume): source_entity_id = random.choice(entities)["id"] target_entity_id = random.choice(entities)["id"] relation = { "id": i + 1, "source_entity_id": source_entity_id, "target_entity_id": target_entity_id, "type": random.choice(["parent", "child", "sibling", "related"]) } relations.append(relation) logging.info(f"Generated {self.volume} relations.") return relations def generate_tasks(self) -> List[Dict[str, Any]]: """ Generates a list of task dictionaries. Returns: A list of task dictionaries. """ tasks = [] for i in range(self.volume): task = { "id": i + 1, "name": f"Task {i + 1}", "description": f"This is a test task number {i + 1}.", "status": random.choice(["open", "in progress", "completed", "blocked"]), "priority": random.choice(["high", "medium", "low"]), "due_date": time.time() + random.randint(0, 86400 * 30) # Up to 30 days from now } tasks.append(task) logging.info(f"Generated {self.volume} tasks.") return tasks def generate_metrics(self) -> List[Dict[str, Any]]: """ Generates a list of metric dictionaries. Returns: A list of metric dictionaries. """ metrics = [] for i in range(self.volume): metric = { "id": i + 1, "name": f"Metric {i + 1}", "value": random.uniform(0, 100), "timestamp": time.time() } metrics.append(metric) logging.info(f"Generated {self.volume} metrics.") return metrics def seed_db(self, entities: List[Dict[str, Any]], relations: List[Dict[str, Any]], tasks: List[Dict[str, Any]], metrics: List[Dict[str, Any]]) -> None: """ Seeds the PostgreSQL database with generated data. Args: entities: A list of entity dictionaries. relations: A list of relation dictionaries. tasks: A list of task dictionaries. metrics: A list of metric dictionaries. """ try: cursor = self.db_conn.cursor() # Create tables if they don't exist (simplified example) cursor.execute(""" CREATE TABLE IF NOT EXISTS entities ( id INTEGER PRIMARY KEY, name VARCHAR(255), description TEXT, created_at FLOAT ); """) cursor.execute(""" CREATE TABLE IF NOT EXISTS relations ( id INTEGER PRIMARY KEY, source_entity_id INTEGER, target_entity_id INTEGER, type VARCHAR(255) ); """) cursor.execute(""" CREATE TABLE IF NOT EXISTS tasks ( id INTEGER PRIMARY KEY, name VARCHAR(255), description TEXT, status VARCHAR(255), priority VARCHAR(255), due_date FLOAT ); """) cursor.execute(""" CREATE TABLE IF NOT EXISTS metrics ( id INTEGER PRIMARY KEY, name VARCHAR(255), value FLOAT, timestamp FLOAT ); """) # Insert data into tables for entity in entities: cursor.execute(""" INSERT INTO entities (id, name, description, created_at) VALUES (%s, %s, %s, %s); """, (entity["id"], entity["name"], entity["description"], entity["created_at"])) for relation in relations: cursor.execute(""" INSERT INTO relations (id, source_entity_id, target_entity_id, type) VALUES (%s, %s, %s, %s); """, (relation["id"], relation["source_entity_id"], relation["target_entity_id"], relation["type"])) for task in tasks: cursor.execute(""" INSERT INTO tasks (id, name, description, status, priority, due_date) VALUES (%s, %s, %s, %s, %s, %s); """, (task["id"], task["name"], task["description"], task["status"], task["priority"], task["due_date"])) for metric in metrics: cursor.execute(""" INSERT INTO metrics (id, name, value, timestamp) VALUES (%s, %s, %s, %s); """, (metric["id"], metric["name"], metric["value"], metric["timestamp"])) self.db_conn.commit() logging.info("Seeded PostgreSQL database.") except psycopg2.Error as e: logging.error(f"Error seeding PostgreSQL database: {e}") raise finally: if cursor: cursor.close() def seed_redis(self, entities: List[Dict[str, Any]]) -> None: """ Seeds the Redis server with generated data. Args: entities: A list of entity dictionaries. """ try: for entity in entities: key = f"entity:{entity['id']}" self.redis_client.hmset(key, entity) # Use hmset for dictionaries logging.info("Seeded Redis.") except redis.exceptions.RedisError as e: logging.error(f"Error seeding Redis: {e}") raise def seed_qdrant(self, entities: List[Dict[str, Any]]) -> None: """ Seeds the Qdrant vector database with generated data. Args: entities: A list of entity dictionaries. """ collection_name = "test_entities" vector_size = 128 # Example vector size try: # Recreate collection if it exists try: self.qdrant_client.recreate_collection( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE) ) except Exception as e: logging.warning(f"Failed to recreate collection, assuming it does not exist: {e}") self.qdrant_client.create_collection( collection_name=collection_name, vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE) ) points = [] for entity in entities: # Generate a random vector (replace with actual embedding generation logic) vector = [random.random() for _ in range(vector_size)] points.append( PointStruct( id=entity["id"], vector=vector, payload=entity # Store the entity data as payload ) ) self.qdrant_client.upsert( collection_name=collection_name, points=points, wait=True # Wait for the operation to complete ) logging.info("Seeded Qdrant.") except Exception as e: logging.error(f"Error seeding Qdrant: {e}") raise def run(self) -> None: """ Runs the data generation and seeding process. """ try: logging.info("Starting data generation...") entities = self.generate_entities() relations = self.generate_relations(entities) tasks = self.generate_tasks() metrics = self.generate_metrics() self.connect_db() self.seed_db(entities, relations, tasks, metrics) self.connect_redis() self.seed_redis(entities) self.connect_qdrant() self.seed_qdrant(entities) logging.info("Data generation and seeding completed successfully.") except Exception as e: logging.error(f"Data generation failed: {e}") raise finally: if self.db_conn: self.db_conn.close() logging.info("Disconnected from PostgreSQL.") if self.redis_client: self.redis_client.close() logging.info("Disconnected from Redis.") if self.qdrant_client: self.qdrant_client.close() logging.info("Disconnected from Qdrant.") if __name__ == "__main__": generator = DataGenerator(volume=10) # Generate 10 test records generator.run()