LangChain: SQL Analyst

# Install: pip install langchain langchain-openai langchain-community sqlalchemy

from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain.chains import create_sql_query_chain
from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
from langchain_core.prompts import PromptTemplate
import os

# --- Config ---
DATABASE_URL = os.environ.get("DATABASE_URL", "sqlite:///./example.db")

# --- Connect to database ---
db = SQLDatabase.from_uri(DATABASE_URL, sample_rows_in_table_info=3)
print(f"Connected. Tables: {db.get_usable_table_names()}")

# --- Build chain ---
llm = ChatOpenAI(model="gpt-4o", temperature=0)

# Generate SQL from natural language
write_query = create_sql_query_chain(llm, db)

# Execute the query
execute_query = QuerySQLDataBaseTool(db=db)

# Answer template
answer_prompt = PromptTemplate.from_template(
    """Given the following user question, SQL query, and SQL result,
answer the user question in plain English. Be concise.

Question: {question}
SQL Query: {query}
SQL Result: {result}
Answer:"""
)

def ask(question: str) -> str:
    """Ask a question about the database in natural language."""
    query = write_query.invoke({"question": question})
    # Safety: only allow SELECT queries
    if not query.strip().upper().startswith("SELECT"):
        return "I can only run SELECT queries for safety."
    result = execute_query.invoke(query)
    answer = llm.invoke(
        answer_prompt.format(question=question, query=query, result=result)
    )
    return f"SQL: {query}\nResult: {result}\nAnswer: {answer.content}"

if __name__ == "__main__":
    print("SQL Analyst ready. Ask questions about your database.")
    while True:
        q = input("\nQuestion: ")
        if q.lower() == "quit":
            break
        print(ask(q))