devjas1 commited on
Commit
986fb2a
·
1 Parent(s): dccf8f7

(FEAT): Add README and implement SQL engine tool for AI agent interaction

Browse files
Files changed (2) hide show
  1. README.md +200 -0
  2. 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?")