sql-query-ai-cli.py

#!/usr/bin/env python3
"""
SQAI (SQL Query AI) CLI Tool

A command-line interface for querying databases using either direct SQL queries
or natural language queries powered by LlamaIndex and Gemini.

Usage:
    python3 sql-query-ai-cli.py -t TABLE_NAME --query "Your query here"
    python3 sql-query-ai-cli.py -t TABLE_NAME -q "Your query here" -s  # For direct SQL
    python3 sql-query-ai-cli.py --interactive  # Interactive mode
    python3 sql-query-ai-cli.py -t TABLE_NAME --interactive  # Interactive mode with table pre-set

Environment Variables:
    DATABASE_URL: Database connection string (default: postgresql://postgres:postgres@localhost:5432/database)
    GEMINI_API_KEY: Your Gemini API key (required for NL queries)
"""

import os
import sys
import argparse
from typing import Optional, List
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.core import SQLDatabase
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.gemini import GeminiEmbedding
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from llama_index.core.llms import ChatMessage, MessageRole
from sqlalchemy import create_engine, text
import tiktoken


class SQLQueryCLI:
    def __init__(self, database_url: str):
        """Initialize the CLI with database connection."""
        self.database_url = database_url
        self.engine = None
        self.sql_database = None
        self.llm = None
        self.query_engine = None
        self._setup_failure_reason = None
        
        # Conversation context and token tracking
        self.conversation_history: List[ChatMessage] = []
        self.memory = None
        self.token_counter = None
        self.callback_manager = None
        self.max_context_tokens = 8000  # Conservative limit for Gemini
        self.warning_threshold = 6000   # Warn at 75% of limit
        
    def connect_to_database(self) -> bool:
        """Establish database connection."""
        try:
            self.engine = create_engine(self.database_url)
            # Test connection
            with self.engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            print(f"✓ Connected to database successfully")
            return True
        except Exception as e:
            print(f"✗ Error connecting to database: {e}")
            return False
    
    def setup_schema_path(self, table_name: str) -> tuple[str, str]:
        """Parse table name and set up schema path if needed."""
        if '.' in table_name:
            schema_name, table_name_only = table_name.split('.', 1)
            try:
                with self.engine.connect() as conn:
                    conn.execute(text(f"SET search_path TO {schema_name}, public"))
                    conn.commit()
                print(f"✓ Set search path to include schema: {schema_name}")
                return schema_name, table_name_only
            except Exception as e:
                print(f"⚠ Warning: Could not set search path for schema {schema_name}: {e}")
                return schema_name, table_name_only
        else:
            return None, table_name
    
    def execute_direct_sql(self, table_name: str, query: str = None) -> bool:
        """Execute direct SQL query (Method 1)."""
        try:
            if query is None:
                # Default query - show all records
                sql_query = f"SELECT * FROM {table_name}"
                print(f"=== All records from {table_name} table (Direct SQL) ===")
            else:
                # Custom SQL query
                sql_query = query
                print(f"=== Executing SQL: {sql_query} ===")
            
            with self.engine.connect() as connection:
                result = connection.execute(text(sql_query))
                rows = result.fetchall()
                columns = result.keys()
                
                if not rows:
                    print("No records found.")
                    return True
                
                # Print column headers
                print(f"Found {len(rows)} records:")
                header = " | ".join(str(col) for col in columns)
                print(header)
                print("-" * len(header))
                
                # Print all records
                for row in rows:
                    print(" | ".join(str(value) if value is not None else "NULL" for value in row))
                
            return True
            
        except Exception as e:
            print(f"✗ Error executing SQL query: {e}")
            return False
    
    def validate_table_exists(self, table_name: str) -> bool:
        """Validate that a table exists by trying a simple query."""
        try:
            # Parse schema and table name
            schema_name, table_name_only = self.setup_schema_path(table_name)
            
            # Try a simple query to check if table exists
            test_query = f"SELECT 1 FROM {table_name} LIMIT 1"
            with self.engine.connect() as connection:
                connection.execute(text(test_query))
                return True
        except Exception:
            return False
    
    def setup_nl_query_engine(self, table_name: str) -> bool:
        """Set up natural language query engine."""
        try:
            # Check for Gemini API key
            if not os.getenv("GEMINI_API_KEY"):
                print("✗ GEMINI_API_KEY environment variable not set.")
                print("You can set it with: export GEMINI_API_KEY='your-api-key-here'")
                return False
            
            # First validate that the table exists
            if not self.validate_table_exists(table_name):
                print(f"⚠ Table {table_name} not found in database")
                self._setup_failure_reason = "table_not_found"
                return False
            
            # Initialize LLM
            self.llm = Gemini(model="models/gemini-2.5-flash", temperature=0.7)
            self.embed_model = GeminiEmbedding(
                model_name="models/embedding-001", 
                api_key=os.getenv("GEMINI_API_KEY")
            )
            print("✓ Initialized Gemini LLM and embedding model")
            
            # Create SQL database wrapper
            self.sql_database = SQLDatabase(self.engine)
            
            # Check available tables and find the correct table name format
            usable_tables = list(self.sql_database.get_usable_table_names())
            print(f"Available tables: {usable_tables}")
            
            # Parse schema and table name
            schema_name, table_name_only = self.setup_schema_path(table_name)
            
            # Determine which table name format to use for the query engine
            # LlamaIndex may discover tables without schema prefix
            target_table = None
            if table_name_only in usable_tables:
                target_table = table_name_only
                print(f"✓ Using table name for query engine: {target_table}")
            elif table_name in usable_tables:
                target_table = table_name
                print(f"✓ Using full table name for query engine: {target_table}")
            else:
                # Fallback to the original table name if not found in usable_tables
                # This can happen with schema-qualified tables
                target_table = table_name_only if table_name_only else table_name
                print(f"✓ Using fallback table name for query engine: {target_table}")
            
            # Create NL Query Engine
            self.query_engine = NLSQLTableQueryEngine(
                sql_database=self.sql_database,
                tables=[target_table],
                llm=self.llm,
                embed_model=self.embed_model
            )
            
            print("✓ Natural language query engine initialized")
            self._setup_failure_reason = None
            return True
            
        except Exception as e:
            print(f"✗ Error setting up NL query engine: {e}")
            self._setup_failure_reason = "setup_error"
            return False
    
    def setup_conversation_context(self):
        """Initialize conversation context and token tracking."""
        try:
            # Initialize token counter with tiktoken for Gemini
            self.token_counter = TokenCountingHandler(
                tokenizer=tiktoken.encoding_for_model("gpt-4").encode  # Use GPT-4 tokenizer as approximation
            )
            self.callback_manager = CallbackManager([self.token_counter])
            
            # Initialize memory buffer with token limit
            self.memory = ChatMemoryBuffer.from_defaults(
                token_limit=self.max_context_tokens,
                tokenizer_fn=tiktoken.encoding_for_model("gpt-4").encode
            )
            
            print("✓ Conversation context and token tracking initialized")
            return True
        except Exception as e:
            print(f"⚠ Warning: Could not initialize conversation context: {e}")
            return False
    
    def get_token_count(self, text: str) -> int:
        """Calculate token count for a given text."""
        try:
            tokenizer = tiktoken.encoding_for_model("gpt-4")
            return len(tokenizer.encode(text))
        except Exception:
            # Fallback: rough estimation (1 token ≈ 4 characters)
            return len(text) // 4
    
    def get_conversation_token_count(self) -> int:
        """Get total token count for current conversation history."""
        if not self.conversation_history:
            return 0
        
        total_tokens = 0
        for message in self.conversation_history:
            # Convert message to string and count tokens
            message_text = str(message.content) if hasattr(message, 'content') else str(message)
            total_tokens += self.get_token_count(message_text)
        
        return total_tokens
    
    def add_to_conversation_history(self, role: MessageRole, content: str):
        """Add a message to conversation history."""
        message = ChatMessage(role=role, content=content)
        self.conversation_history.append(message)
        
        # Check if we need to truncate history to stay within limits
        while self.get_conversation_token_count() > self.max_context_tokens:
            if len(self.conversation_history) > 2:  # Keep at least the last exchange
                self.conversation_history.pop(0)
            else:
                break
    
    def clear_conversation_context(self):
        """Clear conversation history and reset context."""
        self.conversation_history.clear()
        if self.memory:
            self.memory.reset()
        print("✓ Conversation context cleared")
    
    def show_conversation_stats(self):
        """Display conversation statistics."""
        token_count = self.get_conversation_token_count()
        message_count = len(self.conversation_history)
        
        print(f"📊 Conversation Stats:")
        print(f"   Messages: {message_count}")
        print(f"   Tokens: {token_count:,} / {self.max_context_tokens:,}")
        print(f"   Usage: {(token_count/self.max_context_tokens)*100:.1f}%")
        
        if token_count > self.warning_threshold:
            print(f"⚠️  Warning: Approaching token limit! Consider using 'clear' command.")
    
    def build_context_prompt(self, current_query: str) -> str:
        """Build a prompt that includes conversation context."""
        if not self.conversation_history:
            return current_query
        
        context_parts = ["Previous conversation context:"]
        
        # Add recent conversation history
        for i, message in enumerate(self.conversation_history[-6:]):  # Last 6 messages
            role_label = "User" if message.role == MessageRole.USER else "Assistant"
            content = str(message.content) if hasattr(message, 'content') else str(message)
            context_parts.append(f"{role_label}: {content}")
        
        context_parts.append(f"\nCurrent query: {current_query}")
        
        return "\n".join(context_parts)
    
    def execute_nl_query(self, query: str) -> bool:
        """Execute natural language query with conversation context."""
        try:
            print(f"=== Natural Language Query: {query} ===")
            
            # Build query with conversation context
            contextual_query = self.build_context_prompt(query)
            
            # Track tokens before query
            query_tokens = self.get_token_count(contextual_query)
            
            # Execute query
            response = self.query_engine.query(contextual_query)
            response_text = str(response)
            
            # Track tokens after query
            response_tokens = self.get_token_count(response_text)
            total_query_tokens = query_tokens + response_tokens
            
            # Add to conversation history
            self.add_to_conversation_history(MessageRole.USER, query)
            self.add_to_conversation_history(MessageRole.ASSISTANT, response_text)
            
            # Display the generated SQL query
            if hasattr(response, 'metadata') and 'sql_query' in response.metadata:
                sql_query = response.metadata['sql_query']
                print(f"🔍 Generated SQL Query: {sql_query}")
            
            # Display response and token usage
            print(f"Response: {response}")
            print(f"\n🔢 Token Usage:")
            print(f"   Query: {query_tokens:,} tokens")
            print(f"   Response: {response_tokens:,} tokens")
            print(f"   Total this query: {total_query_tokens:,} tokens")
            
            # Show conversation stats and warnings
            conversation_tokens = self.get_conversation_token_count()
            print(f"   Conversation total: {conversation_tokens:,} tokens")
            
            if conversation_tokens > self.warning_threshold:
                print(f"⚠️  Warning: Conversation context is getting large ({conversation_tokens:,} tokens)!")
                print(f"   Consider using 'clear' command to reset context.")
            
            return True
        except Exception as e:
            print(f"✗ Error executing NL query: {type(e).__name__}: {e}")
            # Print more detailed error information for debugging
            import traceback
            print("Full traceback:")
            traceback.print_exc()
            return False
    
    def interactive_mode(self, initial_table_name=None):
        """Run in interactive mode."""
        print("=== SQL Query CLI - Interactive Mode ===")
        print("Type 'quit' or 'exit' to stop")
        print("Commands:")
        print("  sql <query>     - Execute direct SQL query")
        print("  nl <query>      - Execute natural language query")
        print("  table <name>    - Set/change table name")
        print("  show tables     - List available tables")
        print("  clear           - Clear conversation context")
        print("  stats           - Show conversation statistics")
        print("  help           - Show this help")
        print()
        
        # Use initial table name if provided, otherwise prompt for it
        table_name = None
        if initial_table_name:
            print(f"Using table name from command line: {initial_table_name}")
            # Validate that the table exists
            if not self.validate_table_exists(initial_table_name):
                print(f"⚠ Table {initial_table_name} not found in database")
                print("Please enter a valid table name:")
                table_name = None
            else:
                table_name = initial_table_name
        
        # Loop until we get a valid table name (if not already set)
        while not table_name:
            table_name = input("Enter table name: ").strip()
            if not table_name:
                print("Table name is required for interactive mode")
                continue
            
            # First validate that the table exists
            if not self.validate_table_exists(table_name):
                print("Please try again with a valid table name.")
                print()
                table_name = None
                continue
        
        # Set up schema path
        self.setup_schema_path(table_name)
        
        # Try to set up NL query engine (optional in interactive mode)
        nl_available = self.setup_nl_query_engine(table_name)
        
        # Initialize conversation context if NL queries are available
        if nl_available:
            self.setup_conversation_context()
        
        print(f"\nTable set to: {table_name}")
        if nl_available:
            print("Natural language queries are available with conversation context")
        else:
            print("Only direct SQL queries are available")
        print()
        
        while True:
            try:
                user_input = input("Query> ").strip()
                
                if not user_input:
                    continue
                
                if user_input.lower() in ['quit', 'exit']:
                    break
                
                if user_input.lower() == 'help':
                    print("Commands:")
                    print("  sql <query>     - Execute direct SQL query")
                    print("  nl <query>      - Execute natural language query")
                    print("  table <name>    - Set/change table name")
                    print("  show tables     - List available tables")
                    print("  clear           - Clear conversation context")
                    print("  stats           - Show conversation statistics")
                    print("  help           - Show this help")
                    continue
                
                if user_input.lower() == 'clear':
                    self.clear_conversation_context()
                    continue
                
                if user_input.lower() == 'stats':
                    if nl_available:
                        self.show_conversation_stats()
                    else:
                        print("Conversation context not available (NL queries not enabled)")
                    continue
                
                if user_input.lower() == 'show tables':
                    if self.sql_database:
                        tables = list(self.sql_database.get_usable_table_names())
                        print(f"Available tables: {tables}")
                    else:
                        try:
                            with self.engine.connect() as conn:
                                result = conn.execute(text("""
                                    SELECT table_name 
                                    FROM information_schema.tables 
                                    WHERE table_schema = 'public'
                                """))
                                tables = [row[0] for row in result.fetchall()]
                                print(f"Available tables: {tables}")
                        except Exception as e:
                            print(f"Error listing tables: {e}")
                    continue
                
                if user_input.startswith('table '):
                    new_table = user_input[6:].strip()
                    if new_table:
                        # Validate the new table exists
                        if not self.validate_table_exists(new_table):
                            print(f"⚠ Table {new_table} not found in database")
                            continue
                        
                        table_name = new_table
                        self.setup_schema_path(table_name)
                        
                        # Clear conversation context when changing tables
                        if nl_available:
                            self.clear_conversation_context()
                        
                        if nl_available:
                            nl_available = self.setup_nl_query_engine(table_name)
                            if nl_available:
                                self.setup_conversation_context()
                                print(f"✓ Table changed to: {table_name}")
                                print("Natural language queries are available with conversation context")
                            else:
                                print(f"⚠ Table changed to: {table_name}")
                                print("NL setup failed - only direct SQL queries are available")
                        else:
                            print(f"✓ Table changed to: {table_name}")
                    continue
                
                if user_input.startswith('sql '):
                    sql_query = user_input[4:].strip()
                    if sql_query:
                        self.execute_direct_sql(table_name, sql_query)
                    continue
                
                if user_input.startswith('nl '):
                    if not nl_available:
                        print("Natural language queries are not available. Use 'sql' for direct SQL queries.")
                        continue
                    
                    nl_query = user_input[3:].strip()
                    if nl_query:
                        self.execute_nl_query(nl_query)
                    continue
                
                # Default: treat as natural language query if available, otherwise as SQL
                if nl_available:
                    self.execute_nl_query(user_input)
                else:
                    print("Interpreting as SQL query (prefix with 'sql ' to be explicit):")
                    self.execute_direct_sql(table_name, user_input)
                
            except KeyboardInterrupt:
                print("\nGoodbye!")
                break
            except EOFError:
                print("\nGoodbye!")
                break


def main():
    parser = argparse.ArgumentParser(
        description="SQL Query CLI Tool - Query databases with SQL or natural language",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s -t users -q "SELECT * FROM users LIMIT 5" -s
  %(prog)s --table products --query "Show me all products with price > 100"
  %(prog)s --interactive
  %(prog)s -t users --interactive
  
Environment Variables:
  DATABASE_URL    Database connection string
  GEMINI_API_KEY  Your Gemini API key (required for NL queries)
        """
    )
    
    parser.add_argument(
        "-t", "--table",
        help="Table name to query (can include schema: schema.table)"
    )
    
    parser.add_argument(
        "-q", "--query",
        help="Query to execute (SQL if -s flag is used, otherwise natural language)"
    )
    
    parser.add_argument(
        "-s", "--sql",
        action="store_true",
        help="Use direct SQL query (Method 1) instead of natural language"
    )
    
    parser.add_argument(
        "--interactive",
        action="store_true",
        help="Run in interactive mode"
    )
    
    parser.add_argument(
        "--database-url",
        default=os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/database"),
        help="Database connection URL (default: from DATABASE_URL env var)"
    )
    
    args = parser.parse_args()
    
    # Validate arguments
    if not args.interactive and not args.table:
        print("Error: Table name is required for non-interactive mode (use -t/--table)")
        sys.exit(1)
    
    if not args.interactive and not args.query:
        print("Error: Query is required for non-interactive mode (use -q/--query)")
        sys.exit(1)
    
    # Initialize CLI
    cli = SQLQueryCLI(args.database_url)
    
    # Connect to database
    if not cli.connect_to_database():
        sys.exit(1)
    
    # Run interactive mode
    if args.interactive:
        cli.interactive_mode(args.table)
        return
    
    # Run single query mode
    if args.sql:
        # Direct SQL query
        success = cli.execute_direct_sql(args.table, args.query)
    else:
        # Natural language query
        if cli.setup_nl_query_engine(args.table):
            success = cli.execute_nl_query(args.query)
        else:
            # Check why setup failed
            if hasattr(cli, '_setup_failure_reason') and cli._setup_failure_reason == "table_not_found":
                print("✗ Cannot proceed: Table not found in database")
                print("Please verify the table name and ensure it exists in the database")
                sys.exit(1)
            else:
                print("Falling back to direct SQL query...")
                success = cli.execute_direct_sql(args.table, args.query)
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()