agents.rag.db_rag.sql_rag.example

Example usage of SQL RAG Agent.

from typing import Union This module demonstrates various usage patterns for the SQL RAG Agent, from basic queries to advanced configurations. It includes examples for different database types, error handling, and customization options.

Running the Examples:

Basic example:

$ python example.py

With specific database:

$ SQL_DB_TYPE=mysql SQL_DB_NAME=mydb python example.py

With custom query:

$ python example.py --query "Show me top products by revenue"
Examples Included:

  1. Basic usage with default configuration

  2. PostgreSQL with specific tables

  3. SQLite with local file

  4. MySQL with authentication

  5. Complex queries with joins

  6. Error handling and validation

  7. Custom LLM engines

  8. Batch processing

Note

Ensure you have proper database credentials configured either through environment variables or in the code before running.

Functions

basic_example()

Run a basic example with default configuration.

batch_processing_example()

Example of processing multiple queries in batch.

custom_llm_example()

Example with custom LLM configuration.

error_handling_example()

Demonstrate error handling and validation features.

interactive_mode()

Run the agent in interactive mode.

main()

Main function to run examples.

mysql_example()

Example with MySQL database and authentication.

postgresql_example()

Example with PostgreSQL configuration and specific tables.

sqlite_example()

Example with SQLite database file.

Module Contents

agents.rag.db_rag.sql_rag.example.basic_example()

Run a basic example with default configuration.

This example demonstrates the simplest usage of the SQL RAG Agent with minimal configuration, relying on environment variables.

Returns:

Query result including answer and SQL used.

Return type:

Dict[str, Any]

Example Output:
>>> result = basic_example()
✅ Connected to postgresql database
>>> print(result["answer"])
'The database contains the following tables: customers, orders, products...'
agents.rag.db_rag.sql_rag.example.batch_processing_example()

Example of processing multiple queries in batch.

This example shows how to efficiently process multiple queries using the same agent instance, with performance timing.

Returns:

Results for all queries.

Return type:

List[Dict[str, Any]]

Example

>>> results = batch_processing_example()
Processing 5 queries...
Query 1: ✓ (1.2s)
Query 2: ✓ (0.8s)
...
Total time: 5.5s, Average: 1.1s per query
agents.rag.db_rag.sql_rag.example.custom_llm_example()

Example with custom LLM configuration.

This example demonstrates how to customize the LLM engines used for different steps in the workflow.

Returns:

Query result using custom LLM configuration.

Return type:

Dict[str, Any]

Example

>>> result = custom_llm_example()
Using custom LLM configuration...
>>> print(result["answer"])
'Based on the analysis with custom temperature settings...'
agents.rag.db_rag.sql_rag.example.error_handling_example()

Demonstrate error handling and validation features.

This example shows how the agent handles various error conditions including invalid queries, non-existent tables, and SQL errors.

Example

>>> error_handling_example()
Testing error handling...
Query 1 - Invalid domain: This question is not about database...
Query 2 - SQL error corrected successfully
Query 3 - No results found: The database doesn't contain any orders...
Return type:

None

agents.rag.db_rag.sql_rag.example.interactive_mode()

Run the agent in interactive mode.

This function starts an interactive session where users can continuously ask questions about the database.

Example

>>> interactive_mode()
SQL RAG Agent - Interactive Mode
Type 'exit' to quit, 'help' for commands

SQL> What tables do we have? Answer: The database contains tables: customers, orders, products…

SQL> Show me total sales Answer: Total sales amount to $1,234,567…

SQL> exit Goodbye!

Return type:

None

agents.rag.db_rag.sql_rag.example.main()

Main function to run examples.

This function provides a command-line interface for running different examples or custom queries.

Command-line Arguments:

–example: Which example to run (basic, postgresql, sqlite, mysql, error, custom, batch, interactive) –query: Custom query to run –config: Path to JSON config file

Examples

Run basic example:

$ python example.py --example basic

Run custom query:

$ python example.py --query "Show me revenue by product category"

Run interactive mode:

$ python example.py --example interactive

Use custom config:

$ python example.py --config my_config.json --query "Top customers"
Return type:

int | float

agents.rag.db_rag.sql_rag.example.mysql_example()

Example with MySQL database and authentication.

This example shows MySQL configuration with full authentication and custom few-shot examples for better SQL generation.

Returns:

Query result from MySQL database.

Return type:

Dict[str, Any]

Example

>>> result = mysql_example()
Connecting to MySQL database...
✅ Connected to mysql database
>>> print(result["answer"])
'Product sales trend shows 15% growth...'
agents.rag.db_rag.sql_rag.example.postgresql_example()

Example with PostgreSQL configuration and specific tables.

This example shows how to configure the agent for a PostgreSQL database with specific table inclusion and custom domain settings.

Returns:

Query result from PostgreSQL database.

Return type:

Dict[str, Any]

Example

>>> result = postgresql_example()
Configuring PostgreSQL connection...
✅ Connected to postgresql database
>>> print(result["answer"])
'Your top 5 customers by total order value are...'
agents.rag.db_rag.sql_rag.example.sqlite_example()

Example with SQLite database file.

This example demonstrates using a local SQLite database file, which is useful for development and testing.

Returns:

Query result from SQLite database.

Return type:

Dict[str, Any]

Example

>>> result = sqlite_example()
Using SQLite database: ./data/sample.db
✅ Connected to sqlite database
>>> print(result["answer"])
'The total number of active users is 1,234...'