Adaptive RAG: Routing Documents, SQL, and Web

pip install -U langgraph langgraph-checkpoint-sqlite langchain langchain-ollama langchain-community langchain-core pydantic
ollama pull qwen3

import re
from langchain_community.utilities import SQLDatabase
from langchain_ollama import ChatOllama
from langchain_core.tools import tool

llm = ChatOllama(model="gpt-oss", base_url="http://localhost:11434", reasoning=True)

def get_db_connection():
    return SQLDatabase.from_uri('sqlite:///db/employees_db-full-1.0.6.db')

db = get_db_connection()
SCHEMA = db.get_table_info()

@tool
def get_database_schema(table_name: str = None):
    """Get database schema information for SQL query generation."""
    db = get_db_connection()
    if table_name:
        tables = db.get_usable_table_names()
        if table_name.lower() in [t.lower() for t in tables]:
            return db.get_table_info([table_name])
        return f"Error: Table '{table_name}' not found. Available tables: '{', '.join(tables)}'"
    return SCHEMA

@tool
def generate_sql_query(question: str, schema_info: str = None):
    """Generate a SQL SELECT query from a natural language question using database schema."""
    schema_to_use = schema_info if schema_info else SCHEMA
    prompt = f"""Based on this database schema:
                {schema_to_use}
                Generate a SQL query to answer this question: {question}
                Rules:
                - Use only SELECT statements
                - Include only existing columns and tables
                - Limit results to 10 rows unless specified otherwise
                - Use proper SQL syntax for SQLite
                Return only the SQL query, nothing else."""
    sql_query = llm.invoke(prompt).content.strip()
    print(f"[TOOL] Generated SQL Query: {sql_query[:30]}...")
    return sql_query

@tool
def validate_sql_query(query: str):
    """Validate SQL query for safety and syntax before execution."""
    clean_query = re.sub(r'```sql\s*', '', query.strip(), flags=re.IGNORECASE)
    clean_query = re.sub(r'```\s*', '', clean_query, flags=re.IGNORECASE).strip().rstrip(';')

    if not clean_query.lower().startswith('select'):
        return "Error: only 'select' statements are allowed."

    dangerous_keywords = ['INSERT', 'UPDATE', 'DELETE', 'ALTER', 'DROP', 'CREATE', 'TRUNCATE']
    for keyword in dangerous_keywords:
        if keyword in clean_query:
            return f"Error: {keyword} operations are not allowed."

    print("[TOOL] Your sql query is validated. Passed!")
    return clean_query

@tool
def execute_sql_query(sql_query: str):
    """Execute a validated SQL query and return results."""
    db = get_db_connection()
    query = validate_sql_query.invoke(sql_query)
    if query.startswith('Error:'):
        return f"Query '{sql_query}' validation failed with Error: {query}"
    result = db.run(query)
    return f"Query Results: {result}" if result else "Query Executed Successfully but No Result was Found!"

@tool
def fix_sql_error(original_query: str, error_message: str, question):
    """Fix a failed SQL query by analyzing the error and generating a corrected version."""
    fix_prompt = f"""The following SQL query failed:
                    Query: {original_query}
                    Error: {error_message}
                    Original Question: {question}
                    Database Schema:
                    {SCHEMA}
                    Return only the corrected SQL query, nothing else."""
    return llm.invoke(fix_prompt).content.strip()

ALL_SQL_TOOLS = [get_database_schema, generate_sql_query, execute_sql_query, fix_sql_error]

from typing_extensions import TypedDict
from typing import Annotated, List
import operator, os
from langgraph.graph import StateGraph, START, END
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage, SystemMessage
from pydantic import BaseModel, Field
from langchain.agents import create_agent

from scripts import my_tools, mysql_tools
from scripts.nodes import (
    retrieve_node, grade_documents_node, generate_node, transform_query_node,
    should_generate, check_answer_quality, get_latest_user_query
)

llm = ChatOllama(model="qwen3", base_url="http://localhost:11434")

class AgentState(TypedDict):
    messages: Annotated[list, operator.add]
    retrieved_docs: str
    rewritten_queries: List[str]

class RouterQuery(BaseModel):
    datasource: str = Field(description="Datasource to route query to: 'retrieve' for financial documents, 'sql_agent' for employee information, 'web_search' for general knowledge")

def route_question(state):
    query = get_latest_user_query(state['messages'])
    llm_router = llm.with_structured_output(RouterQuery)

    system_prompt = """You are an expert at routing user questions to the appropriate datasource.

                    1. retrieve: Financial documents (10-K, 10-Q, 8-K) about Amazon, Apple, Google
                       - revenue, operating income, cash flow, segments, quarterly/annual performance
                    2. sql_agent: Employee database
                       - employees, departments, salaries, managers, job titles, hiring dates
                    3. web_search: General knowledge outside our databases
                       - current events, general facts, unrelated topics

                    Return the datasource field only."""

    response = llm_router.invoke([SystemMessage(system_prompt), HumanMessage(query)])
    print(f"[DECISION] Routing to {response.datasource}")
    return response.datasource

def sql_agent_node(state):
    query = get_latest_user_query(state['messages'])
    system_prompt = f"""You are an expert SQL analyst working with an employees database.

                    Database Schema:
                    {mysql_tools.SCHEMA}

                    Workflow:
                    1. Use get_database_schema to understand tables (if needed)
                    2. Use generate_sql_query to create SQL
                    3. Use execute_sql_query to run the validated query
                    4. If there's an error, use fix_sql_error and retry (up to 3 times)

                    Provide clear, formatted results in MARKDOWN. Use tables for structured data."""

    agent = create_agent(model=llm, tools=mysql_tools.ALL_SQL_TOOLS, system_prompt=system_prompt)
    result = agent.invoke({'messages': [HumanMessage(query)]})
    print("[SQL] Agent completed!")
    return {'messages': result['messages']}

def web_search_node(state):
    query = get_latest_user_query(state['messages'])
    system_prompt = """You are a helpful web search assistant. Be concise and accurate.
                    Use the web_search tool, synthesize multiple sources, and answer in
                    MARKDOWN (150-250 words) with ## headings, **bold**, and bullet points."""

    agent = create_agent(model=llm, tools=[my_tools.web_search], system_prompt=system_prompt)
    result = agent.invoke({'messages': [HumanMessage(query)]})
    print("[WEB] Agent is completed!")
    return {'messages': result['messages']}

from langgraph.checkpoint.sqlite import SqliteSaver
import sqlite3

db_name = "db/checkpoints.db"
os.makedirs('db', exist_ok=True)

def create_adaptive_rag():
    builder = StateGraph(AgentState)

    builder.add_node('retrieve', retrieve_node)
    builder.add_node('grade_documents', grade_documents_node)
    builder.add_node('generate', generate_node)
    builder.add_node('transform_query', transform_query_node)
    builder.add_node('sql_agent', sql_agent_node)
    builder.add_node('web_search', web_search_node)

    # routed entry point
    builder.add_conditional_edges(START, route_question, ['retrieve', 'sql_agent', 'web_search'])

    # document path (Self-RAG)
    builder.add_edge('retrieve', 'grade_documents')
    builder.add_edge('transform_query', 'retrieve')
    builder.add_conditional_edges('grade_documents', should_generate, ['transform_query', 'generate'])
    builder.add_conditional_edges('generate', check_answer_quality, ['generate', END, 'transform_query'])

    # agent paths
    builder.add_edge('sql_agent', END)
    builder.add_edge('web_search', END)

    conn = sqlite3.connect(db_name, check_same_thread=False)
    checkpointer = SqliteSaver(conn)
    return builder.compile(checkpointer=checkpointer)

agent = create_adaptive_rag()