Building AI Agents that Query Databases Using Natural Language: A Modular Guide

The Talk-to-Database example in the Arindam200/awesome-ai-apps repository demonstrates how to convert natural language questions into executable MySQL queries using LangChain and the Nebius-hosted Qwen-4.5-Air model, wrapped in a Streamlit interface that handles connection persistence and error surfacing.

Building AI agents that query databases using natural language requires orchestrating three distinct layers: a user interface for input capture, a database abstraction for secure connection handling, and an LLM service for semantic translation. The Talk-to-Database implementation provides a production-ready blueprint using Streamlit, PyMySQL, and LangChain to bridge human language and structured query execution while keeping credentials isolated from AI logic.

Architecture Overview

The codebase separates concerns across three dedicated modules to ensure maintainability and security. This architecture isolates database credentials from LLM logic while providing a stateful UI layer that persists configuration across user interactions.

The entry point app.py manages Streamlit session state and orchestrates calls to specialized utilities. The database.py module handles MySQL connection lifecycle management, while ai_services.py contains the natural language processing logic. This separation allows independent testing of SQL generation without requiring live database connections.

Building the Streamlit Interface (app.py)

The front-end layer provides a configuration sidebar and a main interaction area for query generation. It uses st.session_state to persist sensitive credentials and connection parameters across Streamlit reruns, preventing users from re-entering API keys on every interaction.

import streamlit as st
from database import parse_connection_string, get_database_connection, execute_query
from ai_services import translate_to_sql

# Initialize session state for persistence

if "db_config" not in st.session_state:
    st.session_state.db_config = None

with st.sidebar:
    api_key = st.text_input("Nebius API Key", type="password")
    conn_string = st.text_input("MySQL Connection String")
    
    if conn_string:
        st.session_state.db_config = parse_connection_string(conn_string)

When users click "Generate SQL Query", the UI calls translate_to_sql from ai_services.py, passing the natural language question and the static database schema. The sidebar implementation collects the Nebius API key and MySQL URI, then utilizes parse_connection_string to normalize the connection format into keyword arguments required by PyMySQL.

Database Connection Management (database.py)

The database.py module handles MySQL URI parsing, connection pooling, and safe query execution through three core functions that abstract database complexity away from the AI layer.

parse_connection_string converts MySQL URIs (e.g., mysql://user:pass@host:3306/dbname) into keyword arguments compatible with pymysql.connect. This normalization ensures the application accepts standard connection strings while maintaining driver compatibility.

import pymysql
from urllib.parse import urlparse

def parse_connection_string(uri):
    """Parse MySQL URI into connection parameters."""
    parsed = urlparse(uri)
    return {
        "host": parsed.hostname,
        "port": parsed.port or 3306,
        "user": parsed.username,
        "password": parsed.password,
        "database": parsed.path.lstrip('/')
    }

def get_database_connection(config):
    """Create connection from parsed config."""
    return pymysql.connect(**config)

def execute_query(conn, query):
    """Execute SQL and return results as dictionaries."""
    try:
        with conn.cursor(pymysql.cursors.DictCursor) as cursor:
            cursor.execute(query)
            results = cursor.fetchall()
            return results
    except Exception as e:
        raise RuntimeError(f"Query execution failed: {str(e)}")
    finally:
        conn.close()

get_database_connection retrieves the parsed configuration from Streamlit's session state and establishes a live database connection. execute_query runs generated SQL statements via a dictionary cursor, ensuring results serialize properly for the UI, and guarantees connections close via the finally block. All database errors propagate to the UI through st.error calls, preventing silent failures.

Natural Language to SQL Translation (ai_services.py)

The intelligence layer resides in ai_services.py, specifically within the translate_to_sql function. This implementation uses LangChain's ChatPromptTemplate to structure the interaction with the Nebius-hosted Qwen-4.5-Air model hosted on the Nebius platform.

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

DB_SCHEMA = """
Tables:
- category(id, name)
- product(id, name, price, category_id)
- order(id, product_id, quantity, order_date)
"""

def translate_to_sql(question, api_key):
    """Convert natural language to SQL using Qwen-4.5-Air."""
    llm = ChatOpenAI(
        model_name="qwen-4.5-air",
        openai_api_key=api_key,
        openai_api_base="https://api.studio.nebius.com/v1"
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a SQL expert. Given the schema below, write only the SQL query without markdown formatting or explanations.\n\nSchema: {schema}"),
        ("human", "{question}")
    ])
    
    chain = prompt | llm
    response = chain.invoke({
        "schema": DB_SCHEMA,
        "question": question
    })
    
    # Strip any markdown code blocks or backticks

    sql = response.content.replace("```sql", "").replace("```", "").strip()
    return sql

The prompt template injects a static DB_SCHEMA constant—documenting tables like category, product, and order—alongside the user's natural language question. The system prompt strictly instructs the model to output only the raw SQL statement, stripping any markdown formatting that could corrupt query execution.

End-to-End Workflow Integration

The complete workflow chains these operations sequentially. First, the connection string parses in the sidebar. When the user submits a question, translate_to_sql generates the query, which app.py then passes to execute_query along with a live connection from get_database_connection.

Error handling occurs at each boundary. Connection parsing failures surface immediately in the sidebar, SQL generation errors display in the main content area, and database execution failures render through st.error notifications with full traceback details from database.py.

Summary

• The Talk-to-Database implementation uses a three-layer architecture separating UI, database, and AI concerns across app.py, database.py, and ai_services.py.
• Natural language to SQL translation relies on LangChain's ChatPromptTemplate and the Qwen-4.5-Air model, constrained by system prompts to output only executable SQL without markdown formatting.
• Database security is maintained through URI parsing in parse_connection_string and automatic connection closure in execute_query.
• The Streamlit interface uses st.session_state to persist sensitive credentials across interactions while surfacing errors through st.error components rather than silent failures.

Frequently Asked Questions

How does the application prevent SQL injection when building AI agents that query databases?

The implementation relies on PyMySQL's parameterized connection handling through get_database_connection and cursor-based execution in execute_query. While the LLM generates the SQL structure, the database layer treats the generated query as a prepared statement context. Production deployments should add additional input validation layers and consider read-only database users for the connection credentials stored in st.session_state.

Which model does the Talk-to-Database example use for natural language translation?

According to the source code in ai_services.py, the translate_to_sql function utilizes the Nebius-hosted Qwen-4.5-Air model through LangChain's chat interface. The model is specifically instructed via system prompts to return only raw SQL statements without markdown code blocks or explanatory text, ensuring the output can be passed directly to execute_query.

Can this architecture support databases other than MySQL?

While the current database.py implementation uses pymysql.connect for MySQL connections, the modular structure supports extension to PostgreSQL or SQLite. You would modify parse_connection_string to handle alternative URI schemes (like postgresql://) and swap the connection driver in get_database_connection, while keeping ai_services.py unchanged since it generates ANSI-compatible SQL based on the provided schema.

Where are database credentials stored during user sessions?

Credentials exist only in st.session_state within app.py, persisting across Streamlit reruns but never written to disk or logged to console. The parse_connection_string function immediately decomposes the URI into connection parameters stored in memory, and execute_query ensures the connection closes automatically after fetching results, minimizing the window of active credential exposure.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}