tasks/postgres/standard/vectors/vectors_setup.py

"""
Shared Vector Database Setup Utilities

This module provides utilities for setting up a complete PostgreSQL database
with pgvector extension and sample RAG-related tables with vector data.
Used by all vector database tasks.
"""

import os
import logging
import psycopg2
import json
import random
import numpy as np
from typing import List

logger = logging.getLogger(__name__)

def get_connection_params():
    """Get database connection parameters from environment variables."""
    return {
        'host': os.getenv('POSTGRES_HOST', 'localhost'),
        'port': os.getenv('POSTGRES_PORT', '5432'),
        'user': os.getenv('POSTGRES_USERNAME', 'postgres'),
        'password': os.getenv('POSTGRES_PASSWORD', 'password'),
        'database': os.getenv('POSTGRES_DATABASE', 'postgres')
    }


def generate_mock_embedding(dimensions: int = 1536) -> List[float]:
    """Generate a mock embedding vector with specified dimensions."""
    # Generate random values between -1 and 1, then normalize
    vector = np.random.uniform(-1, 1, dimensions)
    # Normalize to unit vector (common practice for embeddings)
    norm = np.linalg.norm(vector)
    if norm > 0:
        vector = vector / norm
    return vector.tolist()


def create_vector_extension():
    """Create the pgvector extension."""
    conn_params = get_connection_params()

    try:
        conn = psycopg2.connect(**conn_params)
        conn.autocommit = True

        with conn.cursor() as cur:
            logger.info("Creating pgvector extension...")
            cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            logger.info("pgvector extension created successfully")

        conn.close()

    except psycopg2.Error as e:
        logger.error(f"Failed to create pgvector extension: {e}")
        raise


def create_vector_tables():
    """Create sample tables with vector columns for RAG applications."""
    conn_params = get_connection_params()

    try:
        conn = psycopg2.connect(**conn_params)
        conn.autocommit = True

        with conn.cursor() as cur:
            logger.info("Creating vector database tables...")

            # Create documents table for document embeddings
            cur.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id SERIAL PRIMARY KEY,
                    title TEXT NOT NULL,
                    content TEXT NOT NULL,
                    source_url TEXT,
                    document_type VARCHAR(50) DEFAULT 'article',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    word_count INTEGER,
                    embedding vector(1536)
                );
            """)

            # Create chunks table for document chunks (common in RAG)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS document_chunks (
                    id SERIAL PRIMARY KEY,
                    document_id INTEGER REFERENCES documents(id) ON DELETE CASCADE,
                    chunk_index INTEGER NOT NULL,
                    chunk_text TEXT NOT NULL,
                    chunk_size INTEGER,
                    overlap_size INTEGER DEFAULT 0,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    embedding vector(1536)
                );
            """)

            # Create queries table for storing user queries and their embeddings
            cur.execute("""
                CREATE TABLE IF NOT EXISTS user_queries (
                    id SERIAL PRIMARY KEY,
                    query_text TEXT NOT NULL,
                    user_id VARCHAR(100),
                    session_id VARCHAR(100),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    response_time_ms INTEGER,
                    embedding vector(1536)
                );
            """)

            # Create embeddings metadata table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS embedding_models (
                    id SERIAL PRIMARY KEY,
                    model_name VARCHAR(100) NOT NULL UNIQUE,
                    provider VARCHAR(50) NOT NULL,
                    dimensions INTEGER NOT NULL,
                    max_tokens INTEGER,
                    cost_per_token DECIMAL(10, 8),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    is_active BOOLEAN DEFAULT TRUE
                );
            """)

            # Create knowledge base table
            cur.execute("""
                CREATE TABLE IF NOT EXISTS knowledge_base (
                    id SERIAL PRIMARY KEY,
                    kb_name VARCHAR(100) NOT NULL,
                    description TEXT,
                    domain VARCHAR(50),
                    language VARCHAR(10) DEFAULT 'en',
                    total_documents INTEGER DEFAULT 0,
                    total_chunks INTEGER DEFAULT 0,
                    total_storage_mb DECIMAL(10, 2),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)

            # Create similarity search results cache
            cur.execute("""
                CREATE TABLE IF NOT EXISTS search_cache (
                    id SERIAL PRIMARY KEY,
                    query_hash VARCHAR(64) NOT NULL,
                    query_text TEXT NOT NULL,
                    results_json JSONB,
                    result_count INTEGER,
                    search_time_ms INTEGER,
                    similarity_threshold DECIMAL(4, 3),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    expires_at TIMESTAMP
                );
            """)

            logger.info("Vector database tables created successfully")

        conn.close()

    except psycopg2.Error as e:
        logger.error(f"Failed to create vector tables: {e}")
        raise


def create_vector_indexes():
    """Create indexes for vector columns and other frequently queried fields."""
    conn_params = get_connection_params()

    try:
        conn = psycopg2.connect(**conn_params)
        conn.autocommit = True

        with conn.cursor() as cur:
            logger.info("Creating vector indexes...")

            # Vector indexes using HNSW (Hierarchical Navigable Small World)
            indexes = [
                ("documents_embedding_idx", "documents", "embedding", "hnsw"),
                ("chunks_embedding_idx", "document_chunks", "embedding", "hnsw"),
                ("queries_embedding_idx", "user_queries", "embedding", "hnsw"),
            ]

            for idx_name, table_name, column_name, method in indexes:
                try:
                    if method == "hnsw":
                        cur.execute(f"""
                            CREATE INDEX IF NOT EXISTS {idx_name}
                            ON {table_name} USING hnsw ({column_name} vector_cosine_ops);
                        """)
                    else:
                        cur.execute(f"""
                            CREATE INDEX IF NOT EXISTS {idx_name}
                            ON {table_name} USING ivfflat ({column_name} vector_cosine_ops) WITH (lists = 100);
                        """)
                    logger.info(f"Created index {idx_name} on {table_name}")
                except psycopg2.Error as e:
                    logger.warning(f"Could not create {method} index {idx_name}: {e}")
                    # Try with IVFFlat as fallback
                    if method == "hnsw":
                        try:
                            cur.execute(f"""
                                CREATE INDEX IF NOT EXISTS {idx_name}_ivf
                                ON {table_name} USING ivfflat ({column_name} vector_cosine_ops) WITH (lists = 100);
                            """)
                            logger.info(f"Created fallback IVFFlat index {idx_name}_ivf on {table_name}")
                        except psycopg2.Error as e2:
                            logger.warning(f"Could not create fallback index: {e2}")

            # Regular indexes for performance
            regular_indexes = [
                ("documents_title_idx", "documents", "title"),
                ("documents_type_idx", "documents", "document_type"),
                ("documents_created_idx", "documents", "created_at"),
                ("chunks_doc_id_idx", "document_chunks", "document_id"),
                ("chunks_index_idx", "document_chunks", "chunk_index"),
                ("queries_user_idx", "user_queries", "user_id"),
                ("queries_created_idx", "user_queries", "created_at"),
                ("cache_hash_idx", "search_cache", "query_hash"),
                ("cache_expires_idx", "search_cache", "expires_at"),
            ]

            for idx_name, table_name, column_name in regular_indexes:
                try:
                    cur.execute(f"CREATE INDEX IF NOT EXISTS {idx_name} ON {table_name} ({column_name});")
                    logger.debug(f"Created regular index {idx_name}")
                except psycopg2.Error as e:
                    logger.warning(f"Could not create regular index {idx_name}: {e}")

            logger.info("Vector indexes created successfully")

        conn.close()

    except psycopg2.Error as e:
        logger.error(f"Failed to create vector indexes: {e}")
        raise


def insert_sample_data():
    """Insert sample data into vector tables."""
    conn_params = get_connection_params()

    try:
        conn = psycopg2.connect(**conn_params)
        conn.autocommit = True

        with conn.cursor() as cur:
            logger.info("Inserting sample data...")

            # Insert embedding models
            embedding_models = [
                ('text-embedding-3-small', 'OpenAI', 1536, 8192, 0.00000002, True),
                ('text-embedding-3-large', 'OpenAI', 3072, 8192, 0.00000013, True),
                ('text-embedding-ada-002', 'OpenAI', 1536, 8192, 0.00000010, False),
                ('all-MiniLM-L6-v2', 'Sentence-Transformers', 384, 512, 0.0, True),
                ('all-mpnet-base-v2', 'Sentence-Transformers', 768, 514, 0.0, True),
            ]

            for model_data in embedding_models:
                cur.execute("""
                    INSERT INTO embedding_models (model_name, provider, dimensions, max_tokens, cost_per_token, is_active)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    ON CONFLICT (model_name) DO NOTHING;
                """, model_data)

            # Insert knowledge bases
            knowledge_bases = [
                ('Technical Documentation', 'Software engineering and API documentation', 'technology'),
                ('Research Papers', 'Academic papers and research publications', 'research'),
                ('Customer Support', 'FAQ and troubleshooting guides', 'support'),
                ('Product Catalog', 'Product descriptions and specifications', 'commerce'),
                ('Legal Documents', 'Contracts, policies, and legal texts', 'legal'),
            ]

            kb_ids = []
            for kb_data in knowledge_bases:
                cur.execute("""
                    INSERT INTO knowledge_base (kb_name, description, domain, total_documents, total_chunks, total_storage_mb)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    RETURNING id;
                """, kb_data + (random.randint(50, 500), random.randint(200, 2000), round(random.uniform(10.5, 250.8), 2)))
                kb_ids.append(cur.fetchone()[0])

            # Insert sample documents
            sample_documents = [
                ("PostgreSQL Performance Tuning", "Comprehensive guide to optimizing PostgreSQL database performance including indexing strategies, query optimization, and configuration tuning.", "https://example.com/pg-performance", "technical_guide"),
                ("Vector Similarity Search", "Understanding vector embeddings and similarity search algorithms for AI applications and recommendation systems.", "https://example.com/vector-search", "technical_guide"),
                ("RAG Implementation Best Practices", "Best practices for implementing Retrieval-Augmented Generation systems using vector databases and large language models.", "https://example.com/rag-practices", "best_practices"),
                ("Database Security Guidelines", "Security considerations and implementation guidelines for PostgreSQL databases in production environments.", "https://example.com/db-security", "security_guide"),
                ("Machine Learning with SQL", "Integrating machine learning workflows with SQL databases and leveraging database extensions for AI applications.", "https://example.com/ml-sql", "tutorial"),
                ("API Documentation Standards", "Standards and best practices for creating comprehensive and user-friendly API documentation.", "https://example.com/api-docs", "documentation"),
                ("Microservices Architecture", "Design patterns and implementation strategies for microservices architecture in modern applications.", "https://example.com/microservices", "architecture_guide"),
                ("Data Pipeline Optimization", "Optimizing data processing pipelines for scalability, reliability, and performance in enterprise environments.", "https://example.com/data-pipelines", "optimization_guide"),
                ("Cloud Database Migration", "Step-by-step guide for migrating on-premises databases to cloud infrastructure with minimal downtime.", "https://example.com/cloud-migration", "migration_guide"),
                ("NoSQL vs SQL Comparison", "Detailed comparison of NoSQL and SQL databases, including use cases, performance characteristics, and selection criteria.", "https://example.com/nosql-sql", "comparison_guide"),
            ]

            doc_ids = []
            for title, content, url, doc_type in sample_documents:
                embedding = generate_mock_embedding(1536)
                word_count = len(content.split())

                cur.execute("""
                    INSERT INTO documents (title, content, source_url, document_type, word_count, embedding)
                    VALUES (%s, %s, %s, %s, %s, %s)
                    RETURNING id;
                """, (title, content, url, doc_type, word_count, embedding))
                doc_ids.append(cur.fetchone()[0])

            # Insert document chunks
            chunk_count = 0
            for doc_id in doc_ids:
                # Generate 3-7 chunks per document
                num_chunks = random.randint(3, 7)
                for chunk_idx in range(num_chunks):
                    chunk_text = f"This is chunk {chunk_idx + 1} of document {doc_id}. " + \
                               "It contains relevant information that would be useful for similarity search and RAG applications. " + \
                               "The content includes technical details, examples, and best practices."
                    chunk_size = len(chunk_text)
                    overlap_size = random.randint(20, 50) if chunk_idx > 0 else 0
                    embedding = generate_mock_embedding(1536)

                    cur.execute("""
                        INSERT INTO document_chunks (document_id, chunk_index, chunk_text, chunk_size, overlap_size, embedding)
                        VALUES (%s, %s, %s, %s, %s, %s);
                    """, (doc_id, chunk_idx, chunk_text, chunk_size, overlap_size, embedding))
                    chunk_count += 1

            # Insert sample user queries
            sample_queries = [
                ("How to optimize PostgreSQL performance?", "user123", "session_abc1"),
                ("What are vector embeddings?", "user456", "session_def2"),
                ("Best practices for RAG implementation", "user789", "session_ghi3"),
                ("Database security checklist", "user123", "session_abc2"),
                ("Machine learning with databases", "user456", "session_def3"),
                ("API documentation examples", "user321", "session_jkl1"),
                ("Microservices design patterns", "user654", "session_mno2"),
                ("Data pipeline best practices", "user987", "session_pqr3"),
                ("Cloud migration strategies", "user111", "session_stu4"),
                ("NoSQL vs SQL databases", "user222", "session_vwx5"),
            ]

            for query_text, user_id, session_id in sample_queries:
                embedding = generate_mock_embedding(1536)
                response_time = random.randint(50, 500)

                cur.execute("""
                    INSERT INTO user_queries (query_text, user_id, session_id, response_time_ms, embedding)
                    VALUES (%s, %s, %s, %s, %s);
                """, (query_text, user_id, session_id, response_time, embedding))

            # Insert some search cache entries
            for i in range(5):
                query_hash = f"hash_{random.randint(100000, 999999)}"
                query_text = f"Sample cached query {i + 1}"
                results = [{"doc_id": random.randint(1, len(doc_ids)), "similarity": round(random.uniform(0.7, 0.95), 3)} for _ in range(3)]
                result_count = len(results)
                search_time = random.randint(10, 100)
                threshold = round(random.uniform(0.6, 0.8), 3)

                cur.execute("""
                    INSERT INTO search_cache (query_hash, query_text, results_json, result_count, search_time_ms, similarity_threshold)
                    VALUES (%s, %s, %s, %s, %s, %s);
                """, (query_hash, query_text, json.dumps(results), result_count, search_time, threshold))

            logger.info(f"Sample data inserted successfully:")
            logger.info(f"   {len(sample_documents)} documents")
            logger.info(f"   {chunk_count} document chunks")
            logger.info(f"   {len(sample_queries)} user queries")
            logger.info(f"   {len(embedding_models)} embedding models")
            logger.info(f"   {len(knowledge_bases)} knowledge bases")

        conn.close()

    except psycopg2.Error as e:
        logger.error(f"Failed to insert sample data: {e}")
        raise


def verify_vector_setup():
    """Verify that the vector database was set up correctly."""
    conn_params = get_connection_params()

    try:
        conn = psycopg2.connect(**conn_params)

        with conn.cursor() as cur:
            logger.info("Verifying vector database setup...")

            # Check extension
            cur.execute("SELECT extname FROM pg_extension WHERE extname = 'vector';")
            if cur.fetchone():
                logger.info("pgvector extension is installed")
            else:
                logger.error("pgvector extension not found")
                return False

            # Check tables and record counts
            tables_to_check = [
                'documents', 'document_chunks', 'user_queries',
                'embedding_models', 'knowledge_base', 'search_cache'
            ]

            table_counts = {}
            for table in tables_to_check:
                cur.execute(f'SELECT COUNT(*) FROM {table}')
                count = cur.fetchone()[0]
                table_counts[table] = count
                logger.info(f"Table {table}: {count} records")

            # Check vector columns
            cur.execute("""
                SELECT table_name, column_name, data_type
                FROM information_schema.columns
                WHERE data_type = 'USER-DEFINED'
                AND udt_name = 'vector'
                ORDER BY table_name, column_name;
            """)

            vector_columns = cur.fetchall()
            logger.info(f"Found {len(vector_columns)} vector columns:")
            for table, column, dtype in vector_columns:
                logger.info(f"   {table}.{column} ({dtype})")

            # Check indexes
            cur.execute("""
                SELECT schemaname, tablename, indexname, indexdef
                FROM pg_indexes
                WHERE indexdef LIKE '%vector%' OR indexdef LIKE '%hnsw%' OR indexdef LIKE '%ivfflat%'
                ORDER BY tablename, indexname;
            """)

            vector_indexes = cur.fetchall()
            logger.info(f"Found {len(vector_indexes)} vector indexes:")
            for schema, table, index, definition in vector_indexes:
                logger.info(f"   {index} on {table}")

            # Test a simple vector similarity query
            mock_embedding = generate_mock_embedding(1536)
            cur.execute("""
                SELECT id, title, embedding <-> %s::vector as distance
                FROM documents
                ORDER BY embedding <-> %s::vector
                LIMIT 3;
            """, (mock_embedding, mock_embedding))

            results = cur.fetchall()
            logger.info(f"Vector similarity query returned {len(results)} results")

        conn.close()
        logger.info("Vector database verification completed successfully")
        return table_counts, vector_columns, vector_indexes

    except psycopg2.Error as e:
        logger.error(f"Verification failed: {e}")
        raise


def prepare_vector_environment():
    """Main function to prepare the vector database environment."""
    logger.info("Preparing vector database environment...")

    try:
        # Create pgvector extension
        create_vector_extension()

        # Create vector tables
        create_vector_tables()

        # Insert sample data first
        insert_sample_data()

        # Create indexes after data insertion for better performance
        create_vector_indexes()

        # Verify the setup
        table_counts, vector_columns, vector_indexes = verify_vector_setup()

        logger.info("Vector database environment prepared successfully!")
        logger.info(f"Total tables created: {len(table_counts)}")
        logger.info(f"Total vector columns: {len(vector_columns)}")
        logger.info(f"Total vector indexes: {len(vector_indexes)}")

        return {
            'table_counts': table_counts,
            'vector_columns': vector_columns,
            'vector_indexes': vector_indexes
        }

    except Exception as e:
        logger.error(f"Failed to prepare vector environment: {e}")
        raise


if __name__ == "__main__":
    # Allow running this module directly for testing
    logging.basicConfig(level=logging.INFO)
    prepare_vector_environment()