Spaces:
Running
Running
devjas1
commited on
Commit
·
986fb2a
1
Parent(s):
dccf8f7
(FEAT): Add README and implement SQL engine tool for AI agent interaction
Browse files- README.md +200 -0
- text_to_sql.py +130 -0
README.md
ADDED
|
@@ -0,0 +1,200 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Building Smart Agents: Turning Natural Language into Database Actions
|
| 2 |
+
|
| 3 |
+
**Goal**: Understand how to build AI agents that can interpret natural language questions, interact with a database using SQL, and provide relevant answers.
|
| 4 |
+
|
| 5 |
+
## What is This Code Doing?
|
| 6 |
+
|
| 7 |
+
At its core, this Python script demonstrates how to create a simple AI agent that can understand questions posed in plain English and then generate and execute SQL queries to retrieve answers from a database. It sets up a simulated scenario with receipts and waiters tables in an in-memory SQLite database, populates them with sample data, and then tasks two different AI models (via `CodeAgent`) to answer questions about this data.
|
| 8 |
+
|
| 9 |
+
**Specifically, the code performs the following actions:**
|
| 10 |
+
|
| 11 |
+
1. Database Setup:
|
| 12 |
+
|
| 13 |
+
- Initializes an in-memory SQLite database using `SQLAlchemy`.
|
| 14 |
+
- Defines two database tables: receipts (with columns for receipt_id, customer_name, price, tip) and waiters (with receipt_id, waiter_name).
|
| 15 |
+
- Populates these tables with sample data.
|
| 16 |
+
|
| 17 |
+
2. Tool Creation (`sql_engine`):
|
| 18 |
+
|
| 19 |
+
- It defines a Python function called `sql_engine` which is decorated as a `@tool`. This means the AI agent can "call" this function to perform specific actions.
|
| 20 |
+
- The `sql_engine` tool takes a SQL query string as input, executes it against the database, and returns the results as a string. Crucially, the tool's description is vital, as it tells the AI agent how and when to use it, including the available tables and their columns.
|
| 21 |
+
|
| 22 |
+
3. AI Agent Instantiation and Task Execution:
|
| 23 |
+
|
| 24 |
+
- It creates instances of `CodeAgent` using different Large Language Models (LLMs) like meta-llama/Llama-3.1-8B-Instruct and Qwen/Qwen2.5-Coder-32B-Instruct.
|
| 25 |
+
- Each agent is provided with the `sql_engine` tool, giving it the capability to interact with the database.
|
| 26 |
+
- The agents are then given natural language prompts, such as "Can you give me the name of the client who got the most expensive receipt?" or "Which waiter got more total money from tips?".
|
| 27 |
+
|
| 28 |
+
4. Problem Solving and Self-Correction:
|
| 29 |
+
|
| 30 |
+
- The output shows how the first agent (Llama-3.1) iteratively attempts to solve the problem. When its initial SQL generation or result parsing fails, it analyzes the error messages (e.g., "syntax error," "could not convert string to float") and revises its approach in subsequent steps. This demonstrates a form of self-correction.
|
| 31 |
+
- The second agent (Qwen2.5) successfully performs a SQL JOIN operation and then processes the results in Python to find the answer.
|
| 32 |
+
|
| 33 |
+
## How Does This Code Work? (The Mechanics of an AI Agent)
|
| 34 |
+
|
| 35 |
+
This code leverages the power of Large Language Models (LLMs) combined with tool-use capabilities to create an AI agent. Here's the "how":
|
| 36 |
+
|
| 37 |
+
1. **The LLM as the "Brain":** The `CodeAgent` uses an LLM (like Llama-3.1 or Qwen2.5) as its core reasoning engine. The LLM's extensive training on vast amounts of text enables it to:
|
| 38 |
+
|
| 39 |
+
- Understand natural language prompts (e.g., "most expensive receipt").
|
| 40 |
+
- Generate code (in this case, SQL queries and Python logic for parsing/calculation) to achieve a goal.
|
| 41 |
+
- Reason about execution errors and potentially self-correct its generated code.
|
| 42 |
+
|
| 43 |
+
2. **Tools as "Hands" or "Senses":** The `@tool` decorator from smolagents (or similar frameworks) allows the LLM to interact with the external world beyond just generating text.
|
| 44 |
+
|
| 45 |
+
- The `sql_engine` function acts as a specialized tool that encapsulates the logic for database interaction.s
|
| 46 |
+
- When the LLM determines that a database query is needed to answer a user's question, it "calls" the `sql_engine` tool, generating the appropriate SQL query string as an argument.
|
| 47 |
+
- The LLM relies heavily on the description attribute of the tool. This description tells the LLM what the tool does, what arguments it expects, and importantly, the schema of the database it can query. A clear and accurate description is paramount for the agent to use the tool correctly.
|
| 48 |
+
|
| 49 |
+
3. **The Agentic Loop:** The `CodeAgent` operates in a loop:
|
| 50 |
+
|
| 51 |
+
- It receives a prompt.
|
| 52 |
+
- It uses the LLM to think about the problem and decide if a tool is needed.
|
| 53 |
+
- If a tool is needed, the LLM generates the code (e.g., sql_engine("SELECT ...")) to use it.
|
| 54 |
+
- This generated code is then executed.
|
| 55 |
+
- The output of the code (e.g., database results, error messages) is fed back to the LLM.
|
| 56 |
+
- The LLM then uses this feedback to either refine its approach (self-correction), generate more code, or formulate a final answer.
|
| 57 |
+
|
| 58 |
+
4. **SQLAlchemy for Database Abstraction:**
|
| 59 |
+
|
| 60 |
+
- `SQLAlchemy` is used here to define the database schema (tables and columns) in a Pythonic way. This makes it easier to manage the database structure programmatically.
|
| 61 |
+
- It also provides an Object Relational Mapper (ORM) and a SQL Expression Language, which simplifies executing SQL queries and handling database connections (though in our sql_engine tool, we're directly using raw SQL for simplicity).
|
| 62 |
+
|
| 63 |
+
## Why Is This Approach Useful for AI Agents?
|
| 64 |
+
|
| 65 |
+
This method of combining LLMs with tools is a powerful paradigm for building sophisticated AI agents for several key reasons:
|
| 66 |
+
|
| 67 |
+
- **Extending LLM Capabilities:** LLMs are great at text understanding and generation, but they don't inherently "know" how to interact with a database, browse the internet, or perform complex calculations. Tools provide these external capabilities, essentially giving the LLM "skills" beyond just language.
|
| 68 |
+
- **Grounding in Factual Data:** By providing a sql_engine tool, the LLM can access up-to-date and specific information stored in a database, preventing hallucinations and ensuring factual accuracy, which LLMs alone sometimes struggle with.
|
| 69 |
+
- **Complex Task Execution:** Agents can break down complex problems into smaller steps, using different tools as needed. For example, in the second question, the agent first used SQL to fetch raw data and then used Python code (generated by the LLM) to perform an aggregation.
|
| 70 |
+
- **Reduced Development Effort:** Instead of hand-coding every possible interaction or query, you define the tools and let the LLM figure out how to use them to achieve the desired outcome from natural language. This abstracts away much of the conditional logic usually required in traditional programming.
|
| 71 |
+
- **Adaptability and Flexibility:** If your database schema changes, you primarily update the tool's description. The LLM, given the updated description, can often adapt its SQL generation without needing extensive code changes.
|
| 72 |
+
- **Self-Correction and Robustness:** The ability of the agent to analyze `Execution logs` and correct its own errors makes the system more robust and capable of handling unexpected issues.
|
| 73 |
+
|
| 74 |
+
## How to Create Your Own AI Agents (Step-by-Step Tutorial)
|
| 75 |
+
|
| 76 |
+
Let's break down the process of creating an AI agent similar to the one demonstrated:
|
| 77 |
+
|
| 78 |
+
### Step 1: Define Your Data and Environment
|
| 79 |
+
|
| 80 |
+
First, you need to decide what external systems your agent will interact with. In this example, it's a database.
|
| 81 |
+
|
| 82 |
+
**Set up your database (or other external system):**
|
| 83 |
+
|
| 84 |
+
- Decide on your database type (e.g., SQLite, PostgreSQL, MySQL).
|
| 85 |
+
- Use a library like SQLAlchemy to define your database schema (tables, columns, data types, primary keys).
|
| 86 |
+
- Populate your database with initial data.
|
| 87 |
+
|
| 88 |
+
```python
|
| 89 |
+
from sqlalchemy import create_engine, MetaData, Table, Column, String, Integer, Float
|
| 90 |
+
|
| 91 |
+
engine = create_engine("sqlite:///:memory:") # In-memory database for this example
|
| 92 |
+
metadata_obj = MetaData()
|
| 93 |
+
|
| 94 |
+
# Define your tables
|
| 95 |
+
|
| 96 |
+
receipts = Table(
|
| 97 |
+
"receipts",
|
| 98 |
+
metadata_obj,
|
| 99 |
+
Column("receipt_id", Integer, primary_key=True),
|
| 100 |
+
Column("customer_name", String(255)),
|
| 101 |
+
Column("price", Float),
|
| 102 |
+
Column("tip", Float),
|
| 103 |
+
)
|
| 104 |
+
|
| 105 |
+
# Add other tables as needed, e.g., 'waiters'
|
| 106 |
+
|
| 107 |
+
metadata_obj.create_all(engine) # Create tables in the database
|
| 108 |
+
|
| 109 |
+
# Insert initial data
|
| 110 |
+
|
| 111 |
+
# ... (as shown in the example code)
|
| 112 |
+
```
|
| 113 |
+
|
| 114 |
+
### Step 2: Create Tools for Your Agent
|
| 115 |
+
|
| 116 |
+
Tools are the interfaces your AI agent uses to interact with external systems. Each tool should be a Python function.
|
| 117 |
+
|
| 118 |
+
- **Define a Python function for your tool:** This function should encapsulate the logic for a specific action (e.g., querying the database, making an API call, reading a file).
|
| 119 |
+
- **Decorate it with `@tool`:** This tells the smolagents framework that this function is an agent tool.
|
| 120 |
+
- **Provide a clear and detailed docstring:** This docstring is crucial. The LLM will read this description to understand what the tool does, its parameters, and what information it provides. Include:
|
| 121 |
+
- A high-level explanation of the tool's purpose.
|
| 122 |
+
- A detailed description of the data it can access or modify (e.g., table schemas, API endpoints).
|
| 123 |
+
- Examples if necessary.
|
| 124 |
+
|
| 125 |
+
```python
|
| 126 |
+
from smolagents import tool
|
| 127 |
+
from sqlalchemy import text # Make sure to import text for raw SQL execution
|
| 128 |
+
|
| 129 |
+
`@tool`
|
| 130 |
+
def sql_engine(query: str) -> str:
|
| 131 |
+
"""
|
| 132 |
+
Allows you to perform SQL queries on the database.
|
| 133 |
+
Returns a string representation of the query results.
|
| 134 |
+
|
| 135 |
+
Available tables and their schemas:
|
| 136 |
+
Table 'receipts':
|
| 137 |
+
Columns:
|
| 138 |
+
- receipt_id: INTEGER (Primary Key)
|
| 139 |
+
- customer_name: VARCHAR(255)
|
| 140 |
+
- price: FLOAT
|
| 141 |
+
- tip: FLOAT
|
| 142 |
+
|
| 143 |
+
Table 'waiters':
|
| 144 |
+
Columns:
|
| 145 |
+
- receipt_id: INTEGER (Primary Key)
|
| 146 |
+
- waiter_name: VARCHAR(16) (Primary Key)
|
| 147 |
+
|
| 148 |
+
Args:
|
| 149 |
+
query: The SQL query string to execute. This should be correct SQL.
|
| 150 |
+
Example: "SELECT customer_name FROM receipts WHERE price > 10.0"
|
| 151 |
+
"""
|
| 152 |
+
output = ""
|
| 153 |
+
with engine.connect() as con:
|
| 154 |
+
rows = con.execute(text(query)) # Use text() for raw SQL queries
|
| 155 |
+
for row in rows:
|
| 156 |
+
output += "\n" + str(row)
|
| 157 |
+
return output
|
| 158 |
+
```
|
| 159 |
+
|
| 160 |
+
- **Tip:** Keep the tool's output as consistent and parseable as possible. The LLM will need to interpret this string.
|
| 161 |
+
|
| 162 |
+
### Step 3: Instantiate Your AI Agent
|
| 163 |
+
|
| 164 |
+
Now, you bring everything together by creating your `CodeAgent` and providing it with the necessary components.
|
| 165 |
+
|
| 166 |
+
- **Choose your LLM:** Select a suitable language model. Different models have different capabilities and cost implications.
|
| 167 |
+
- **Initialize `CodeAgent`:** Pass your list of tools and your chosen LLM.
|
| 168 |
+
|
| 169 |
+
```python
|
| 170 |
+
from smolagents import `CodeAgent`, InferenceClientModel
|
| 171 |
+
|
| 172 |
+
agent = `CodeAgent`(
|
| 173 |
+
tools=[sql_engine], # Provide the list of tools your agent can use
|
| 174 |
+
model=InferenceClientModel(model_id="meta-llama/Llama-3.1-8B-Instruct"), # Or another model
|
| 175 |
+
)
|
| 176 |
+
```
|
| 177 |
+
|
| 178 |
+
### Step 4: Give Your Agent a Task
|
| 179 |
+
|
| 180 |
+
Once your agent is set up, you can give it natural language prompts, and it will attempt to solve them using its tools and reasoning.
|
| 181 |
+
|
| 182 |
+
- **Run the agent with a prompt:**
|
| 183 |
+
|
| 184 |
+
```python
|
| 185 |
+
agent.run("Can you tell me the customer with the highest tip amount?")
|
| 186 |
+
```
|
| 187 |
+
|
| 188 |
+
The agent will then go through its loop: interpret the question, decide if sql_engine is needed, generate SQL, execute it, process results, and formulate an answer.
|
| 189 |
+
|
| 190 |
+
### Step 5: Iterate and Refine
|
| 191 |
+
|
| 192 |
+
Building effective agents is an iterative process.
|
| 193 |
+
|
| 194 |
+
- **Review Agent Output and Execution logs:** Pay close attention to the `Execution logs`. If the agent fails, analyze the error messages to understand why.
|
| 195 |
+
- **Improve Tool Descriptions:** Often, agents fail because the tool's description isn't clear enough, or doesn't provide all the necessary context (e.g., table schemas, data types). Make your tool descriptions as precise and comprehensive as possible.
|
| 196 |
+
- **Refine Prompts:** Sometimes, rephrasing the initial prompt can help the agent understand the task better.
|
| 197 |
+
- **Consider Model Choice:** Different LLMs excel at different tasks. If one model struggles, experimenting with another might yield better results.
|
| 198 |
+
- **Handle Complex Output Parsing:** As seen in the example, parsing the string output from tools can be tricky. The agent might need to generate additional Python code to properly extract information from the tool's raw string output.
|
| 199 |
+
|
| 200 |
+
By following these steps, you can start building powerful AI agents that can automate tasks, answer questions, and interact with various systems using natural language, making them incredibly versatile for many applications.
|
text_to_sql.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from dotenv import load_dotenv
|
| 2 |
+
|
| 3 |
+
load_dotenv()
|
| 4 |
+
|
| 5 |
+
from sqlalchemy import (
|
| 6 |
+
create_engine,
|
| 7 |
+
MetaData,
|
| 8 |
+
Table,
|
| 9 |
+
Column,
|
| 10 |
+
String,
|
| 11 |
+
Integer,
|
| 12 |
+
Float,
|
| 13 |
+
insert,
|
| 14 |
+
text, # Import text to be used in sql_engine tool
|
| 15 |
+
)
|
| 16 |
+
from sqlalchemy import inspect
|
| 17 |
+
from smolagents import tool, CodeAgent, InferenceClientModel
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
engine = create_engine("sqlite:///:memory:")
|
| 21 |
+
metadata_obj = MetaData()
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def insert_rows_into_table(data_rows, target_table, engine=engine):
|
| 25 |
+
for row in data_rows:
|
| 26 |
+
stmt = insert(target_table).values(**row)
|
| 27 |
+
with engine.begin() as connection:
|
| 28 |
+
connection.execute(stmt)
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
table_name = "receipts"
|
| 32 |
+
# Corrected receipts table definition
|
| 33 |
+
receipts = Table(
|
| 34 |
+
table_name,
|
| 35 |
+
metadata_obj,
|
| 36 |
+
Column("receipt_id", Integer, primary_key=True),
|
| 37 |
+
Column("customer_name", String(255)), # Added customer_name column
|
| 38 |
+
Column("price", Float),
|
| 39 |
+
Column("tip", Float),
|
| 40 |
+
)
|
| 41 |
+
metadata_obj.create_all(engine)
|
| 42 |
+
|
| 43 |
+
rows = [
|
| 44 |
+
{"receipt_id": 1, "customer_name": "Alex Payne", "price": 12.06, "tip": 1.20},
|
| 45 |
+
{"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
|
| 46 |
+
{"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
|
| 47 |
+
{"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
|
| 48 |
+
]
|
| 49 |
+
insert_rows_into_table(rows, receipts)
|
| 50 |
+
|
| 51 |
+
inspector = inspect(engine)
|
| 52 |
+
columns_info = [(col["name"], col["type"]) for col in inspector.get_columns("receipts")]
|
| 53 |
+
|
| 54 |
+
table_description = "Columns:\n" + "\n".join(
|
| 55 |
+
[f" - {name}: {col_type}" for name, col_type in columns_info]
|
| 56 |
+
)
|
| 57 |
+
print(table_description)
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
@tool
|
| 61 |
+
def sql_engine(query: str) -> str:
|
| 62 |
+
"""
|
| 63 |
+
Allows you to perform SQL queries on the table. Returns a string representation of
|
| 64 |
+
the result. The table is named 'receipts'. It's description is as follows:
|
| 65 |
+
Columns:
|
| 66 |
+
- receipt_id: INTEGER
|
| 67 |
+
- customer_name: STRING
|
| 68 |
+
- price: FLOAT
|
| 69 |
+
- tip: FLOAT
|
| 70 |
+
|
| 71 |
+
Args:
|
| 72 |
+
query: The query to perform. This should be correct SQL.
|
| 73 |
+
|
| 74 |
+
"""
|
| 75 |
+
output = ""
|
| 76 |
+
with engine.connect() as con:
|
| 77 |
+
rows = con.execute(text(query))
|
| 78 |
+
for row in rows:
|
| 79 |
+
output += "\n" + str(row)
|
| 80 |
+
|
| 81 |
+
return output
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
agent = CodeAgent(
|
| 85 |
+
tools=[sql_engine],
|
| 86 |
+
model=InferenceClientModel(model_id="meta-llama/Llama-3.1-8B-Instruct"),
|
| 87 |
+
)
|
| 88 |
+
agent.run("Can you give me the name of the client who got the most expensive receipt?")
|
| 89 |
+
|
| 90 |
+
table_name = "waiters"
|
| 91 |
+
waiters = Table(
|
| 92 |
+
table_name,
|
| 93 |
+
metadata_obj,
|
| 94 |
+
Column("receipt_id", Integer, primary_key=True),
|
| 95 |
+
Column("waiter_name", String(16), primary_key=True),
|
| 96 |
+
)
|
| 97 |
+
metadata_obj.create_all(engine)
|
| 98 |
+
|
| 99 |
+
rows = [
|
| 100 |
+
{"receipt_id": 1, "waiter_name": "Corey Johnson"},
|
| 101 |
+
{"receipt_id": 2, "waiter_name": "Michael Watts"},
|
| 102 |
+
{"receipt_id": 3, "waiter_name": "Michael Watts"},
|
| 103 |
+
{"receipt_id": 4, "waiter_name": "Margaret James"},
|
| 104 |
+
]
|
| 105 |
+
insert_rows_into_table(rows, waiters)
|
| 106 |
+
|
| 107 |
+
updated_description = """Allows you to perform SQL queries on the table. Beware that this tool's output is a string representation of the execution output.
|
| 108 |
+
It can use the following tables:"""
|
| 109 |
+
|
| 110 |
+
inspector = inspect(engine)
|
| 111 |
+
for table in ["receipts", "waiters"]:
|
| 112 |
+
columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
|
| 113 |
+
|
| 114 |
+
table_description = f"Table '{table}':\n"
|
| 115 |
+
|
| 116 |
+
table_description += "Columns:\n" + "\n".join(
|
| 117 |
+
[f" - {name}: {col_type}" for name, col_type in columns_info]
|
| 118 |
+
)
|
| 119 |
+
updated_description += "\n\n" + table_description
|
| 120 |
+
|
| 121 |
+
print(updated_description)
|
| 122 |
+
|
| 123 |
+
sql_engine.description = updated_description
|
| 124 |
+
|
| 125 |
+
agent = CodeAgent(
|
| 126 |
+
tools=[sql_engine],
|
| 127 |
+
model=InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
agent.run("Which waiter got more total money from tips?")
|