agents.rag.db_rag.sql_rag.utils¶

Utility functions for SQL RAG Agent.

This module provides helper functions for database operations, toolkit creation, error handling, and schema exploration. These utilities support the main agent workflow with reusable functionality.

Example

Creating SQL toolkit:

>>> from haive.agents.rag.db_rag.sql_rag.utils import create_sql_toolkit
>>> from haive.agents.rag.db_rag.sql_rag.config import SQLDatabaseConfig
>>>
>>> db_config = SQLDatabaseConfig(db_uri="sqlite:///example.db")
>>> toolkit = create_sql_toolkit(db_config)
>>> tools = toolkit.get_tools()
>>> print(f"Available tools: {[tool.name for tool in tools]}")
Available tools: ['sql_db_query', 'sql_db_schema', 'sql_db_list_tables', ...]

Functions¶

`create_sql_toolkit`(db_config[, llm_config])	Create a SQL Database Toolkit using the provided configuration.
`create_tool_node_with_fallback`(tools)	Create a ToolNode with a fallback to handle errors gracefully.
`explore_database_schema`(db)	Explore the database schema thoroughly to get comprehensive information.
`get_all_toolkit_tools`(toolkit)	Get all tools from a SQLDatabaseToolkit.
`handle_tool_error`(state)	Handle tool execution errors by surfacing them to the agent.

Module Contents¶

agents.rag.db_rag.sql_rag.utils.create_sql_toolkit(db_config, llm_config=None)¶

Create a SQL Database Toolkit using the provided configuration.

This function creates a LangChain SQLDatabaseToolkit that provides various tools for interacting with SQL databases including query execution, schema inspection, and table listing.

Parameters:

db_config (SQLDatabaseConfig) – Configuration for the SQL database connection.
llm_config (Optional[AugLLMConfig]) – Configuration for the LLM to use with the toolkit. If None, defaults to OpenAI GPT-4.

Returns:

Configured toolkit instance with database tools.

Return type:

SQLDatabaseToolkit

Raises:

ValueError – If database connection fails or is invalid.

Example

Creating toolkit with custom LLM:

>>> from haive.core.engine.aug_llm import AugLLMConfig
>>>
>>> llm_config = AugLLMConfig(
...     model="gpt-4",
...     temperature=0.1  # Low temperature for SQL generation
... )
>>>
>>> toolkit = create_sql_toolkit(db_config, llm_config)
>>>
>>> # Use toolkit to query schema
>>> schema_tool = next(t for t in toolkit.get_tools() if t.name == "sql_db_schema")
>>> schema = schema_tool.invoke({"table_names": "customers"})

agents.rag.db_rag.sql_rag.utils.create_tool_node_with_fallback(tools)¶

Create a ToolNode with a fallback to handle errors gracefully.

This function wraps tools in a ToolNode with error handling that surfaces errors back to the agent for correction rather than failing.

Parameters:: tools (Union[BaseTool, List[BaseTool]]) – Tool or list of tools to create a node for.
Returns:: ToolNode wrapped with fallback error handling.
Return type:: RunnableWithFallbacks

Example

>>> query_tool = toolkit.get_tool("sql_db_query")
>>> safe_tool_node = create_tool_node_with_fallback(query_tool)
>>>
>>> # If query fails, error is captured and returned
>>> result = safe_tool_node.invoke({
...     "messages": [AIMessage(tool_calls=[{"name": "sql_db_query", "args": {"query": "INVALID SQL"}}])]
... })

agents.rag.db_rag.sql_rag.utils.explore_database_schema(db)¶

Explore the database schema thoroughly to get comprehensive information.

This function performs a deep exploration of the database schema, including tables, columns, sample data, and potential relationships. It’s used to provide context for SQL generation.

Parameters:

db (SQLDatabase) – SQLDatabase instance to explore.

Returns:

Detailed schema information including:

tables: List of table names
table_info: Detailed schema for each table
table_samples: Sample rows from each table
relationships: Detected foreign key relationships
dialect: SQL dialect being used

Return type:

Dict[str, Any]

Example

>>> db = SQLDatabase.from_uri("sqlite:///northwind.db")
>>> schema_info = explore_database_schema(db)
>>>
>>> print(f"Found {len(schema_info['tables'])} tables")
Found 8 tables
>>>
>>> print(f"Dialect: {schema_info['dialect']}")
Dialect: sqlite
>>>
>>> # Show detected relationships
>>> for rel in schema_info['relationships']:
...     print(f"{rel['from_table']}.{rel['column']} -> {rel['to_table']}")
orders.customer_id -> customers
order_details.product_id -> products

Note

Relationship detection uses heuristics based on naming conventions and may not catch all foreign keys. For production use, consider querying database metadata directly.

agents.rag.db_rag.sql_rag.utils.get_all_toolkit_tools(toolkit)¶

Get all tools from a SQLDatabaseToolkit.

Parameters:: toolkit (SQLDatabaseToolkit) – The toolkit instance to extract tools from.
Returns:: List of all available tools in the toolkit.
Return type:: List[BaseTool]

Example

>>> toolkit = create_sql_toolkit(db_config)
>>> tools = get_all_toolkit_tools(toolkit)
>>> for tool in tools:
...     print(f"Tool: {tool.name} - {tool.description}")
Tool: sql_db_query - Execute a SQL query against the database
Tool: sql_db_schema - Get the schema of specific tables
Tool: sql_db_list_tables - List all tables in the database

agents.rag.db_rag.sql_rag.utils.handle_tool_error(state)¶

Handle tool execution errors by surfacing them to the agent.

This function processes tool execution errors and formats them for the agent to understand and potentially correct.

Parameters:: state (Dict[str, Any]) – Current state with error information.
Returns:: Modified state with formatted error messages.
Return type:: Dict[str, Any]

Example

>>> state = {
...     "error": ValueError("Table 'users' does not exist"),
...     "messages": [AIMessage(tool_calls=[{"id": "123", "name": "sql_db_query"}])]
... }
>>> result = handle_tool_error(state)
>>> print(result["messages"][0].content)
'Error: ValueError("Table \'users\' does not exist")\\nPlease fix your mistakes.'