Agent Foundry
All Problems

#9. Build a Database Query Tool

EasyTool Calling

The Problem

Your database assistant agent can query an employee database, but the current tool implementation is dangerously naive. It executes any SQL string the model generates — including DROP TABLE or DELETE FROM — with zero validation. It also returns raw Python tuples like [(1, 'Alice', 'Engineering', 95000)] instead of human-readable results. Your job is to fix the tool implementation so it only allows SELECT queries, prevents SQL injection, formats results with column headers, and handles empty result sets gracefully.

Examples

Example 1

User input: Show me all engineers and their salaries

Current (bad) output: [(1, 'Alice', 'Engineering', 95000), (3, 'Charlie', 'Engineering', 88000)] (raw tuples, no column names)

Expected (good) output:

Results (2 rows):
| id | name    | department  | salary |
|----|---------|-------------|--------|
| 1  | Alice   | Engineering | 95000  |
| 3  | Charlie | Engineering | 88000  |

Example 2

User input: (attacker tries) DROP TABLE employees; --

Current (bad) output: The table is dropped and all data is lost.

Expected (good) output: Error: Only SELECT queries are allowed. DROP, DELETE, INSERT, UPDATE, and ALTER statements are blocked.

Example 3

User input: Show me employees in the Legal department

Current (bad) output: [] (empty list, no explanation)

Expected (good) output: No results found. The query returned 0 rows.

Your Task

Fix the query_database tool so that it:

  • Only allows SELECT statements — rejects DROP, DELETE, INSERT, UPDATE, ALTER, etc.
  • Formats results as a readable table with column headers.
  • Returns a clear message when no rows match the query.
  • Has a proper docstring explaining its purpose, parameters, and restrictions.

Do not change the database schema, the agent setup, or the prompt.

Evaluation

Submissions are checked for the following:

  • SQL injection prevented: Dangerous SQL operations (DROP, DELETE, INSERT, UPDATE) are rejected with a clear error.
  • Results formatted: Query results include column headers and are formatted as readable text.
  • Handles empty results: Queries that return no rows produce a helpful "no results" message instead of an empty list.

Constraints

  • The tool must prevent SQL injection attacks
  • Only SELECT queries are allowed
  • Results must be formatted as a readable string, not raw tuples
Starter Code
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool
import sqlite3

llm = ChatOpenAI(model="gpt-4o-mini")

# Set up in-memory database with sample data
conn = sqlite3.connect(":memory:")
cursor = conn.cursor()
cursor.execute("CREATE TABLE employees (id INTEGER, name TEXT, department TEXT, salary INTEGER)")
cursor.executemany("INSERT INTO employees VALUES (?, ?, ?, ?)", [
    (1, "Alice", "Engineering", 95000),
    (2, "Bob", "Marketing", 72000),
    (3, "Charlie", "Engineering", 88000),
    (4, "Diana", "Sales", 67000),
    (5, "Eve", "Marketing", 78000),
])
conn.commit()

# BUG: This tool has no SQL injection protection, no query validation,
# and returns raw tuples instead of formatted results
@tool
def query_database(sql: str) -> str:
    """Run a SQL query against the employees database."""
    result = cursor.execute(sql).fetchall()
    return str(result)

tools = [query_database]

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a database assistant. Use the query tool to answer questions about employees."),
    ("human", "{input}"),
    ("placeholder", "{agent_scratchpad}"),
])

agent = create_tool_calling_agent(llm, tools, prompt)
executor = AgentExecutor(agent=agent, tools=tools)

result = executor.invoke({"input": "Show me all engineers and their salaries"})
print(result["output"])
Open in Google Colab
Evaluation Criteria0/3