iris-vector-rag 0.2.1__py3-none-any.whl

common/database_schema_manager.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+"""
+Database Schema Manager for RAG Templates
+Provides centralized, config-driven table and column name resolution.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TableConfig:
+    """Configuration for a database table."""
+
+    name: str
+    alias: Optional[str] = None
+    description: str = ""
+    columns: Dict[str, str] = field(default_factory=dict)
+    indexes: List[Dict[str, Any]] = field(default_factory=list)
+
+
+@dataclass
+class SchemaConfig:
+    """Complete database schema configuration."""
+
+    name: str
+    description: str = ""
+    tables: Dict[str, TableConfig] = field(default_factory=dict)
+    column_mappings: Dict[str, Any] = field(default_factory=dict)
+    data_types: Dict[str, str] = field(default_factory=dict)
+
+
+class DatabaseSchemaManager:
+    """
+    Centralized manager for database schema configuration.
+    Eliminates hardcoded table/column names throughout the codebase.
+    """
+
+    def __init__(self, config_path: Optional[str] = None):
+        """Initialize with schema configuration."""
+        if config_path is None:
+            # Default to config/database_schema.yaml in project root
+            project_root = Path(__file__).parent.parent
+            config_path = project_root / "config" / "database_schema.yaml"
+
+        self.config_path = Path(config_path)
+        self.schema_config = self._load_schema_config()
+
+        # Cache for performance
+        self._table_cache = {}
+        self._column_cache = {}
+
+        logger.info(f"✅ Database schema manager initialized from {config_path}")
+
+    def _load_schema_config(self) -> SchemaConfig:
+        """Load schema configuration from YAML file."""
+        try:
+            with open(self.config_path, "r") as f:
+                config_data = yaml.safe_load(f)
+
+            # Parse schema info
+            schema_info = config_data.get("schema", {})
+            schema_name = schema_info.get("name", "RAG")
+            schema_desc = schema_info.get("description", "")
+
+            # Parse tables
+            tables = {}
+            tables_data = config_data.get("tables", {})
+            for table_key, table_data in tables_data.items():
+                tables[table_key] = TableConfig(
+                    name=table_data.get("name", table_key),
+                    alias=table_data.get("alias"),
+                    description=table_data.get("description", ""),
+                    columns=table_data.get("columns", {}),
+                    indexes=table_data.get("indexes", []),
+                )
+
+            return SchemaConfig(
+                name=schema_name,
+                description=schema_desc,
+                tables=tables,
+                column_mappings=config_data.get("column_mappings", {}),
+                data_types=config_data.get("data_types", {}),
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to load schema config from {self.config_path}: {e}")
+            raise
+
+    def get_schema_name(self) -> str:
+        """Get the schema name (e.g., 'RAG')."""
+        return self.schema_config.name
+
+    def get_table_name(self, table_key: str, fully_qualified: bool = True) -> str:
+        """
+        Get the actual table name for a logical table key.
+
+        Args:
+            table_key: Logical name (e.g., 'source_documents')
+            fully_qualified: Whether to include schema prefix (e.g., 'RAG.SourceDocuments')
+
+        Returns:
+            Actual table name
+        """
+        if table_key in self._table_cache:
+            table_name = self._table_cache[table_key]
+        else:
+            if table_key not in self.schema_config.tables:
+                raise ValueError(f"Unknown table key: {table_key}")
+
+            table_config = self.schema_config.tables[table_key]
+            table_name = table_config.name
+            self._table_cache[table_key] = table_name
+
+        if fully_qualified:
+            return f"{self.schema_config.name}.{table_name}"
+        return table_name
+
+    def get_column_name(self, table_key: str, column_key: str) -> str:
+        """
+        Get the actual column name for a logical column key.
+
+        Args:
+            table_key: Logical table name
+            column_key: Logical column name
+
+        Returns:
+            Actual column name
+        """
+        cache_key = f"{table_key}.{column_key}"
+        if cache_key in self._column_cache:
+            return self._column_cache[cache_key]
+
+        if table_key not in self.schema_config.tables:
+            raise ValueError(f"Unknown table key: {table_key}")
+
+        table_config = self.schema_config.tables[table_key]
+        if column_key not in table_config.columns:
+            raise ValueError(
+                f"Unknown column key '{column_key}' for table '{table_key}'"
+            )
+
+        column_name = table_config.columns[column_key]
+        self._column_cache[cache_key] = column_name
+        return column_name
+
+    def get_qualified_column(self, table_key: str, column_key: str) -> str:
+        """Get fully qualified column name (table.column)."""
+        table_name = self.get_table_name(table_key, fully_qualified=False)
+        column_name = self.get_column_name(table_key, column_key)
+        return f"{table_name}.{column_name}"
+
+    def standardize_column_name(self, column_name: str) -> str:
+        """
+        Standardize a column name using the mappings.
+
+        Args:
+            column_name: Original column name
+
+        Returns:
+            Standardized column name
+        """
+        mappings = self.schema_config.column_mappings
+
+        # Check each mapping category
+        for mapping_key, mapping_data in mappings.items():
+            if isinstance(mapping_data, dict) and "standard" in mapping_data:
+                variants = mapping_data.get("variants", [])
+                if column_name in variants:
+                    return mapping_data["standard"]
+
+        # Return original if no mapping found
+        return column_name
+
+    def get_data_type(self, type_key: str) -> str:
+        """Get data type definition."""
+        return self.schema_config.data_types.get(type_key, "VARCHAR(255)")
+
+    def build_create_table_sql(self, table_key: str) -> str:
+        """Generate CREATE TABLE SQL for a table."""
+        if table_key not in self.schema_config.tables:
+            raise ValueError(f"Unknown table key: {table_key}")
+
+        table_config = self.schema_config.tables[table_key]
+        table_name = self.get_table_name(table_key, fully_qualified=True)
+
+        # Build column definitions
+        column_defs = []
+        for col_key, col_name in table_config.columns.items():
+            data_type = self.get_data_type(col_key)
+            column_defs.append(f"    {col_name} {data_type}")
+
+        sql = f"CREATE TABLE IF NOT EXISTS {table_name} (\n"
+        sql += ",\n".join(column_defs)
+        sql += "\n);"
+
+        return sql
+
+    def get_all_tables(self) -> Dict[str, str]:
+        """Get mapping of all table keys to actual names."""
+        return {
+            key: self.get_table_name(key, fully_qualified=True)
+            for key in self.schema_config.tables.keys()
+        }
+
+    def validate_table_exists(self, table_key: str) -> bool:
+        """Check if a table configuration exists."""
+        return table_key in self.schema_config.tables
+
+    def get_table_info(self, table_key: str) -> Dict[str, Any]:
+        """Get complete information about a table."""
+        if table_key not in self.schema_config.tables:
+            raise ValueError(f"Unknown table key: {table_key}")
+
+        table_config = self.schema_config.tables[table_key]
+        return {
+            "key": table_key,
+            "name": table_config.name,
+            "alias": table_config.alias,
+            "description": table_config.description,
+            "fully_qualified_name": self.get_table_name(
+                table_key, fully_qualified=True
+            ),
+            "columns": table_config.columns,
+            "indexes": table_config.indexes,
+        }
+
+
+# Global schema manager instance
+_schema_manager = None
+
+
+def get_schema_manager() -> DatabaseSchemaManager:
+    """Get the global schema manager instance."""
+    global _schema_manager
+    if _schema_manager is None:
+        _schema_manager = DatabaseSchemaManager()
+    return _schema_manager
+
+
+# Convenience functions for common operations
+def get_table(table_key: str, fully_qualified: bool = True) -> str:
+    """Convenience function to get table name."""
+    return get_schema_manager().get_table_name(table_key, fully_qualified)
+
+
+def get_column(table_key: str, column_key: str) -> str:
+    """Convenience function to get column name."""
+    return get_schema_manager().get_column_name(table_key, column_key)
+
+
+def get_qualified_column(table_key: str, column_key: str) -> str:
+    """Convenience function to get qualified column name."""
+    return get_schema_manager().get_qualified_column(table_key, column_key)
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Test the schema manager
+    logging.basicConfig(level=logging.INFO)
+
+    manager = DatabaseSchemaManager()
+
+    print("🔍 Testing Database Schema Manager")
+    print("=" * 50)
+
+    # Test table names
+    print("Table Names:")
+    for table_key in ["source_documents", "document_chunks", "document_entities"]:
+        name = manager.get_table_name(table_key)
+        print(f"  {table_key} → {name}")
+
+    print("\nColumn Names:")
+    for table_key in ["source_documents", "document_chunks"]:
+        for col_key in ["id", "content"]:
+            try:
+                col_name = manager.get_column_name(table_key, col_key)
+                print(f"  {table_key}.{col_key} → {col_name}")
+            except ValueError as e:
+                print(f"  {table_key}.{col_key} → ERROR: {e}")
+
+    print("\nCreate Table SQL:")
+    print(manager.build_create_table_sql("source_documents"))

common/db_init_with_indexes.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+"""
+Complete database initialization script with all tables and indexes.
+This script creates a fresh RAG database with all optimizations included.
+"""
+
+import logging
+import sys
+from pathlib import Path
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from .iris_connection_manager import get_iris_connection
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def initialize_complete_rag_database(schema: str = "RAG"):
+    """
+    Initialize complete RAG database with all tables and indexes.
+
+    Args:
+        schema: Database schema name (default: "RAG")
+    """
+
+    conn = get_iris_connection()
+    cursor = conn.cursor()
+
+    try:
+        logger.info(f"🚀 Initializing complete RAG database in schema: {schema}")
+
+        # Read and execute the complete SQL schema
+        sql_file_path = Path(__file__).parent / "db_init_complete.sql"
+
+        if sql_file_path.exists():
+            with open(sql_file_path, "r") as f:
+                sql_content = f.read()
+
+            # Replace RAG with the specified schema
+            sql_content = sql_content.replace("RAG.", f"{schema}.")
+
+            # Split by semicolons and execute each statement
+            statements = [
+                stmt.strip() for stmt in sql_content.split(";") if stmt.strip()
+            ]
+
+            # Separate table creation from index creation
+            table_statements = []
+            index_statements = []
+
+            for statement in statements:
+                if statement and not statement.startswith("--"):
+                    if statement.upper().startswith("CREATE INDEX"):
+                        index_statements.append(statement)
+                    else:
+                        table_statements.append(statement)
+
+            # Execute table creation statements first
+            for i, statement in enumerate(table_statements):
+                try:
+                    cursor.execute(statement)
+                    logger.debug(
+                        f"✅ Executed table statement {i+1}/{len(table_statements)}"
+                    )
+                except Exception as e:
+                    if (
+                        "already exists" in str(e).lower()
+                        or "duplicate" in str(e).lower()
+                    ):
+                        logger.debug(f"⚠️ Table statement {i+1} - object already exists")
+                    else:
+                        logger.warning(f"⚠️ Table statement {i+1} failed: {e}")
+
+            # Use the index utility for index creation
+            if index_statements:
+                logger.info("Creating indexes with proper error handling...")
+                failed_indexes = []
+                for statement in index_statements:
+                    try:
+                        # Replace "CREATE INDEX IF NOT EXISTS" with "CREATE INDEX"
+                        statement = statement.replace("IF NOT EXISTS", "").replace(
+                            "if not exists", ""
+                        )
+                        cursor.execute(statement)
+                        logger.debug(f"✅ Created index: {statement[:50]}...")
+                    except Exception as e:
+                        error_str = str(e).lower()
+                        if any(
+                            indicator in error_str
+                            for indicator in [
+                                "already exists",
+                                "duplicate",
+                                "index exists",
+                                "name already used",
+                            ]
+                        ):
+                            logger.debug(
+                                f"⚠️ Index already exists (ignored): {statement[:50]}..."
+                            )
+                        else:
+                            logger.warning(
+                                f"⚠️ Index creation failed: {statement[:50]}... Error: {e}"
+                            )
+                            failed_indexes.append(statement)
+
+                if failed_indexes:
+                    logger.warning(f"⚠️ {len(failed_indexes)} indexes failed to create")
+                else:
+                    logger.info("✅ All indexes created successfully")
+
+            logger.info(f"✅ Schema initialization completed for {schema}")
+        else:
+            logger.error(f"❌ SQL file not found: {sql_file_path}")
+            return False
+
+        # Verify tables were created
+        logger.info("🔍 Verifying table creation...")
+
+        expected_tables = [
+            "SourceDocuments",
+            "DocumentChunks",
+        ]
+
+        for table in expected_tables:
+            try:
+                cursor.execute(f"SELECT COUNT(*) FROM {schema}.{table}")
+                count = cursor.fetchone()[0]
+                logger.info(f"✅ {schema}.{table}: {count:,} rows")
+            except Exception as e:
+                logger.warning(f"❌ {schema}.{table}: {e}")
+
+        logger.info("🎉 Complete RAG database initialization successful!")
+        return True
+
+    except Exception as e:
+        logger.error(f"❌ Error during database initialization: {e}")
+        return False
+    finally:
+        cursor.close()
+
+
+def create_schema_if_not_exists(schema: str = "RAG"):
+    """Create schema if it doesn't exist"""
+    conn = get_iris_connection()
+    cursor = conn.cursor()
+
+    try:
+        cursor.execute(f"CREATE SCHEMA IF NOT EXISTS {schema}")
+        logger.info(f"✅ Schema {schema} ready")
+    except Exception as e:
+        logger.warning(f"Schema creation: {e}")
+    finally:
+        cursor.close()
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Initialize complete RAG database")
+    parser.add_argument(
+        "--schema", default="RAG", help="Database schema name (default: RAG)"
+    )
+    args = parser.parse_args()
+
+    # Create schema first
+    create_schema_if_not_exists(args.schema)
+
+    # Initialize database
+    success = initialize_complete_rag_database(args.schema)
+
+    if success:
+        print(
+            f"🎉 Database initialization completed successfully for schema: {args.schema}"
+        )
+        print("📋 All tables and indexes are ready for:")
+        print("   - BasicRAG, ReRanking, CRAG")
+    else:
+        print("❌ Database initialization failed")
+        sys.exit(1)

common/db_vector_utils.py

@@ -0,0 +1,181 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+def insert_vector(
+    cursor: Any,
+    table_name: str,
+    vector_column_name: str,
+    vector_data: List[float],
+    target_dimension: int,
+    key_columns: Dict[str, Any],
+    additional_data: Optional[Dict[str, Any]] = None,
+) -> bool:
+    """
+    Inserts a record with a vector embedding into a specified table.
+    The vector is truncated to the target_dimension and inserted using TO_VECTOR(?).
+
+    Args:
+        cursor: Database cursor.
+        table_name: Name of the table to insert into (e.g., "RAG.DocumentTokenEmbeddings").
+        vector_column_name: Name of the column that stores the vector.
+        vector_data: The raw embedding vector (list of floats).
+        target_dimension: The dimension to truncate/pad the vector to.
+        key_columns: Dictionary of primary key or identifying column names and their values.
+                     (e.g., {"doc_id": "id1", "token_index": 0})
+        additional_data: Optional dictionary of other column names and their values.
+                         (e.g., {"token_text": "example"})
+
+    Returns:
+        True if insertion was successful, False otherwise.
+    """
+    # Validate cursor handle
+    if cursor is None:
+        logger.error(
+            f"DB Vector Util: Cannot insert vector into table '{table_name}': cursor is NULL"
+        )
+        return False
+
+    if not isinstance(vector_data, list) or not all(
+        isinstance(x, (float, int)) for x in vector_data
+    ):
+        logger.error(
+            f"DB Vector Util: Invalid vector_data format for table '{table_name}'. "
+            f"Expected list of floats/ints. Got type: {type(vector_data)}. Skipping insertion."
+        )
+        return False
+
+    # Truncate or pad the vector to the target dimension
+    processed_vector = vector_data[:target_dimension]
+    if len(processed_vector) < target_dimension:
+        logger.warning(
+            f"DB Vector Util: Original vector length ({len(vector_data)}) for table '{table_name}', column '{vector_column_name}' "
+            f"is less than target dimension ({target_dimension}). Padding with zeros."
+        )
+        processed_vector.extend([0.0] * (target_dimension - len(processed_vector)))
+
+    # Format as bracketed comma-separated string for IRIS TO_VECTOR() function
+    embedding_str = "[" + ",".join(map(str, processed_vector)) + "]"
+
+    all_columns_dict = {}
+    all_columns_dict.update(key_columns)
+    if additional_data:
+        all_columns_dict.update(additional_data)
+
+    # Separate vector column from other data for SQL construction
+    other_column_names = [col for col in all_columns_dict.keys()]
+    other_column_values = [all_columns_dict[col] for col in other_column_names]
+
+    column_names_sql = ", ".join(other_column_names + [vector_column_name])
+
+    # IMPORTANT: TO_VECTOR() does NOT accept parameter markers (?, :param, etc.)
+    # Must embed the vector string directly in SQL
+    # Use FLOAT to match how test data was created (see tests/fixtures/embedding_generator.py:247)
+    placeholders_list = ["?" for _ in other_column_names] + [
+        f"TO_VECTOR('{embedding_str}', FLOAT, {target_dimension})"
+    ]
+    placeholders_sql = ", ".join(placeholders_list)
+
+    # Use MERGE for upsert functionality to handle duplicates
+    # Build MERGE statement for IRIS
+    key_conditions = " AND ".join(
+        [f"target.{col} = source.{col}" for col in key_columns.keys()]
+    )
+    update_assignments = ", ".join(
+        [
+            f"{col} = source.{col}"
+            for col in other_column_names
+            if col not in key_columns
+        ]
+    )
+
+    # Separate approach: try INSERT first, if it fails due to constraint, try UPDATE
+    sql_query = (
+        f"INSERT INTO {table_name} ({column_names_sql}) VALUES ({placeholders_sql})"
+    )
+    # Don't pass embedding_str as a parameter - it's embedded in SQL above
+    params = other_column_values
+
+    try:
+        logger.debug(f"DB Vector Util: Executing INSERT: {sql_query}")
+        logger.debug(f"DB Vector Util: Parameters: {params}")
+        logger.debug(
+            f"DB Vector Util: Embedding string length: {len(embedding_str)} chars"
+        )
+        logger.debug(f"DB Vector Util: Vector dimension: {target_dimension}")
+        cursor.execute(sql_query, params)
+        return True
+    except Exception as e:
+        # Check for connection handle issues
+        error_str = str(e).lower()
+        if "_handle is null" in error_str or "handle is null" in error_str:
+            logger.error(
+                f"DB Vector Util: Database connection handle is NULL during vector insertion: {e}"
+            )
+            return False
+
+        # Check if it's a unique constraint violation
+        if "UNIQUE" in str(e) or "constraint failed" in str(e):
+            logger.debug(
+                f"DB Vector Util: INSERT failed due to duplicate key, attempting UPDATE..."
+            )
+
+            # Build UPDATE statement
+            set_clauses = []
+            update_params = []
+
+            # Add non-key columns to SET clause
+            for col in other_column_names:
+                if col not in key_columns:
+                    set_clauses.append(f"{col} = ?")
+                    update_params.append(all_columns_dict[col])
+
+            # Add vector column to SET clause - TO_VECTOR doesn't accept parameters
+            # Use FLOAT to match how test data was created (see tests/fixtures/embedding_generator.py:247)
+            set_clauses.append(
+                f"{vector_column_name} = TO_VECTOR('{embedding_str}', FLOAT, {target_dimension})"
+            )
+            # Don't append embedding_str to params - it's embedded in SQL above
+
+            # Add key columns to WHERE clause
+            where_clauses = []
+            for col, val in key_columns.items():
+                where_clauses.append(f"{col} = ?")
+                update_params.append(val)
+
+            if set_clauses and where_clauses:
+                update_sql = f"UPDATE {table_name} SET {', '.join(set_clauses)} WHERE {' AND '.join(where_clauses)}"
+
+                try:
+                    logger.debug(f"DB Vector Util: Executing UPDATE: {update_sql}")
+                    cursor.execute(update_sql, update_params)
+                    return True
+                except Exception as update_error:
+                    # Check for connection handle issues in UPDATE
+                    update_error_str = str(update_error).lower()
+                    if (
+                        "_handle is null" in update_error_str
+                        or "handle is null" in update_error_str
+                    ):
+                        logger.error(
+                            f"DB Vector Util: Database connection handle is NULL during UPDATE: {update_error}"
+                        )
+                    else:
+                        logger.error(
+                            f"DB Vector Util: UPDATE also failed: {update_error}"
+                        )
+                    return False
+            else:
+                logger.error(f"DB Vector Util: Could not build UPDATE statement")
+                return False
+        else:
+            logger.error(
+                f"DB Vector Util: Error inserting vector into table '{table_name}', column '{vector_column_name}': {e}"
+            )
+            logger.error(f"DB Vector Util: Key columns: {key_columns}")
+            logger.error(
+                f"DB Vector Util: Failing embedding string (first 100 chars): {embedding_str[:100] if 'embedding_str' in locals() else 'NOT_SET'}"
+            )
+            return False