Chat with Databases using LangChain

# Install required packages
pip install langchain langchain-community langchain-openai langgraph

# Import necessary modules
import os
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain import hub
from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
from langgraph.prebuilt import create_react_agent

# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = "your-api-key-here"

# Connect to an SQLite Database
from sqlalchemy import create_engine
from sqlalchemy.pool import StaticPool

# For an existing SQLite database
db_path = "chinook.db"  # Path to your SQLite database
engine = create_engine(f"sqlite:///{db_path}")
db = SQLDatabase(engine)

# Alternatively, for in-memory testing with example data
def get_in_memory_db():
    import sqlite3
    import requests

    url = "https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql"
    response = requests.get(url)
    sql_script = response.text

    connection = sqlite3.connect(":memory:", check_same_thread=False)
    connection.executescript(sql_script)

    return create_engine(
        "sqlite://",
        creator=lambda: connection,
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )

engine = get_in_memory_db()
db = SQLDatabase(engine)

# Initialize Language Model
llm = ChatOpenAI(model="gpt-4")  # You can also use gpt-3.5-turbo for lower cost

# Initialize SQL Database Toolkit
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
toolkit = SQLDatabaseToolkit(db=db, llm=llm)

# Get the tools from the toolkit
tools = toolkit.get_tools()

# Get a prompt from LangChain Hub
prompt_template = hub.pull("langchain-ai/sql-agent-system-prompt")
system_message = prompt_template.format(dialect="SQLite", top_k=5)

# Create the agent
agent_executor = create_react_agent(
    llm,
    tools,
    state_modifier=system_message
)

# views.py for Django integration
from django.http import JsonResponse
from django.views.decorators.csrf import csrf_exempt
import json

@csrf_exempt
def chat_with_db(request):
    if request.method == 'POST':
        try:
            data = json.loads(request.body)
            user_message = data.get('message', '')

            # Process the user message with our agent
            response = agent_executor.invoke(
                {"messages": [("user", user_message)]}
            )

            # Extract the final AI message
            ai_message = response["messages"][-1].content

            return JsonResponse({
                'success': True,
                'message': ai_message
            })

        except Exception as e:
            return JsonResponse({
                'success': False,
                'error': str(e)
            })

    return JsonResponse({
        'success': False,
        'error': 'Only POST requests are supported'
    })