python -m venv venv
source venv/bin/activate
pip install sqlalchemy openai sqlparse pandas
export OPENAI_API_KEY="sk-your-key"

# schema/introspector.py
from sqlalchemy import create_engine, inspect, MetaData
from typing import Dict, List

class SchemaIntrospector:
    def __init__(self, connection_string: str):
        self.engine = create_engine(connection_string)
        self.inspector = inspect(self.engine)
        self.metadata = MetaData()
        self.metadata.reflect(self.engine)

    def get_schema_summary(self) -> str:
        tables = self.inspector.get_table_names()
        summary_parts = []
        for table in tables:
            columns = self.inspector.get_columns(table)
            col_defs = [f"  {c['name']} ({c['type']})" for c in columns]
            pk = self.inspector.get_pk_constraint(table)
            fks = self.inspector.get_foreign_keys(table)
            fk_defs = [f"  FK: {fk['constrained_columns']} -> {fk['referred_table']}.{fk['referred_columns']}" for fk in fks]
            summary_parts.append(f"Table: {table}\nColumns:\n" + "\n".join(col_defs))
            if pk.get("constrained_columns"):
                summary_parts.append(f"  Primary Key: {pk['constrained_columns']}")
            summary_parts.extend(fk_defs)
        return "\n".join(summary_parts)

    def get_table_ddl(self, table_name: str) -> str:
        table = self.metadata.tables.get(table_name)
        if table:
            from sqlalchemy import CreateTable
            return str(CreateTable(table).compile(self.engine))
        return f"Table {table_name} not found"

    def get_sample_data(self, table_name: str, n: int = 3) -> str:
        import pandas as pd
        df = pd.read_sql_table(table_name, self.engine)
        return df.head(n).to_string()

    def get_relationships(self) -> List[Dict]:
        relationships = []
        for table in self.inspector.get_table_names():
            fks = self.inspector.get_foreign_keys(table)
            for fk in fks:
                relationships.append({
                    "from_table": table,
                    "from_columns": fk["constrained_columns"],
                    "to_table": fk["referred_table"],
                    "to_columns": fk["referred_columns"],
                })
        return relationships

# generation/sql_generator.py
from openai import OpenAI
from typing import Dict

class SQLGenerator:
    def __init__(self, model: str = "gpt-4-turbo-preview"):
        self.client = OpenAI()
        self.model = model

    def generate(
        self, question: str, schema: str, sample_data: str = ""
    ) -> Dict:
        prompt = f"""You are a SQL expert. Generate a SQL query based on the natural language question.

Database Schema:
{schema}

{f"Sample data:\\n{sample_data}" if sample_data else ""}

Question: {question}

Rules:
1. Use only SELECT statements
2. Use table aliases for readability
3. Include appropriate JOINs
4. Use LIMIT for large result sets
5. Handle NULLs appropriately
6. Use proper aggregation functions

Return JSON:
{{
    "sql": "the SQL query",
    "explanation": "plain English explanation of the query",
    "assumptions": ["any assumptions made"]
}}"""
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a SQL expert generating queries from natural language."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.0,
        )
        import json
        try:
            return json.loads(response.choices[0].message.content)
        except:
            return {"sql": response.choices[0].message.content, "explanation": "", "assumptions": []}

# execution/validator.py
import sqlparse
from typing import Tuple

class QueryValidator:
    FORBIDDEN_KEYWORDS = {"INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "CREATE", "TRUNCATE", "GRANT", "REVOKE"}

    def validate(self, sql: str) -> Tuple[bool, str]:
        parsed = sqlparse.parse(sql)
        if not parsed:
            return False, "Invalid SQL syntax"
        statement = parsed[0]
        stmt_type = statement.get_type()
        if stmt_type and stmt_type.upper() in self.FORBIDDEN_KEYWORDS:
            return False, f"Forbidden operation: {stmt_type}"
        for token in statement.tokens:
            if token.ttype is sqlparse.tokens.Keyword:
                if token.value.upper() in self.FORBIDDEN_KEYWORDS:
                    return False, f"Forbidden keyword: {token.value}"
        return True, "Valid"

    def explain_query(self, sql: str) -> str:
        return f"Query will: {sql}"

# execution/executor.py
import pandas as pd
from sqlalchemy import create_engine, text
from typing import Dict

class QueryExecutor:
    def __init__(self, connection_string: str, timeout: int = 30):
        self.engine = create_engine(connection_string)
        self.timeout = timeout

    def execute(self, sql: str) -> Dict:
        try:
            with self.engine.connect() as conn:
                df = pd.read_sql(text(sql), conn)
                return {
                    "success": True,
                    "data": df.to_dict("records"),
                    "columns": list(df.columns),
                    "row_count": len(df),
                    "preview": df.head(20).to_string(),
                }
        except Exception as e:
            return {"success": False, "error": str(e), "data": []}

# agent.py
from schema.introspector import SchemaIntrospector
from generation.sql_generator import SQLGenerator
from execution.validator import QueryValidator
from execution.executor import QueryExecutor
from typing import Dict

class TextToSQLAgent:
    def __init__(self, connection_string: str, model: str = "gpt-4-turbo-preview"):
        self.introspector = SchemaIntrospector(connection_string)
        self.generator = SQLGenerator(model)
        self.validator = QueryValidator()
        self.executor = QueryExecutor(connection_string)

    def query(self, question: str) -> Dict:
        schema = self.introspector.get_schema_summary()
        result = self.generator.generate(question, schema)
        sql = result.get("sql", "")
        is_valid, validation_msg = self.validator.validate(sql)
        if not is_valid:
            return {"success": False, "error": f"Invalid query: {validation_msg}", "sql": sql}
        exec_result = self.executor.execute(sql)
        return {
            "success": exec_result["success"],
            "question": question,
            "sql": sql,
            "explanation": result.get("explanation", ""),
            "data": exec_result.get("data", []),
            "row_count": exec_result.get("row_count", 0),
            "error": exec_result.get("error"),
        }

    def explain_database(self) -> str:
        return self.introspector.get_schema_summary()

import pytest
from schema.introspector import SchemaIntrospector

def test_schema_introspection():
    introspector = SchemaIntrospector("sqlite:///test.db")
    schema = introspector.get_schema_summary()
    assert len(schema) > 0

def test_query_validation():
    validator = QueryValidator()
    valid, _ = validator.validate("SELECT * FROM users LIMIT 10")
    assert valid
    valid, _ = validator.validate("DROP TABLE users")
    assert not valid

# main.py
from agent import TextToSQLAgent

def main():
    agent = TextToSQLAgent("sqlite:///company.db")
    print("Text-to-SQL Agent Ready\n")
    while True:
        question = input("Ask a question (quit to exit): ").strip()
        if question.lower() == "quit":
            break
        result = agent.query(question)
        if result["success"]:
            print(f"\nSQL: {result['sql']}")
            print(f"Rows: {result['row_count']}")
            print(f"\n{result['data'][:5]}\n")
        else:
            print(f"\nError: {result['error']}\n")

if __name__ == "__main__":
    main()