Dynamic SQL Agent Enhancement

[jmanhype]

🚀 Dynamic SQL Agent Enhancement

📋 Overview

This PR enhances Ember with a dynamic SQL Agent capable of handling any database schema without hardcoded patterns. The agent dynamically explores database structures, generates appropriate SQL queries, and provides definitive answers for any question about the data.

✨ Key Improvements

1️⃣ Dynamic Database Schema Exploration

• 🔍 Added functionality to dynamically explore any database schema
• 📊 Retrieves all tables, columns, and sample data automatically
• 🔗 Attempts to infer relationships between tables
• 🌐 Adapts to any database structure without hardcoding

2️⃣ Improved Query Generation

• 💻 Generates SQL queries based on the complete database schema
• 📝 Explicitly instructs the LLM to use actual table names, not placeholders
• ✅ Validates queries against the actual database schema
• 🔄 Automatically corrects queries with invalid column or table names

3️⃣ Enhanced Answer Generation

• 📈 Generates clear, definitive answers based on query results
• 🔍 Includes specific details from query results in the answers
• 🧠 Adapts to any type of question without relying on predefined patterns
• 📚 Distills SQL query results and adds them to the LLM context

4️⃣ Interactive Streamlit Interface

• 🖥️ Added a web-based UI for interacting with the SQL Agent
• 💬 Chat-based interface for asking questions in natural language
• 📊 Displays generated SQL queries and their results
• 🎛️ Supports multiple LLM providers (OpenAI, Anthropic, DeepMind)
• 📤 Allows exporting chat history for documentation

🛠️ Technical Details

• 🔒 Uses SQLAlchemy's text() function for safer SQL query execution
• 🗂️ Implements dynamic schema mapping to understand any database structure
• 🧰 Added query validation and correction mechanisms
• 📊 Enhanced the answer generation process to provide more definitive responses
• 🌐 Built a responsive Streamlit interface for easy interaction

🧪 Testing

The enhanced SQL Agent has been tested with various queries and database schemas, demonstrating its ability to:

• 🔄 Dynamically map and understand any database structure
• 💡 Generate appropriate SQL queries for any question
• 📋 Provide clear, definitive answers based on query results

📂 Project Structure

src/ember/examples/sql_agent/
├── README.md                # Comprehensive documentation
├── __init__.py              # Package initialization
├── knowledge/               # Knowledge data for examples
├── load_f1_data.py          # Formula 1 data loader
├── requirements.txt         # Dependencies
├── run_sql_agent.py         # Command-line example runner
├── app.py                   # Streamlit web interface
├── sql_agent.py             # Main SQL Agent implementation
└── utils.py                 # Utility functions

📝 Example Usage

from ember.examples.sql_agent.sql_agent import SQLAgent

# Initialize the agent with a database connection
agent = SQLAgent(database_url="sqlite:///your_database.db")

# Ask a question about your data
result = agent.query("Who won the most races in the 2019 season?")
print(result["answer"])

🌐 Streamlit Interface

# First, set up your environment with the required API key
export OPENAI_API_KEY='your-api-key'  # or ANTHROPIC_API_KEY, DEEPMIND_API_KEY

# Then run the Streamlit app
streamlit run src/ember/examples/sql_agent/app.py

📊 Formula 1 Example Dataset

This example includes a Formula 1 dataset with tables for:

• 🏎️ Drivers championship
• 🏆 Constructors championship
• 🏁 Race results
• 🥇 Race wins
• ⏱️ Fastest laps

🔗 Related Issues

This PR addresses the need for a more dynamic and flexible SQL Agent that can handle any database schema without hardcoding.