text2sql.py

from llama_index.core.program import LLMTextCompletionProgram
import pandas as pd
from typing import List, Tuple, Optional, Dict, Any, Set
from decimal import Decimal
from datetime import datetime, date
from sqlalchemy import text
import json
import os
import uuid
import psycopg2
from private_gpt.settings.settings import Settings
from private_gpt.components.llm.llm_component import LLMComponent
from private_gpt.core_logger import logger
from private_gpt.components.embedding.embedding_component import EmbeddingComponent
from sqlalchemy import (
    create_engine,
    MetaData
)
import re
import time
from llama_index.core.objects import (
    SQLTableNodeMapping,
    ObjectIndex,
    SQLTableSchema,
)
from llama_index.core import SQLDatabase
from llama_index.core.retrievers import SQLRetriever
from typing import List
from llama_index.core.query_pipeline import FnComponent
from llama_index.core import PromptTemplate
from llama_index.core.bridge.pydantic import Field
from llama_index.core.llms import ChatMessage, MessageRole
from llama_index.core.query_pipeline import CustomQueryComponent
# from llama_index.core.prompts.default_prompts import DEFAULT_TEXT_TO_SQL_PROMPT

from llama_index.core.query_pipeline import FnComponent
from llama_index.core.llms import ChatResponse
from llama_index.core.query_pipeline import (
    QueryPipeline as QP,
    InputComponent,
)
from injector import inject
from llama_index.vector_stores.qdrant import QdrantVectorStore
import qdrant_client.models
from .query_sanitize import QuerySanitizer, TableInfo
import qdrant_client
from .qdrant_cache import QdrantQueryCache
from llama_index.core.prompts.prompt_type import PromptType
from private_gpt.server.chat.redis_history import RedisHistoryHandler
from private_gpt.server.google_calendar.load_config import load_config
from private_gpt.server.google_calendar.user_functions import *

config = load_config()

POSTGRES_PASSWORD = os.getenv('POSTGRES_PASSWORD')
POSTGRES_HOST = os.getenv('POSTGRES_HOST', 'localhost')
POSTGRES_PORT = '5432'
POSTGRES_DB = 'postgres'
QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
REDIS_TTL = config["redis_timeout"]

FORBIDDEN_NAMES = {
        # SQL Reserved Keywords
        'abort', 'absolute', 'access', 'action', 'add', 'admin', 'after', 'aggregate', 'alias', 
        'all', 'allocate', 'alter', 'analyse', 'analyze', 'and', 'any', 'array', 'as', 'asc', 
        'assertion', 'assignment', 'asymmetric', 'at', 'authorization', 'backward', 'before', 
        'begin', 'between', 'bigint', 'binary', 'bit', 'boolean', 'both', 'by', 'cache', 
        'call', 'cascade', 'cascaded', 'case', 'cast', 'catalog', 'char', 'character', 
        'characteristics', 'check', 'checkpoint', 'class', 'close', 'cluster', 'coalesce', 
        'collate', 'column', 'commit', 'committed', 'concurrently', 'condition', 'connect', 
        'connection', 'constraint', 'constraints', 'continue', 'convert', 'copy', 'cost', 
        'create', 'cross', 'csv', 'current', 'current_catalog', 'current_date', 'current_role', 
        'current_schema', 'current_time', 'current_timestamp', 'current_user', 'cursor', 
        'cycle', 'data', 'database', 'day', 'deallocate', 'dec', 'decimal', 'declare', 
        'default', 'defaults', 'deferrable', 'deferred', 'definer', 'delete', 'delimiter', 
        'delimiters', 'desc', 'dictionary', 'disconnect', 'distinct', 'do', 'domain', 'double', 
        'drop', 'each', 'else', 'end', 'escape', 'except', 'exclude', 'exclusive', 'execute', 
        'exists', 'explain', 'external', 'extract', 'false', 'fetch', 'filter', 'first', 
        'float', 'following', 'for', 'force', 'foreign', 'forward', 'freeze', 'from', 'full', 
        'function', 'functions', 'global', 'grant', 'granted', 'greatest', 'group', 'handler', 
        'having', 'header', 'hold', 'hour', 'identity', 'if', 'ilike', 'immediate', 'immutable', 
        'implicit', 'in', 'including', 'increment', 'index', 'indexes', 'inherit', 'inherits', 
        'initially', 'inner', 'input', 'insensitive', 'insert', 'instead', 'int', 'integer', 
        'intersect', 'interval', 'into', 'invoker', 'is', 'isnull', 'isolation', 'join', 'key', 
        'language', 'large', 'last', 'leading', 'least', 'left', 'level', 'like', 'limit', 
        'listen', 'load', 'local', 'localtime', 'localtimestamp', 'location', 'lock', 'login', 
        'logout', 'mapping', 'match', 'maxvalue', 'minute', 'minvalue', 'mode', 'month', 
        'move', 'names', 'national', 'natural', 'nchar', 'new', 'next', 'no', 'none', 'not', 
        'nothing', 'notify', 'null', 'nullif', 'numeric', 'object', 'of', 'off', 'offset', 
        'oids', 'old', 'on', 'only', 'operator', 'option', 'or', 'order', 'ordinality', 
        'outer', 'over', 'overlaps', 'overlay', 'owned', 'owner', 'parser', 'partial', 
        'partition', 'passing', 'password', 'path', 'pg_attribute', 'pg_class', 'pg_type', 
        'pg_tables', 'pg_views', 'pg_indexes', 'pg_user', 'pg_roles', 'pg_proc', 'pg_stat', 
        'pg_database', 'pg_namespace', 'placing', 'plans', 'policy', 'position', 'preceding', 
        'precision', 'prepare', 'primary', 'prior', 'privileges', 'procedure', 'procedures', 
        'program', 'quote', 'range', 'read', 'real', 'reassign', 'recheck', 'record', 
        'recursive', 'ref', 'references', 'referencing', 'refresh', 'reindex', 'relative', 
        'release', 'rename', 'repeatable', 'replace', 'replica', 'reset', 'restart', 
        'restrict', 'return', 'returns', 'revoke', 'right', 'role', 'rollback', 'row', 'rows', 
        'rule', 'savepoint', 'schema', 'scroll', 'search', 'second', 'security', 'select', 
        'sequence', 'sequences', 'serial', 'server', 'session', 'set', 'sets', 'share', 
        'show', 'similar', 'simple', 'smallint', 'snapshot', 'some', 'stable', 'standalone', 
        'start', 'statement', 'statistics', 'stdin', 'stdout', 'storage', 'strict', 'strip', 
        'subquery', 'substring', 'symmetric', 'sysid', 'system', 'table', 'tables', 'tablespace', 
        'temp', 'template', 'temporary', 'then', 'time', 'timestamp', 'to', 'trailing', 
        'transaction', 'transform', 'trigger', 'true', 'truncate', 'trusted', 'type', 'types', 
        'unbounded', 'uncommitted', 'unencrypted', 'union', 'unique', 'unknown', 'unlisten', 
        'until', 'update', 'user', 'using', 'vacuum', 'valid', 'validate', 'validator', 'value', 
        'values', 'varchar', 'variadic', 'varying', 'verbose', 'version', 'view', 'volatile', 
        'when', 'where', 'whitespace', 'window', 'with', 'within', 'without', 'work', 'wrapper', 
        'write', 'xml', 'xmlagg', 'xmlattributes', 'xmlbinary', 'xmlconcat', 'xmlelement', 
        'xmlexists', 'xmlforest', 'xmlnamespaces', 'xmlparse', 'xmlpi', 'xmlroot', 'xmlserialize', 
        'year', 'zone',
        
        # Common Data Types
        'serial', 'bigserial', 'int', 'int2', 'int4', 'int8', 'float4', 'float8', 'numeric', 
        'boolean', 'char', 'text', 'timestamp', 'time', 'date', 'interval', 'uuid'
    }

rewrite = (
    "Please write a query to a semantic search engine using the current conversation.\n"
    "\n"
    "\n"
    "{chat_history_str}"
    "\n"
    "\n"
    "Latest message: {query_str}\n"
    'Query:"""\n'
)

class ChatInputComponent(CustomQueryComponent):
    llm: Any = Field(description="Language model to use for query refinement")
    system_prompt: Optional[str] = Field(
        default="Your task is to analyze the current query in the context of the chat history "
               "and refine it to be more precise and complete. Maintain the original intent "
               "while incorporating relevant context from previous interactions. Output only "
               "the refined query without explanations.",
        description="System prompt for query refinement"
    )

    def _validate_component_inputs(
        self, input: Dict[str, Any]
    ) -> Dict[str, Any]:
        logger.info("Validating inputs")
        return input

    @property
    def _input_keys(self) -> Set[str]:
        return {"chat_history", "query_str"}

    @property
    def _output_keys(self) -> Set[str]:
        return {"query_str"}  # Changed from 'response' to 'query_str'

    def _prepare_context(
        self,
        chat_history: List[ChatMessage],
        query_str: str
    ) -> List[ChatMessage]:
        messages = []
        
        # Add system prompt
        if self.system_prompt:
            messages.append(ChatMessage(role="system", content=self.system_prompt))
        
        # Add chat history
        for msg in chat_history:
            messages.append(msg)
            
        # Add current query
        messages.append(ChatMessage(role="user", content=query_str))
        
        return messages

    def _run_component(self, **kwargs) -> Dict[str, Any]:
        chat_history = kwargs.get("chat_history", [])
        query_str = kwargs.get("query_str", "")
        
        try:
            # Prepare context with chat history and current query
            context = self._prepare_context(
                chat_history=chat_history,
                query_str=query_str
            )
            if chat_history:
            # Get refined query from LLM
                refined_query = self.llm.chat(messages=context)
            else:
                refined_query = query_str
            
            # Return the refined query
            return {"query_str": refined_query}
            
        except Exception as e:
            logger.error(f"Error in query refinement: {str(e)}")
            # If there's an error, return original query
            return {"query_str": query_str}
                
    # async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
    #     logger.info("arun")
    #     chat_history = kwargs["chat_history"]
    #     query_str = kwargs["query_str"]
    #     logger.info(query_str)
    #     logger.info(chat_history)
    #     prepared_context = self._prepare_context(chat_history, query_str)
    #     response = await self.llm.achat(prepared_context)
    #     logger.info(response)
    #     return {"response": response}

class InitializeClients:
    conn_str = None
    query_cache = None
    engine = None
    conn = None
    cur = None
    metadata_obj = None
    filtered_table_infos = None
    sql_database = None
    obj_retriever = None
    sql_retriever = None
    collection_name = "text2sql_query_cache"

    @classmethod
    def initialize(cls):
        # print("initialize")
        cls.conn_str = f'postgresql://postgres:{POSTGRES_PASSWORD}@{POSTGRES_HOST}/{POSTGRES_DB}'
        
        # Add error handling for Qdrant initialization
        if QDRANT_URL:
            try:
                cls.client = qdrant_client.QdrantClient(
                    url=QDRANT_URL,
                    api_key=QDRANT_API_KEY
                )
                print(f"Successfully connected to Qdrant at {QDRANT_URL}")
            except Exception as e:
                print(f"Warning: Failed to initialize Qdrant connection: {str(e)}")
                # print("The application will continue without query caching.")
                cls.client = None
        else:
            cls.client = None
        return cls

    @classmethod
    def _create_collection(cls):
        """Explicitly create Qdrant collection"""
        try:
            collections = cls.client.get_collections()
            # print(f"Current collections: {collections}")
            cls.client.get_collection(cls.collection_name)
        except Exception:
            # print(f"creating collection: {cls.collection_name}")
            cls.client.create_collection(
                collection_name=cls.collection_name,
                vectors_config=qdrant_client.models.VectorParams(
                    size=1024,
                    distance=qdrant_client.models.Distance.COSINE
                )
            )
            
        cls.vector_store = QdrantVectorStore(
                client=cls.client, 
                collection_name=cls.collection_name
            )
    
    @classmethod
    def initialize_settings(cls):
        logger.info("Initializing database schemas and user mapping...")
        
        try:
            cls.engine = create_engine(cls.conn_str)
            # print("create engine")
            cls.conn = psycopg2.connect(
                host=POSTGRES_HOST,
                password=POSTGRES_PASSWORD,
                user='postgres',
                database=POSTGRES_DB
            )
            # print("pyscop connect")
            cls.conn.autocommit = True

            cls.cur = cls.conn.cursor()
            
            # Create necessary schemas
            cls.cur.execute("CREATE SCHEMA IF NOT EXISTS csv_data")
            cls.cur.execute("CREATE SCHEMA IF NOT EXISTS query_history")
            cls.cur.execute("CREATE SCHEMA IF NOT EXISTS auth")
            cls.conn.commit()
            # Initialize metadata object and reflect schemas
            cls.metadata_obj = MetaData()
            cls.metadata_obj.reflect(bind=cls.engine, schema='csv_data')

            # cls.metadata_obj.reflect(bind=cls.engine, schema='auth')
            cls._create_collection()
            
            logger.info("Database initialization completed successfully")
            
        except Exception as e:
            logger.error(f"Error during database initialization: {str(e)}")
            raise RuntimeError("Failed to initialize database connections") from e
        
    @classmethod
    def load_existing_database(cls):
        """
        Retrieve existing database from PostgreSQL's csv_data schema and generate
        table summaries using LLM.
        """
        print("Retrieving existing database from PostgreSQL csv_data schema...")
        
        try:
            cls.table_names = []
            cls.table_infos = []
            
            # Get list of tables in csv_data schema
            tables_query = """
                SELECT tablename 
                FROM pg_catalog.pg_tables 
                WHERE schemaname = 'csv_data';
            """
            cls.cur.execute(tables_query)
            tables = cls.cur.fetchall()
            
            if not tables:
                logger.info("No existing tables found in csv_data schema")
                return
            
            for table_data in tables:
                table_name = table_data[0]
                
                fallback_table_info = TableInfo(
                    table_name=table_name,
                    table_summary=f"Table containing data from {table_name}"
                )
                cls.table_infos.append(fallback_table_info)
                
            logger.info(f"Successfully loaded {len(cls.table_infos)} tables from csv_data schema")
                    
        except Exception as e:
            logger.error(f"Error loading database schema: {str(e)}")
            cls.table_infos = []
            raise
        
        finally:
            if hasattr(cls, 'cur') and cls.cur is not None:
                cls.cur.close()

    @classmethod
    def check_user_mapping(cls, email):
        """
        Check if a user email exists in the mapping table and return their UUID.
        If the email doesn't exist, create a new mapping.
        """
        try:
            cls.cur = cls.conn.cursor()
            
            # Check if email exists
            cls.cur.execute("""
                SELECT id, login_count 
                FROM auth.email_mapping 
                WHERE email = %s
            """, (email,))
            
            result = cls.cur.fetchone()
            
            if result:
                # Update existing user's login information
                user_id, login_count = result
                cls.cur.execute("""
                    UPDATE auth.email_mapping 
                    SET last_login_date = NOW(),
                        has_logged_in = TRUE,
                        login_count = %s
                    WHERE id = %s
                """, (login_count + 1, user_id))
            else:
                # Create new user mapping
                user_id = uuid.uuid4()
                cls.cur.execute("""
                    INSERT INTO auth.email_mapping 
                    (id, email, first_login_date, last_login_date, has_logged_in, login_count)
                    VALUES (%s, %s, NOW(), NOW(), TRUE, 1)
                """, (str(user_id), email))
            
            cls.conn.commit()
            return str(user_id)
            
        except Exception as e:
            cls.conn.rollback()
            logger.error(f"Error in user mapping operation: {str(e)}")
            raise
        finally:
            if cls.cur is not None:
                cls.cur.close()    
    
    @classmethod
    def get_table_infos(cls):
        return cls.table_infos
    
    @classmethod
    def get_table_names(cls):
        return cls.table_names
    
    @classmethod
    def get_engine(cls):
        if not cls.engine:
            raise RuntimeError("Engine not initialized")
        return cls.engine
    
    @classmethod
    def get_metadata(cls):
        if not cls.metadata_obj:
            raise RuntimeError("Metadata not initialized")
        return cls.metadata_obj
    
    @classmethod
    def get_conn(cls):
        if not cls.conn:
            raise RuntimeError("Connection not initialized")
        return cls.conn
    
    @classmethod
    def get_query(cls):
        return cls.client
    
    @classmethod
    def get_vector_store(cls):
        return cls.vector_store
        
    @classmethod
    def close(cls):
        if cls.conn:
            cls.conn.close()
        if cls.engine:
            cls.engine.dispose()
        if cls.client:
            cls.client.close()

class Text2SQL:
    settings: Settings
    @inject
    def __init__(self, llmcomponent: LLMComponent, settings: Settings, embeddng: EmbeddingComponent, ):
        # self.csv_path = csv_path
        self.llm_service = llmcomponent
        self.settings = settings
        self.llm = self.llm_service.llm
        # self.structured_llm = self.llm.as_structured_llm()
        self.query_sanitizer = QuerySanitizer(self.llm)
        self.sql_query = None
        self.embedding = embeddng
        
        self.engine = InitializeClients.get_engine()
        self.metadata_obj = InitializeClients.get_metadata()
        self.conn = InitializeClients.get_conn()
        self.query_client = InitializeClients.get_query()
        self.vector_store = InitializeClients.get_vector_store()
        self.table_infos = InitializeClients.get_table_infos()
        self.table_names = InitializeClients.get_table_names()
        self.query_cache = QdrantQueryCache(self.embedding.embedding_model, self.vector_store)
        
        prompt_str = """\
        Give me a summary of the table with the following JSON format.

        - The original table name is provided and must remain unchanged.
        - Do NOT output a generic table name (e.g. table, my_table).

        Original Table Name: {table_name}
        Table:
        {table_str}

        Summary: """

        self.program = LLMTextCompletionProgram.from_defaults(
            output_cls=TableInfo,
            llm=self.llm,
            prompt_template_str=prompt_str,
        )

    def sanitize_column_name(self, col_name: str) -> str:
        """
        Sanitize column names for PostgreSQL compatibility.
        
        Args:
            col_name (str): Original column name
            
        Returns:
            str: Sanitized column name that is PostgreSQL compatible
        
        Features:
        - Converts to lowercase
        - Replaces spaces and special characters with underscores
        - Ensures name starts with a letter
        - Truncates to PostgreSQL's length limit
        - Handles duplicate underscores
        - Ensures unique names by adding numeric suffix if needed
        - Avoids forbidden names by appending a harmless suffix
        """
        
        if not isinstance(col_name, str):
            col_name = str(col_name)
            
        # Convert to lowercase
        clean_name = col_name.lower().strip()
        
        # Replace spaces and special characters with underscores
        clean_name = re.sub(r'[^a-z0-9_]', '_', clean_name)
        
        # Replace multiple underscores with single underscore
        clean_name = re.sub(r'_+', '_', clean_name)
        
        # Remove leading and trailing underscores
        clean_name = clean_name.strip('_')
        
        # Ensure name starts with a letter
        if not clean_name[0].isalpha():
            clean_name = 'tmp_' + clean_name
            
        # PostgreSQL has a 63-byte limit for identifiers
        if len(clean_name) > 63:
            # If truncating would create a name ending with underscore, remove it
            clean_name = clean_name[:63].rstrip('_')
            
        # Ensure we have at least one character
        if not clean_name:
            clean_name = 'column'
        
        # Check if the sanitized name is a forbidden PostgreSQL keyword
        if clean_name in FORBIDDEN_NAMES:
            clean_name = clean_name + '_col'  # Append a harmless suffix
        
        return clean_name

    def process_data(self, file_paths: list[str], uuid: str) -> list[pd.DataFrame]:
        """Load CSV, Excel, and PDF data into PostgreSQL's csv_data schema using optimized table creation.
        
        Args:
            file_paths: List of paths to CSV, Excel, or PDF files to process
            uuid: User UUID to prefix table names
            
        Returns:
            list[pd.DataFrame]: List of processed dataframes
        """
        dfs = []
        self.metadata_obj = MetaData()
        self.table_names = []
        self.ingestion_timestamps = []
        
        for file_path in file_paths:
            file_name = os.path.basename(file_path)
            logger.info(f"Processing file: {file_path}")
            
            try:
                if file_path.endswith(".csv"):
                    df = pd.read_csv(file_path)
                    df.columns = [self.sanitize_column_name(col) for col in df.columns]
                    
                    # Create table name with full UUID prefix
                    base_table_name = self.sanitize_column_name(os.path.splitext(file_name)[0])
                    table_name = f"{uuid}_{base_table_name}"
                    
                    # If base name is too long, truncate it while keeping full UUID
                    if len(table_name) > 63:
                        max_base_length = 63 - len(uuid) - 1  # -1 for underscore
                        base_table_name = base_table_name[:max_base_length]
                        table_name = f"{uuid}_{base_table_name}"
                    
                    self.table_names.append(table_name)
                    ingestion_timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                    self.ingestion_timestamps.append(ingestion_timestamp)
                    
                    self.create_table_from_dataframe_(
                        df=df,
                        table_name=table_name,
                    )
                    logger.info(f"Created table csv_data.{table_name}")
                    dfs.append(df)

                elif file_path.endswith((".xlsx", ".xls")):
                    excel_file = pd.ExcelFile(file_path)
                    sheet_names = excel_file.sheet_names
                    
                    for sheet_name in sheet_names:
                        logger.info(f"Processing sheet: {sheet_name}")
                        df = pd.read_excel(file_path, sheet_name=sheet_name)
                        if df.empty:
                            logger.info(f"Skipping empty sheet: {sheet_name}")
                            continue
                        
                        df.columns = [self.sanitize_column_name(col) for col in df.columns]
                        base_name = self.sanitize_column_name(os.path.splitext(file_name)[0])
                        sheet_name_clean = self.sanitize_column_name(sheet_name)
                        
                        # Create table name with full UUID prefix
                        table_name = f"{uuid}_{base_name}_{sheet_name_clean}"
                        
                        # If combined name is too long, truncate the base name and sheet name while keeping full UUID
                        if len(table_name) > 63:
                            remaining_space = 63 - len(uuid) - 2  # -2 for two underscores
                            # Split remaining space between base name and sheet name
                            max_part_length = remaining_space // 2
                            base_name = base_name[:max_part_length]
                            sheet_name_clean = sheet_name_clean[:max_part_length]
                            table_name = f"{uuid}_{base_name}_{sheet_name_clean}"
                        
                        self.table_names.append(table_name)
                        ingestion_timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                        self.ingestion_timestamps.append(ingestion_timestamp)
                        
                        self.create_table_from_dataframe_(
                            df=df,
                            table_name=table_name
                        )
                        logger.info(f"Created table csv_data.{table_name}")
                        dfs.append(df)

                # elif file_path.endswith(".pdf"):
                #     # Process PDF using Camelot
                #     logger.info(f"Attempting to extract tables from PDF: {file_name}")
                    
                #     tables_lattice = camelot.read_pdf(
                #         file_path,
                #         pages='all',
                #         flavor='lattice'
                #     )
                #     tables_stream = camelot.read_pdf(
                #         file_path,
                #         pages='all',
                #         flavor='stream'
                #     )
                #     tables = tables_lattice + tables_stream
                #     if tables:
                #         for i, table in enumerate(tables):
                #             df = table.df
                #             df.columns = [self.sanitize_column_name(col) for col in df.iloc[0]]  # Use first row as header
                #             df = df[1:]  # Remove the first row which is now headers
                            
                #             base_name = self.sanitize_column_name(os.path.splitext(file_name)[0])
                #             table_name = f"{base_name}_table_{i}"
                            
                #             if len(table_name) > 63:
                #                 table_name = table_name[:63]
                                
                #             self.table_names.append(table_name)
                #             ingestion_timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                #             self.ingestion_timestamps.append(ingestion_timestamp)
                #             self.create_table_from_dataframe_(
                #                 df=df,
                #                 table_name=table_name
                #             )
                #             logger.info(f"Created table csv_data.{table_name}")
                #             dfs.append(df)
                #     else:
                #         logger.info(f"No tables found in PDF: {file_name}")
                        
                else:
                    logger.info(f"Unsupported file format: {file_name}")
                    
            except Exception as e:
                logger.info(f"Error processing {file_name}: {str(e)}")
                continue
                        
        return dfs

    def create_table_from_dataframe_(self, df: pd.DataFrame, table_name: str, schema: str = 'csv_data'):
        """
        Upload DataFrame to PostgreSQL, replacing NaN/null with empty string.
        
        Args:
            df (pd.DataFrame): Input DataFrame to upload
            table_name (str): Destination table name
            schema (str, optional): Database schema. Defaults to 'csv_data'.
        """
        try:
            # Sanitize column names
            df = df.copy()
            df.columns = [self.sanitize_column_name(col) for col in df.columns]

            # Replace NaN and None with empty string
            # df = df.fillna('')
            
            # Write DataFrame to PostgreSQL
            df.to_sql(
                name=table_name, 
                con=self.engine, 
                schema=schema, 
                if_exists='replace',  # Replace existing table
                index=False,  # Don't write index as a column
                dtype=None  # Let SQLAlchemy infer types
            )

        except Exception as e:
            logger.error(f"Error uploading DataFrame to PostgreSQL: {e}")
            raise

    def extract_file_metadata(self, file_paths: List[str], ingestion_timestamps: List[str], uuid: str) -> pd.DataFrame:
        """
        Extract metadata for multiple files and update in database using uuid as user_id.
        If user exists, deletes all their existing records before inserting new ones.
        
        Args:
            file_paths (List[str]): List of paths to the files
            ingestion_timestamps (List[str]): List of timestamps of ingestion
            uuid (str): UUID for both user identification and file prefix
            
        Returns:
            pd.DataFrame: DataFrame containing metadata for all files
        """
        try:
            all_metadata_rows = []
            
            for file_path, ingestion_timestamp in zip(file_paths, ingestion_timestamps):
                try:
                    # Get file stats
                    stat_info = os.stat(file_path)
                    
                    if file_path.endswith((".xlsx", ".xls")):
                        # Handle Excel files with multiple sheets
                        excel_file = pd.ExcelFile(file_path)
                        sheet_names = excel_file.sheet_names
                        
                        for sheet_name in sheet_names:
                            df = pd.read_excel(file_path, sheet_name=sheet_name, nrows=5)
                            base_name = self.sanitize_column_name(os.path.basename(file_path[:-5]))
                            sheet_name_clean = self.sanitize_column_name(sheet_name)
                            table_name = f"{uuid}_{base_name}_{sheet_name_clean}"

                            if not df.empty:
                                all_metadata_rows.append({
                                    'user_id': uuid,
                                    'filename': table_name,
                                    'sheet_name': sheet_name,
                                    'file_path': file_path,
                                    'file_size_bytes': stat_info.st_size,
                                    'created_at': datetime.fromtimestamp(stat_info.st_ctime).strftime('%Y-%m-%d %H:%M:%S'),
                                    'modified_at': datetime.fromtimestamp(stat_info.st_mtime).strftime('%Y-%m-%d %H:%M:%S'),
                                    'ingested_at': ingestion_timestamp,
                                    'num_columns': len(df.columns),
                                    'columns': ', '.join(df.columns),
                                    'first_5_rows': df.head().to_json(),
                                    'file_type': 'excel'
                                })
                        
                    elif file_path.endswith('csv'):
                        # Handle CSV files
                        df = pd.read_csv(file_path, nrows=5)
                        all_metadata_rows.append({
                            'user_id': uuid,
                            'filename': f"{uuid}_{self.sanitize_column_name(os.path.basename(file_path[:-4]))}",
                            'sheet_name': '',
                            'file_path': file_path,
                            'file_size_bytes': stat_info.st_size,
                            'created_at': datetime.fromtimestamp(stat_info.st_ctime).strftime('%Y-%m-%d %H:%M:%S'),
                            'modified_at': datetime.fromtimestamp(stat_info.st_mtime).strftime('%Y-%m-%d %H:%M:%S'),
                            'ingested_at': ingestion_timestamp,
                            'num_columns': len(df.columns),
                            'columns': ', '.join(df.columns),
                            'first_5_rows': df.head().to_json(),
                            'file_type': 'csv'
                        })
                        
                except Exception as e:
                    print(f"Error processing file {file_path}: {e}")
                    continue
            
            if not all_metadata_rows:
                print("No valid files were processed")
                return None
                
            # Create combined metadata DataFrame
            metadata_df = pd.DataFrame(all_metadata_rows)

            # Handle database operations
            with self.engine.connect() as connection:
                # Check if user exists
                check_query = text("SELECT EXISTS(SELECT 1 FROM csv_data.file_metadata WHERE user_id = :user_id)")
                result = connection.execute(check_query, {"user_id": uuid})
                user_exists = result.scalar()

                if user_exists:
                    # Delete all existing records for this user
                    delete_query = text("DELETE FROM csv_data.file_metadata WHERE user_id = :user_id")
                    connection.execute(delete_query, {"user_id": uuid})

                # Insert all new records
                metadata_df.to_sql('file_metadata', 
                                connection, 
                                schema='csv_data',
                                if_exists='append', 
                                index=False)
                
                connection.commit()
            
            return metadata_df
                
        except Exception as e:
            print(f"Error in metadata extraction process: {e}")
            return None

    def dump_tableinfo(self, uploaded_path, uuid):
        """
        Process and store table information for all data files.
        Creates table summaries and maintains unique table names.
        """
        self.table_infos = []
        self.dfs = self.process_data(uploaded_path, uuid)
        
        for idx, df in enumerate(self.dfs):
            # Get sample of data for table summary
            df_str = df.head(10).to_csv()
            
            # Generate table info using LLM
            table_info = self.program(
                table_str=df_str,
                table_name=self.table_names[idx],
            )
            table_name = table_info.table_name
            logger.info(f"Processed table: {table_name}")

            self.table_infos.append(table_info)
            
        logger.info(f"Generated table info for {len(self.table_infos)} tables in csv_data schema")
    
    def create_table(self, file_paths: list[str], uuid):
        """Create tables in PostgreSQL csv_data schema and handle metadata tracking.
        
        Args:
            file_paths: List of paths to files being processed
        """
        logger.info("Creating tables in PostgreSQL csv_data schema")
        
        create_table_query = """
            CREATE TABLE IF NOT EXISTS csv_data.file_metadata (
                user_id VARCHAR(255) NOT NULL,
                filename VARCHAR(255) NOT NULL,
                sheet_name VARCHAR(255),
                file_path TEXT,
                file_size_bytes BIGINT,
                created_at TIMESTAMP,
                modified_at TIMESTAMP,
                ingested_at TIMESTAMP,
                num_columns INTEGER,
                columns TEXT,
                first_5_rows JSONB,
                file_type VARCHAR(50)
            );
            """
            
        with self.conn.cursor() as cursor:
            cursor.execute(create_table_query)
            self.conn.commit()
        
        if isinstance(file_paths, str):
            file_paths = [file_paths]
            
        uuid_prefix = self.sanitize_query(uuid[:12])

        self.dump_tableinfo(file_paths, uuid_prefix)
        
        self.extract_file_metadata(
            file_paths, 
            self.ingestion_timestamps,
            uuid_prefix
        )
            
        file_metadata_tableinfo = TableInfo(
            table_name="file_metadata",
            table_summary="Metadata about processed files"
        )
        self.table_infos.append(file_metadata_tableinfo)
        
        # try:
        #     self.table_names = []
        #     self.table_infos = []
            
        #     # Get list of tables in csv_data schema
        #     tables_query = """
        #         SELECT tablename 
        #         FROM pg_catalog.pg_tables 
        #         WHERE schemaname = 'csv_data';
        #     """
        #     with self.conn.cursor() as cursor:
        #         cursor.execute(tables_query)
        #         tables = cursor.fetchall()
                
        #         if not tables:
        #             print("No existing tables found in csv_data schema")
        #             return
                
        #         for table_data in tables:
        #             table_name = table_data[0]
                    
        #             # Get sample data for table summary
        #             sample_query = f"""
        #                 SELECT * 
        #                 FROM csv_data."{table_name}" 
        #                 LIMIT 10;
        #             """
        #             try:
        #                 # Get column names
        #                 cursor.execute(f"""
        #                     SELECT column_name 
        #                     FROM information_schema.columns 
        #                     WHERE table_schema = 'csv_data' 
        #                     AND table_name = '{table_name}'
        #                 """)
        #                 columns = [col[0] for col in cursor.fetchall()]
                        
        #                 # Get sample data
        #                 cursor.execute(sample_query)
        #                 sample_data = cursor.fetchall()
                        
        #                 # Convert to DataFrame and then to CSV string
        #                 import pandas as pd
        #                 df_sample = pd.DataFrame(sample_data, columns=columns)
        #                 df_str = df_sample.to_csv(index=False)
                        
        #                 while True:
        #                     # Generate table info using LLM
        #                     table_info = self.program(
        #                         table_str=df_str,
        #                         table_name=table_name
        #                     )
        #                     # generated_table_name = table_info.table_name
        #                     print(f"Processed table: {table_info.table_name}")
                            
        #                     # Ensure table name is unique
        #                     if table_info.table_name not in self.table_names:
        #                         self.table_names.append(table_name)
        #                         break
        #                     else:
        #                         print(f"Table name {table_name} already exists, trying again.")
        #                         pass
                        
        #                 self.table_infos.append(table_info)
                        
        #             except Exception as e:
        #                 print(f"Warning: Error processing table {table_name}: {str(e)}")
        #             # Add basic table info if LLM summary fails
        #             fallback_table_info = TableInfo(
        #                 table_name=table_name,
        #                 table_summary=f"Table containing data from {table_name}"
        #             )
        #             self.table_infos.append(fallback_table_info)
                    
        #             # print(self.table_infos)
        #         print(f"Successfully loaded {len(self.table_infos)} tables from csv_data schema:")
                    
        # except Exception as e:
        #     print(f"Error loading database schema: {str(e)}")
        #     self.table_infos = []
        #     raise
        
        # finally:
        #     if hasattr(self, 'cur') and self.cur is not None:
        #         self.cur.close()

    def load_existing_database(self):
        """
        Retrieve existing database from PostgreSQL's csv_data schema and generate
        table summaries using LLM.
        """
        print("Retrieving existing database from PostgreSQL csv_data schema...")
        
        try:
            self.table_names = []
            self.table_infos = []
            
            # Get list of tables in csv_data schema
            tables_query = """
                SELECT tablename 
                FROM pg_catalog.pg_tables 
                WHERE schemaname = 'csv_data';
            """
            with self.conn.cursor() as cur:
                cur.execute(tables_query)
                tables = cur.fetchall()
                
                if not tables:
                    logger.info("No existing tables found in csv_data schema")
                    return
                
                for table_data in tables:
                    table_name = table_data[0]
                    
                    fallback_table_info = TableInfo(
                        table_name=table_name,
                        table_summary=f"Table containing data from {table_name}"
                    )
                    self.table_infos.append(fallback_table_info)
                    
                self.metadata_obj = MetaData()
                self.metadata_obj.reflect(bind=self.engine, schema='csv_data')
                    
                logger.info(f"Successfully loaded {len(self.table_infos)} tables from csv_data schema")
                    
        except Exception as e:
            logger.error(f"Error loading database schema: {str(e)}")
            self.table_infos = []
            raise
        
        finally:
            if hasattr(self, 'cur') and self.cur is not None:
                self.cur.close()

    def retrieveSQL(self, user_id):
        """Retrieve SQL database focusing on latest ingestion date and including file_metadata."""
                
        try:
            # First ensure csv_data schema exists
            create_schema_query = """
            CREATE SCHEMA IF NOT EXISTS csv_data;
            """
            
            with self.conn.cursor() as cursor:
                cursor.execute(create_schema_query)
                self.conn.commit()

            latest_tables_query = """
                SELECT filename as table_name 
                FROM csv_data.file_metadata
                WHERE user_id = %s
                UNION
                SELECT 'file_metadata' as table_name
            """

            with self.conn.cursor() as cursor:
                cursor.execute(latest_tables_query, (user_id,))
                self.latest_tables = [row[0] for row in cursor.fetchall()]
    
            # print(self.latest_tables, self.table_infos)
            self.filtered_table_infos = [
                t for t in self.table_infos 
                if t.table_name in self.latest_tables
            ]
            
            # Initialize SQL database
            self.sql_database = SQLDatabase(self.engine, schema='csv_data')

            # Create table node mapping and schema objects
            table_node_mapping = SQLTableNodeMapping(self.sql_database)
            table_schema_objs = [
                SQLTableSchema(
                    table_name=t.table_name,
                    context_str=t.table_summary
                )
                for t in self.filtered_table_infos
            ]
            print("filtered infos:   ", self.filtered_table_infos)
            
            obj_index = ObjectIndex.from_objects(
                table_schema_objs,
                table_node_mapping,
                embed_model=self.embedding.embedding_model
            )
            self.obj_retriever = obj_index.as_retriever(similarity_top_k=3)
            self.sql_retriever = SQLRetriever(self.sql_database)
            
            logger.info(f"Successfully retrieved {len(self.filtered_table_infos)} tables from latest ingestion")
            if len(self.filtered_table_infos) == 0:
                raise ValueError("No tables found from latest ingestion")
            
        except Exception as e:
            logger.info(f"Error in retrieveSQL: {str(e)}")
            raise

    def get_table_context_str(self, table_schema_objs: List[SQLTableSchema]):
        """Get table context string."""
        context_strs = []
        for table_schema_obj in table_schema_objs:
            # print(table_schema_obj.model_json_schema)
            table_info = self.sql_database.get_single_table_info(
                table_schema_obj.table_name
            )
            if table_schema_obj.context_str:
                table_opt_context = " The table description is: "
                table_opt_context += table_schema_obj.context_str
                table_info += table_opt_context
                
            # print(table_info)

            context_strs.append(table_info)
        return "\n\n".join(context_strs)
    
    def sanitize_query(self, query: str) -> str:
        """
        Comprehensive query sanitization process
        
        Sanitization steps:
        1. Lowercase conversion
        2. Remove special characters
        3. Remove excessive whitespace
        4. Remove stop words
        5. Apply query guardrails
        6. Remove potentially harmful patterns
        """
        # Lowercase conversion
        # sanitized_query = query.lower()
        
        # Remove special characters and non-alphanumeric chars
        sanitized_query = re.sub(r'[^a-z0-9\s]', '', query)
        
        # Remove multiple whitespaces
        sanitized_query = re.sub(r'\s+', ' ', sanitized_query)
        
        # Predefined stop words to remove
        stop_words = set([
            'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 
            'to', 'for', 'of', 'with', 'by', 'from', 'up', 'about', 
            'into', 'over', 'after', 'beneath', 'under', 'above'
        ])
        
        # Remove stop words
        # word_tokens = sanitized_query.split()
        # filtered_query = [word for word in word_tokens if word not in stop_words]
        # sanitized_query = ' '.join(filtered_query)
        
        # Apply query guardrails
        # Remove potentially sensitive or dangerous terms
        dangerous_terms = [
            'drop', 'delete', 'truncate', 'alter', 
            'create', 'insert', 'update', 'merge',
            'execute', 'script', 'shell', 'system',
            'sudo', 'root', 'admin', 'password'
        ]
        
        # Remove dangerous terms
        for term in dangerous_terms:
            sanitized_query = sanitized_query.replace(term, '')
        
        # Length limitation
        max_query_length = 255
        sanitized_query = sanitized_query[:max_query_length]
        
        # Final trim
        sanitized_query = sanitized_query.strip()
        
        return sanitized_query
    
    def process_query(self, query: str) -> Tuple[bool, Optional[str], Optional[List[str]]]:
        """
        Process and sanitize a natural language query.
        Returns (is_safe, processed_query, clarifying_questions)
        """
        sanitized_query = self.sanitize_query(query)
        
        # Check for SQL injection patterns
        if self.query_sanitizer.check_sql_injection(sanitized_query):
            return False, None, ["Query contains potentially harmful patterns and was rejected."]

        # Generate clarifying questions if needed
        logger.info(self.filtered_table_infos, self.metadata_obj.tables)
        clarification = self.query_sanitizer.generate_clarifying_questions(
            sanitized_query,
            self.filtered_table_infos,
            self.metadata_obj
        )

        if clarification.needs_clarification:
            return False, sanitized_query, clarification.questions

        return True, sanitized_query, None

    def parse_response_to_sql(self, response: ChatResponse) -> str:
        """Parse response to SQL, normalize schema prefixes and add them where missing."""
        response_content = response.message.content
        
        # Extract SQL query from response
        sql_query_start = response_content.find("SQLQuery:")
        if sql_query_start != -1:
            response_content = response_content[sql_query_start:]
            if response_content.startswith("SQLQuery:"):
                response_content = response_content[len("SQLQuery:"):]
        
        # Remove SQL result section if present
        sql_result_start = response_content.find("SQLResult:")
        if sql_result_start != -1:
            response_content = response_content[:sql_result_start]
        
        sql_query = response_content.strip().strip("```").strip()
        
        # Remove standalone "sql" word (case-insensitive)
        sql_query = re.sub(r'\bsql\b', '', sql_query, flags=re.IGNORECASE)
        
        logger.info(f"Original SQL query: {sql_query}")
        logger.info(f"Latest tables: {self.latest_tables}")
        
        for table in self.latest_tables:
            # Escape table name for quoting if it contains special characters or starts with a number
            if table[0].isdigit() or "*" in table:
                table_pattern = f'"{table}"'  # Table pattern without schema
            else:
                table_pattern = table  # No need to quote if the table is simple
            
            schema_pattern = f'csv_data.{table_pattern}'  # The table pattern with schema
            
            # Step 1: Ensure correct schema qualification
            if table_pattern in sql_query and schema_pattern not in sql_query:
                sql_query = sql_query.replace(table_pattern, schema_pattern)
            elif table in sql_query:  # For unquoted table names
                sql_query = sql_query.replace(table, schema_pattern)
        
        # Clean up any extra whitespace created by removing "sql"
        sql_query = ' '.join(sql_query.split())
        
        self.sql_query = sql_query
        logger.info(f"Final SQL query: {sql_query}")
        
        return sql_query

    def identify_column_type(self, data_dict: List[Dict[str, Any]]) -> tuple[List[Dict[str, Any]], Dict[str, str]]:
        """
        Identifies the data type of columns and transforms the data dictionary.
        Converts Decimal values to floats during processing.
        Treats date-like column names as numerical if they contain date-formattable or float values.
        
        Args:
            data_dict (List[Dict[str, Any]]): List of dictionaries containing the data
        
        Returns:
            tuple: (transformed_data_dict, column_types)
                - transformed_data_dict (List[Dict[str, Any]]): Data with converted types
                - column_types (Dict[str, str]): Dictionary mapping column names to their types
        """
        if not data_dict:
            return [], {}
        
        columns = data_dict[0].keys()
        column_types = {}
        transformed_data = []
        
        for col_name in columns:
            values = [item[col_name] for item in data_dict]
            clean_values = [v for v in values if v is not None and not pd.isna(v)]
            
            if not clean_values:
                column_types[col_name] = 'unknown'
                continue
            
            # Check for boolean type first
            if all(isinstance(v, bool) or (isinstance(v, str) and v.lower() in ['true', 'false', 't', 'f', 'yes', 'no', 'y', 'n'])
                for v in clean_values):
                column_types[col_name] = 'boolean'
                continue
            
            # Check for datetime types first
            datetime_count = sum(1 for v in clean_values if isinstance(v, (datetime, pd.Timestamp, date)))
            if datetime_count / len(clean_values) > 0.8:
                column_types[col_name] = 'datetime'
                continue
                
            # Modified date column handling for string dates
            date_keywords = ['date', 'time', 'year', 'month', 'day', 'dt', 'timestamp']
            if any(keyword in col_name.lower() for keyword in date_keywords):
                date_count = 0
                for val in clean_values[:100]:  # Sample first 100 values
                    if isinstance(val, (datetime, pd.Timestamp, date)):
                        date_count += 1
                        continue
                        
                    try:
                        if isinstance(val, str):
                            pd.to_datetime(val)
                            date_count += 1
                    except (ValueError, TypeError):
                        continue
                
                sample_size = min(len(clean_values), 100)
                if date_count / sample_size > 0.8:
                    column_types[col_name] = 'datetime'
                    continue
            
            # Check for numeric type
            numeric_count = sum(1 for val in clean_values if isinstance(val, (int, float, Decimal)))
            str_numeric_count = 0
            for val in clean_values:
                if isinstance(val, str):
                    try:
                        float(val.replace(',', ''))
                        str_numeric_count += 1
                    except ValueError:
                        continue
            
            if (numeric_count + str_numeric_count) / len(clean_values) > 0.8:
                column_types[col_name] = 'numeric'
                continue
            
            # Check for categorical with numerical representation
            unique_ratio = len(set(clean_values)) / len(clean_values)
            if unique_ratio < 0.2:  # If less than 20% unique values, likely categorical
                column_types[col_name] = 'categorical'
                continue
            
            # Check if all remaining values are strings
            if all(isinstance(v, str) for v in clean_values):
                column_types[col_name] = 'categorical'
                continue
            
            column_types[col_name] = 'unknown'
        
        # Transform the data based on identified types
        for item in data_dict:
            transformed_item = {}
            for col_name, value in item.items():
                if value is None or pd.isna(value):
                    transformed_item[col_name] = None
                elif column_types[col_name] == 'numeric' and isinstance(value, Decimal):
                    transformed_item[col_name] = float(value)
                else:
                    transformed_item[col_name] = value
            transformed_data.append(transformed_item)
        # print(transformed_data, column_types)
        
        return transformed_data, column_types
    
    def suggest_and_format_charts(self, data_dict: List[Dict[str, Any]], chart_type: Optional[str] = None) -> Dict[str, Any]:
        """
        Formats data dictionary and determines/validates chart type with enhanced data type detection.
        Each data point contains x/y fields for column names and values stored under column names.
        
        Args:
            data_dict (List[Dict[str, Any]]): List of dictionaries containing the data
            chart_type (str, optional): Type of chart ('line', 'bar', 'pie', 'scatter', 'grouped_bar')
        
        Returns:
            Dict[str, Any]: Dictionary containing formatted data, suggested chart type, and field identifiers
        """
        
        def _format_value(value):
            """Helper method to format values consistently"""
            if isinstance(value, (Decimal, float)):
                return float(value)
            elif isinstance(value, (date, datetime)):
                return value.isoformat()
            return value
        if not data_dict or len(data_dict) < 2:  # Require at least 2 data points
            return {
                "data": None,
                "type": None,
            }
        
        # Convert Decimal to float in DataFrame creation
        df = pd.DataFrame(data_dict).map(lambda x: float(x) if isinstance(x, Decimal) else x)
        transformed_data, column_types = self.identify_column_type(data_dict)
        
        numerical_cols = [col for col, type_ in column_types.items() if type_ == 'numeric']
        categorical_cols = [col for col, type_ in column_types.items() if type_ == 'categorical']
        datetime_cols = [col for col, type_ in column_types.items() if type_ == 'datetime']
        # Select columns based on data distribution
        selected_numerical = numerical_cols[-1] if numerical_cols else None
        selected_categorical = categorical_cols[0] if categorical_cols else None
        
        selected_datetime = None
        if datetime_cols:
            date_ranges = {col: (df[col].max() - df[col].min()).total_seconds() 
                        for col in datetime_cols}
            selected_datetime = max(date_ranges.items(), key=lambda x: x[1])[0]
        
        # print(selected_categorical, selected_datetime, selected_numerical)
        suggested_type = chart_type
        
        if chart_type is None:
            if selected_categorical and selected_numerical and selected_datetime:
                suggested_type = 'bar'
            elif selected_categorical and selected_numerical and len(df) <= 15:
                suggested_type = 'pie'
            elif selected_datetime and selected_numerical:
                suggested_type = 'line'
        
        formatted_data = None
        
        if suggested_type == 'line':
            if selected_datetime and selected_numerical:
                formatted_data = [{
                    selected_datetime: _format_value(record[selected_datetime]),
                    selected_numerical: _format_value(record[selected_numerical])
                } for record in transformed_data]
        
        elif suggested_type == 'bar':
            if selected_numerical:
                if selected_categorical:
                    formatted_data = [{
                        selected_categorical: _format_value(record[selected_categorical]),
                        selected_numerical: _format_value(record[selected_numerical])
                    } for record in transformed_data]
        
        elif suggested_type == 'pie':
            if selected_numerical and selected_categorical:
                formatted_data = [{
                    selected_categorical: _format_value(record[selected_categorical]),
                    selected_numerical: _format_value(record[selected_numerical])
                } for record in transformed_data]
        
        return {
            "data": formatted_data,
            "type": suggested_type,
        }
    
    def extract_chart_data(self, original_query: str) -> Dict[str, Any]:
        """
        Extract chart-friendly data with query-aware chart type selection
        """
        logger.info("extracting chart data")
        try:
            # with self.engine.connect() as connection:
            #     result = connection.execute(text(sql_query))
            #     columns = result.keys()
            #     rows = result.fetchall()
                
            data_dict = [dict(zip(self.cols, row)) for row in self.query_results]

            # Filter out rows with None values
            data_dict = [row for row in data_dict if all(row.values())]
            
            chart_type = None
            
            # Check for explicit chart type mentions in query
            line_keywords = ['trend', 'over time', 'timeline', 'progression', 'line']
            pie_keywords = ['proportion', 'breakdown', 'composition', 'pie']
            bar_keywords = ['bar', 'vertical', 'horizontal']
            
            for keyword in line_keywords:
                if keyword.lower() in original_query.lower():
                    chart_type = 'line'
                    break
            
            for keyword in pie_keywords:
                if keyword.lower() in original_query.lower():
                    chart_type = 'pie'
                    break
                
            for keyword in bar_keywords:
                if keyword.lower() in original_query.lower():
                    chart_type = 'bar'
                    break
            
            transformed_data = self.suggest_and_format_charts(data_dict, chart_type)
            
            return transformed_data
        
        except Exception as e:
            logger.info(f"Cannot plot the requested chart: {e}")
            return None
    
    def store_query_history(self, query: str = None, response: str = None, sql_query: str = None, chart_data: dict = None, uuid: str = None):
        """
        Store query history in PostgreSQL database under query_history schema.
        
        Args:
            query (str): Original natural language query
            response (str): Generated response
            sql_query (str, optional): Generated SQL query
            chart_data (dict, optional): Generated chart data
            uuid (str, optional): Unique identifier to prefix the table name
        """
        try:
            # Validate UUID is provided
            if not uuid:
                raise ValueError("UUID must be provided to create a unique table name")
            
            safe_uuid = uuid
            
            # Create table name with UUID prefix
            table_name = f"query_history.queries_{safe_uuid}"
           
            with self.conn.cursor() as cursor:
                
                # Create query_history table if it doesn't exist
                create_table_sql = f"""
                CREATE TABLE IF NOT EXISTS {table_name} (
                    id SERIAL PRIMARY KEY,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    original_query TEXT,
                    sql_query TEXT,
                    response TEXT,
                    chart_data JSONB,
                    execution_time INTERVAL
                )
                """
                cursor.execute(create_table_sql)
                
                # Insert the query history
                insert_sql = f"""
                INSERT INTO {table_name} 
                    (original_query, sql_query, response, chart_data)
                VALUES 
                    (%s, %s, %s, %s)
                """
                
                # Convert chart_data to JSON string if it exists
                chart_data_json = json.dumps(chart_data) if chart_data else None
                
                cursor.execute(insert_sql, (
                    query,
                    sql_query,
                    response,
                    chart_data_json
                ))
                
                self.conn.commit()
                logger.info(f"Query history stored successfully in table {table_name}")
                
        except Exception as e:
            self.conn.rollback()
            logger.error(f"Error storing query history: {str(e)}")
            raise
    
    def link_qp(self, current_time_str):
        """Creates query pipeline for SQL generation"""
        logger.info("Initializing query pipeline...")
        
        table_parser_component = FnComponent(fn=self.get_table_context_str)
        sql_parser_component = FnComponent(fn=self.parse_response_to_sql)
        
        TEXT_TO_SQL_TMPL = (
            "Given an input question, create a syntactically correct {dialect} query following these rules:\n\n"
            "1. Column Selection:\n"
            " - Never query for all columns (*)\n"
            " - Only select columns relevant to the question\n"
            " - Always qualify column names with table names when needed\n"
            " - Only use columns visible in the schema description\n"
            " - Check type of data in column then generate the query.\n\n"
            "2. Date Handling:\n"
            " - Current date is: {current_time_str}\n"
            " - For text date columns:\n"
            " * Always cast to timestamp using ::timestamp\n"
            " * Example: date_col::timestamp >= '{current_time_str}'::timestamp\n"
            " - For date extraction:\n"
            " * Use: EXTRACT(YEAR FROM date_col::timestamp) or DATE_PART('year', date_col::timestamp)\n"
            " - For date ranges:\n"
            " * Use: date_col::timestamp >= '{current_time_str}'::timestamp - INTERVAL 'N days'\n"
            " - For historical data:\n"
            " * Always clip the data to current date if the table has date_col column: \n"
            " - date_col::timestamp >= CURRENT_DATA - INTERVAL 'N days'\n"
            " - Always handle timezone-aware timestamps (ending in Z) appropriately\n\n"
            "3. Query Structure:\n"
            " - Query should be executable in PostgreSQL\n"
            " - use ::numeric whenever using ROUND(x::numeric, y) function.\n"
            " - Include proper GROUP BY clauses for aggregations\n"
            " - Use appropriate ORDER BY for temporal data\n"
            " - Apply LIMIT only when explicitly needed\n"
            " - Don't include words that have no meaning in the SQL language, just provide the SQL query generated\n\n"
            "Response Format:\n"
            "Question: <input question>\n"
            "SQLQuery: <postgresql executable generated SQL query>\n"
            "SQLResult: <query results>\n"
            "Answer: <string formattable final answer>\n\n"
            "Available Tables and Schema:\n"
            "{schema}\n\n"
            "Question: {query_str}\n"
            "SQLQuery: "
            )

        TEXT_TO_SQL_PROMPT = PromptTemplate(
            TEXT_TO_SQL_TMPL,
            prompt_type=PromptType.TEXT_TO_SQL
        )
        
        text2sql_prompt = TEXT_TO_SQL_PROMPT.partial_format(
            dialect=self.engine.dialect.name,
            current_time_str=current_time_str
        )
        
        # Debug checkpoint 1
        logger.info(f"Using SQL dialect: {self.engine.dialect.name}")
        
        qp = QP(
            modules={
                "input": InputComponent(),
                "table_retriever": self.obj_retriever,
                "table_output_parser": table_parser_component,
                "text2sql_prompt": text2sql_prompt,
                "text2sql_llm": self.llm,
                "sql_output_parser": sql_parser_component,
            },
            verbose=True,
        )
        
        # Build pipeline
        qp.add_chain(["input", "table_retriever", "table_output_parser"])
        qp.add_link("input", "text2sql_prompt", dest_key="query_str")
        qp.add_link("table_output_parser", "text2sql_prompt", dest_key="schema")
        qp.add_chain(["text2sql_prompt", "text2sql_llm", "sql_output_parser"])
        
        # Debug checkpoint 2
        logger.info("Query pipeline constructed successfully")
        return qp
    
    def link_qp_history(self, current_time_str):
        """Creates query pipeline for SQL generation with chat history"""
        logger.info("Initializing query pipeline...")
        
        table_parser_component = FnComponent(fn=self.get_table_context_str)
        sql_parser_component = FnComponent(fn=self.parse_response_to_sql)

        TEXT_TO_SQL_TMPL = (
            "Given an input question, create a syntactically correct {dialect} query following these rules:\n\n"
            "1. Column Selection:\n"
            " - Never query for all columns (*)\n"
            " - Only select columns relevant to the question\n"
            " - Always qualify column names with table names when needed\n"
            " - Only use columns visible in the schema description\n"
            " - Check type of data in column then generate the query.\n\n"
            "2. Date Handling:\n"
            " - Current date is: {current_time_str}\n"
            " - For text date columns:\n"
            " * Always cast to timestamp using ::timestamp\n"
            " * Example: date_col::timestamp >= '{current_time_str}'::timestamp\n"
            " - For date extraction:\n"
            " * Use: EXTRACT(YEAR FROM date_col::timestamp) or DATE_PART('year', date_col::timestamp)\n"
            " - For date ranges:\n"
            " * Use: date_col::timestamp >= '{current_time_str}'::timestamp - INTERVAL 'N days'\n"
            " - For historical data:\n"
            " * Always clip the data to current date if the table has date_col column: \n"
            " - date_col::timestamp >= CURRENT_DATE - INTERVAL 'N days'\n"
            " - Always handle timezone-aware timestamps (ending in Z) appropriately\n\n"
            "3. Query Structure:\n"
            " - Query should be executable in PostgreSQL\n"
            " - use ::numeric whenever using ROUND(x::numeric, y) function.\n"
            " - Include proper GROUP BY clauses for aggregations\n"
            " - Use appropriate ORDER BY for temporal data\n"
            " - Apply LIMIT only when explicitly needed\n"
            " - Don't include words that have no meaning in the SQL language, just provide the SQL query generated\n\n"
            "Available Tables and Schema:\n"
            "{schema}\n\n"
            "Question: {query_str}\n\n"
            "Response Format:\n"
            "Question: {query_str}\n"
            "SQLQuery: <postgresql executable generated SQL query>\n"
            "SQLResult: <query results>\n"
            "Answer: <string formattable final answer>\n\n"
        )

        TEXT_TO_SQL_PROMPT = PromptTemplate(
            TEXT_TO_SQL_TMPL,
            prompt_type=PromptType.TEXT_TO_SQL
        )

        text2sql_prompt = TEXT_TO_SQL_PROMPT.partial_format(
            dialect=self.engine.dialect.name,
            current_time_str=current_time_str
        )

        logger.info(f"Using SQL dialect: {self.engine.dialect.name}")

        # Pass user chat history into the input component
        chat_input_component = ChatInputComponent(
            llm=self.llm, 
            system_prompt=(
                "Your task is to analyze the current query in the context of the chat history and refine it to be more precise and complete." 
                "If the query seems complete and not genuinely ambigous without the context of chat history, do not refine it."
                "Maintain the original intent while incorporating relevant context from previous interactions." 
                "Do not modify, expand, or interpret location names or other specific terms provided in the original query."
                "Output only the refined query without explanations.\n"
                "Response: <professional concise refined query>"
            ),
        )

        qp = QP(
            modules={
                "chat_input": chat_input_component,
                "table_retriever": self.obj_retriever,
                "table_output_parser": table_parser_component,
                "text2sql_prompt": text2sql_prompt,
                "text2sql_llm": self.llm,
                "sql_output_parser": sql_parser_component,
            },
            verbose=True,
        )

        # Simplified pipeline without rewrite_template
        qp.add_chain(["chat_input", "table_retriever", "table_output_parser"])
        qp.add_link("chat_input", "text2sql_prompt", dest_key="query_str")
        qp.add_link("table_output_parser", "text2sql_prompt", dest_key="schema")
        qp.add_chain(["text2sql_prompt", "text2sql_llm", "sql_output_parser"])

        logger.info("Query pipeline constructed successfully")
        return qp
    
    def verify_sql_query(self, query: str, sql_query: str, query_error: str) -> Dict[str, Any]:
        """
        Verify and correct SQL query if needed
        
        Args:
            query (str): Original natural language query
            sql_query (str): Generated SQL query
        
        Returns:
            Dict containing verification confidence and corrected query
        """
        verification_prompt = f"""
        You are a SQL expert. Evaluate this SQL query and if needed, provide a corrected postgresql executable version
        given the query execution error. Dont modify schema names appended to table names.
        Dont remove quotes.
        
        Original Query Execution error: "{query_error}"
        Natural Language Query: "{query}"
        Generated SQL Query: "{sql_query}"

        Return ONLY a single-line dictionary with no formatting or newlines:
        {{"confidence": <float>, "Reason": <list of str> explaining the changes made and why, SQLQuery": "<single line SQL query>"}}
        """
        
        try:
            response = self.llm.complete(verification_prompt)

            response_text = response.text if hasattr(response, 'text') else str(response)
          
            # response_text = response_text.replace('\n', '').replace('    ', ' ')
          
            verification_data = json.loads(response_text)
         
            # Ensure expected keys are present
            verification_data.setdefault('confidence', 0.5)
            verification_data.setdefault('SQLQuery', sql_query)  # Default to original query
            
            return verification_data
            
        except Exception as e:
            logger.info(f"Error in SQL query verification: {e}")
            return {
                'confidence': 0.5,
                'SQLQuery': sql_query  # Return original query if verification fails
            }

    def generate_response(self, query_str: str, sql_query: str) -> str:
        """
        Generates a natural language response based on SQL results using the LLM.
        If the query fails, attempts verification and runs corrected query.
        
        Args:
            query_str (str): Original user query
            sql_query (str): The SQL query to be executed
        
        Returns:
            str: Natural language response
            
        Raises:
            Exception: If query execution fails and cannot be corrected
        """
        try:
            with self.engine.connect() as connection:
                try:
                    # First attempt with original query
                    with connection.begin():
                        result = connection.execute(text(sql_query))
                        rows = result.fetchall()
                        columns = result.keys()
                except Exception as query_error:
                    logger.error(f"Initial query failed: {str(query_error)}")
                    
                    # Get verified/corrected query
                    verification_result = self.verify_sql_query(query_str, sql_query, str(query_error))
                    
                    if verification_result['confidence'] > 0.8 and verification_result['SQLQuery'] != sql_query:
                        logger.info("Attempting execution with corrected query")
                        # Start a fresh transaction for the corrected query
                        with connection.begin():
                            result = connection.execute(text(verification_result['SQLQuery']))
                            rows = result.fetchall()
                            columns = result.keys()
                            sql_query = verification_result['SQLQuery']
                    else:
                        raise query_error

                # Store results
                self.query_results = rows
                self.cols = columns
                self.sql_query = sql_query

                # Format results
                formatted_rows = []
                for row in rows:
                    if hasattr(row, '_asdict'):
                        formatted_rows.append(row._asdict())
                    elif hasattr(row, 'items'):
                        formatted_rows.append(dict(row.items()))
                    else:
                        formatted_rows.append(tuple(row))
                
                rows_str = "\n".join(str(row) for row in formatted_rows)
                # print(rows_str)
                if not rows_str:
                    return "Requested data not present in database!"

                response_synthesis_prompt = (
                    "Given an input question, synthesize a professional concise human response from the query results.\n"
                    "If query asks to plot, show or visualize something.\n"
                    "Please do not output any type of code and/or charts.\n"
                    "Do not start the response with something like 'Based on'.\n"
                    "Do not return anything is SQL Response is empty'.\n"
                    f"Query: {query_str}\n"
                    f"SQL: {sql_query}\n"
                    f"SQL Response: {rows_str}\n"
                    "Response: "
                )
                
                logger.info(f"Generated prompt: {response_synthesis_prompt}")
                return self.llm.complete(response_synthesis_prompt)
                    
        except Exception as e:
            logger.error(f"Error in generate_response: {str(e)}")
            raise
    
    async def execute_query(self, query: str, uuid, current_time_str, redis, thread_key) -> Tuple[bool, Optional[str], Optional[List[str]]]:
        """
        Enhanced query execution with SQL verification and proper transaction handling
        """
        chart_data = None
        try:
            # First, sanitize and check the query
            uuid_prefix = self.sanitize_query(uuid[:12])
            is_safe, processed_query, questions = self.process_query(query)
            
            if uuid_prefix not in ["5", "9"]:
                print("not pizza guys")
                self.history_handler = RedisHistoryHandler(redis)
                
                try:
                    ts = str(time.time())
                    await self.history_handler.store_data(
                        user_id=uuid_prefix,  # Make sure you have the user_id from your context
                        thread_key=thread_key,  # If you have a thread_key, otherwise pass None
                        sender='user',
                        text=encrypt_data(processed_query),
                        ttl=REDIS_TTL,
                        ts=ts
                    )
                except Exception as e:
                    logger.error(f"Error storing user message in Redis: {str(e)}")
               
                if not is_safe:               
                    try:
                        logger.info("Retrieving User chat history from Redis.")
                        decrypted_messages = await self.history_handler.get_decrypted_messages(uuid_prefix, thread_key)                        
                        user_chat_history = self.history_handler.convert_to_chat_messages(decrypted_messages)
                        logger.debug(f"Created User History Chat Messages")
                        
                        print("unclear query")
                        qp = self.link_qp_history(current_time_str)
                        response = qp.run(query_str=processed_query, chat_history=user_chat_history)
                        
                    except Exception as e:
                        logger.error(f"Error creating user history chat Messages: {str(e)}")
                        user_chat_history = []
                else:
                    qp = self.link_qp(current_time_str)
                    response = qp.run(query_str=processed_query)
            else:
                qp = self.link_qp(current_time_str)
                response = qp.run(query_str=processed_query)
               
            response = self.generate_response(processed_query, self.sql_query)
            chart_data = self.extract_chart_data(processed_query)
           
            # Store query in history
            self.store_query_history(
                query=processed_query,
                response=str(response),
                sql_query=self.sql_query,
                chart_data=chart_data,
                uuid=uuid_prefix
            )
            
            logger.info("Query Stored in DB.")
            self.query_cache.store_query_mapping(query, self.sql_query, str(response), chart_data)
            
            return True, str(response), None, None, None, chart_data
                    
        except Exception as e:
            logger.info(f"Error details: {str(e)}")
            return False, "Sorry! Error occured during processing of your query!", [f"Error processing query: {str(e)}"], None, None, None

if __name__=="__main__":
    
    qdrant_url=os.getenv("QDRANT_URL")
    qdrant_api_key=os.getenv("QDRANT_API_KEY")
    query = "revenue generated by each item in percentages in the last 7 days"
    import pytz
    current_time_str = datetime.now(pytz.UTC).strftime("%Y-%m-%d %H:%M:%S.000Z")
    agent = Text2SQL()
    agent.retrieveSQL()
    
    result = agent.execute_query(query, 4, current_time_str)
    print(result)