agents.rag.db_rag.sql_rag.config¶

Configuration module for SQL RAG Agent.

This module provides configuration classes for setting up SQL database connections and agent behavior. It supports multiple database types and includes extensive customization options for the agent workflow.

Example

Basic PostgreSQL configuration:

>>> from haive.agents.rag.db_rag.sql_rag.config import SQLRAGConfig, SQLDatabaseConfig
>>>
>>> # Configure database connection
>>> db_config = SQLDatabaseConfig(
...     db_type="postgresql",
...     db_host="localhost",
...     db_name="sales_db",
...     db_user="admin",
...     db_password="secure_password",
...     include_tables=["orders", "customers", "products"]
... )
>>>
>>> # Configure agent
>>> agent_config = SQLRAGConfig(
...     domain_name="sales",
...     db_config=db_config,
...     hallucination_check=True,
...     max_iterations=3
... )

Note

Database credentials can be provided through environment variables for security. See SQLDatabaseConfig for supported variables.

Classes¶

`SQLDatabaseConfig`	Configuration for SQL database connections.
`SQLRAGConfig`	Configuration for the SQL RAG Agent.

Module Contents¶

class agents.rag.db_rag.sql_rag.config.SQLDatabaseConfig(/, **data)¶

Bases: pydantic.BaseModel

Configuration for SQL database connections.

This class handles connection details for various SQL database types including PostgreSQL, MySQL, SQLite, and MS SQL Server. Credentials can be provided directly or through environment variables.

Parameters:: data (Any)

db_type¶

Database type - postgresql, mysql, sqlite, mssql.

Type:: str

db_uri¶

Complete connection URI if provided directly.

Type:: Optional[str]

db_user¶

Database username (env: SQL_DB_USER).

Type:: str

db_password¶

Database password (env: SQL_DB_PASSWORD).

Type:: str

db_host¶

Database host (env: SQL_DB_HOST).

Type:: str

db_port¶

Database port (env: SQL_DB_PORT).

Type:: str

db_name¶

Database name (env: SQL_DB_NAME).

Type:: str

include_tables¶

Tables to include (None = all).

Type:: Optional[List[str]]

exclude_tables¶

Tables to exclude from queries.

Type:: List[str]

sample_rows_in_table_info¶

Sample rows to show in schema.

Type:: int

custom_query¶

Custom query for schema retrieval.

Type:: Optional[str]

Example

Using environment variables:

>>> # Set environment variables:
>>> # SQL_DB_TYPE=postgresql
>>> # SQL_DB_HOST=localhost
>>> # SQL_DB_NAME=mydb
>>> # SQL_DB_USER=user
>>> # SQL_DB_PASSWORD=pass
>>>
>>> config = SQLDatabaseConfig()
>>> db = config.get_sql_db()
✅ Connected to postgresql database

Direct configuration:

>>> config = SQLDatabaseConfig(
...     db_type="mysql",
...     db_host="mysql.example.com",
...     db_port="3306",
...     db_name="analytics",
...     db_user="analyst",
...     db_password="secure_pass",
...     include_tables=["sales", "customers"]
... )

Using connection URI:

>>> config = SQLDatabaseConfig(
...     db_uri="sqlite:///path/to/database.db"
... )

Note

For SQLite databases, db_name should be the file path. For other databases, ensure the appropriate driver is installed: - PostgreSQL: psycopg2 - MySQL: pymysql - MS SQL: pyodbc

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

get_connection_string()¶

Generate a connection string based on the database type.

Returns:: Properly formatted connection string for the database.
Return type:: str
Raises:: ValueError – If the database type is not supported.

Example

>>> config = SQLDatabaseConfig(
...     db_type="postgresql",
...     db_user="user",
...     db_password="pass",
...     db_host="localhost",
...     db_port="5432",
...     db_name="mydb"
... )
>>> print(config.get_connection_string())
'postgresql+psycopg2://user:pass@localhost:5432/mydb'

get_db_schema()¶

Retrieve the schema and basic table info from the database.

Returns:

Dictionary containing:

tables: List of table names
dialect: SQL dialect being used
table_info: Dictionary of table schemas

Return type:

Dict[str, Any]

Example

>>> config = SQLDatabaseConfig(db_uri="sqlite:///sales.db")
>>> schema = config.get_db_schema()
>>> print(f"Database dialect: {schema['dialect']}")
Database dialect: sqlite
>>> print(f"Tables: {', '.join(schema['tables'])}")
Tables: customers, orders, products

get_sql_db()¶

Create and return a SQLDatabase object for interacting with the database.

Returns:: Connected database object or None if connection fails.
Return type:: Optional[SQLDatabase]

Example

>>> config = SQLDatabaseConfig(db_uri="sqlite:///test.db")
>>> db = config.get_sql_db()
✅ Connected to sqlite database at sqlite:///test.db
>>> tables = db.get_usable_table_names()
>>> print(f"Found {len(tables)} tables")
Found 5 tables

Note

The method handles backward compatibility for different versions of langchain_community.utilities.SQLDatabase.

class agents.rag.db_rag.sql_rag.config.SQLRAGConfig¶

Bases: haive.core.engine.agent.agent.AgentConfig

Configuration for the SQL RAG Agent.

This class configures the behavior of the SQL RAG agent including database connection, LLM engines, validation settings, and workflow parameters.

engines¶

LLM engines for each workflow step.

Type:: Dict[str, AugLLMConfig]

llm_config¶

Default LLM configuration.

Type:: LLMConfig

domain_name¶

Domain specialization (e.g., “sales”, “inventory”).

Type:: str

domain_categories¶

Valid categories for domain routing.

Type:: List[str]

state_schema¶

State schema class for the agent.

Type:: Any

db_config¶

Database connection configuration.

Type:: SQLDatabaseConfig

input_schema¶

Input schema for agent invocation.

Type:: Any

output_schema¶

Output schema for agent results.

Type:: Any

hallucination_check¶

Enable hallucination detection.

Type:: bool

answer_grading¶

Enable answer quality grading.

Type:: bool

examples_path¶

Path to few-shot examples JSON.

Type:: Optional[str]

domain_examples¶

Few-shot examples.

Type:: Dict[str, List[Dict[str, str]]]

max_iterations¶

Maximum SQL correction attempts.

Type:: int

Example

Complete configuration example:

>>> config = SQLRAGConfig(
...     domain_name="e-commerce",
...     domain_categories=["sales", "inventory", "customers"],
...     db_config=SQLDatabaseConfig(
...         db_type="postgresql",
...         db_name="ecommerce_db",
...         include_tables=["orders", "products", "customers"]
...     ),
...     hallucination_check=True,
...     answer_grading=True,
...     max_iterations=3,
...     domain_examples={
...         "e-commerce": [
...             {
...                 "question": "Top selling products",
...                 "query": "SELECT p.name, SUM(oi.quantity) as sold FROM products p JOIN order_items oi ON p.id = oi.product_id GROUP BY p.id ORDER BY sold DESC LIMIT 10"
...             }
...         ]
...     }
... )

Using custom engines:

>>> from haive.core.engine.aug_llm import AugLLMConfig
>>>
>>> custom_engines = {
...     "generate_sql": AugLLMConfig(
...         name="custom_sql_generator",
...         temperature=0.1,  # Lower temperature for SQL
...         model="gpt-4"
...     ),
...     **default_sql_engines  # Keep other defaults
>>> }
>>>
>>> config = SQLRAGConfig(
...     engines=custom_engines,
...     domain_name="analytics"
... )

Raises:: ValueError – If required engines are missing from configuration.

classmethod check_required_engines(v)¶

Validate that all required engines are present.

Parameters:: v (dict[str, haive.core.engine.aug_llm.AugLLMConfig]) – Dictionary of engine configurations.
Returns:: Validated engine configurations.
Return type:: Dict[str, AugLLMConfig]
Raises:: ValueError – If any required engines are missing.