Database Agents with MCP and LangChain

# sql_server.py
import sqlite3
from mcp.server.fastmcp import FastMCP
from sqlalchemy import create_engine, text
from langchain_community.utilities import SQLDatabase

# Initialize the MCP server with a descriptive name
mcp = FastMCP("DatabaseQueryServer")

# --- Database Connection Setup ---
DB_FILE = "Chinook.db"
engine = create_engine(f"sqlite:///{DB_FILE}")
db = SQLDatabase(engine=engine)

@mcp.tool()
def list_tables() -> str:
    """
    Returns a comma-separated list of available table names in the database.
    Use this tool first to understand what tables you can query.
    """
    return ", ".join(db.get_usable_table_names())

@mcp.tool()
def get_schema_for_tables(tables: str) -> str:
    """
    Returns the schema (CREATE TABLE statement) for a comma-separated list of table names.
    Example input: "Album, Artist"
    """
    try:
        table_list = [table.strip() for table in tables.split(',')]
        return db.get_table_info(table_names=table_list)
    except Exception as e:
        return f"Error retrieving schema: {e}."

@mcp.tool()
def execute_safe_query(sql_query: str) -> str:
    """
    Executes a read-only SQL query and returns the result.
    IMPORTANT: This tool ONLY supports SELECT statements.
    """
    if not sql_query.strip().upper().startswith("SELECT"):
        return "Error: Only SELECT queries are allowed."
    
    try:
        with engine.connect() as connection:
            result = connection.execute(text(sql_query))
            return str(result.fetchall())
    except Exception as e:
        return f"Error executing query: {e}"

if __name__ == "__main__":
    mcp.run()

# agent_graph.py
from typing import List, TypedDict
from langchain_core.messages import BaseMessage
from langgraph.prebuilt import create_react_agent
from langgraph.graph import StateGraph, END
from langchain_community.agent_toolkits import SQLDatabaseToolkit

# 1. Define Agent State
class AgentState(TypedDict):
    question: str
    chat_history: List
    agent_scratchpad: List

# 2. Setup Toolkit and Agent Runnable
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()
agent_runnable = create_react_agent(llm, tools)

# 3. Define Graph Nodes
def run_agent_reasoning(state: AgentState):
    agent_outcome = agent_runnable.invoke(state)
    return {"agent_scratchpad": agent_outcome.return_values["messages"]}

def execute_tools(state: AgentState):
    tool_calls = state["agent_scratchpad"][-1].tool_calls
    # ... logic to execute tool calls ...
    return {"agent_scratchpad": [tool_output]}

def should_continue(state: AgentState):
    if "final_answer" in state["agent_scratchpad"][-1].content:
        return "end"
    return "continue"

# 4. Build the Graph
workflow = StateGraph(AgentState)
workflow.add_node("agent", run_agent_reasoning)
workflow.add_node("action", execute_tools)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {"continue": "action", "end": END},
)
workflow.add_edge("action", "agent")

# 5. Compile and run
app = workflow.compile()
inputs = {"question": "How many employees are there?", "chat_history": []}
for s in app.stream(inputs):
    print(s)