75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
from sqlalchemy import (
|
|
create_engine,
|
|
MetaData,
|
|
Table,
|
|
Column,
|
|
String,
|
|
Integer,
|
|
Float,
|
|
insert,
|
|
inspect,
|
|
text,
|
|
)
|
|
|
|
engine = create_engine("sqlite:///:memory:")
|
|
metadata_obj = MetaData()
|
|
|
|
# create city SQL table
|
|
table_name = "receipts"
|
|
receipts = Table(
|
|
table_name,
|
|
metadata_obj,
|
|
Column("receipt_id", Integer, primary_key=True),
|
|
Column("customer_name", String(16), primary_key=True),
|
|
Column("price", Float),
|
|
Column("tip", Float),
|
|
)
|
|
metadata_obj.create_all(engine)
|
|
|
|
rows = [
|
|
{"receipt_id": 1, "customer_name": "Alan Payne", "price": 12.06, "tip": 1.20},
|
|
{"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
|
|
{"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
|
|
{"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
|
|
]
|
|
for row in rows:
|
|
stmt = insert(receipts).values(**row)
|
|
with engine.begin() as connection:
|
|
cursor = connection.execute(stmt)
|
|
|
|
inspector = inspect(engine)
|
|
columns_info = [(col["name"], col["type"]) for col in inspector.get_columns("receipts")]
|
|
|
|
table_description = "Columns:\n" + "\n".join([f" - {name}: {col_type}" for name, col_type in columns_info])
|
|
print(table_description)
|
|
|
|
from smolagents import tool
|
|
|
|
@tool
|
|
def sql_engine(query: str) -> str:
|
|
"""
|
|
Allows you to perform SQL queries on the table. Returns a string representation of the result.
|
|
The table is named 'receipts'. Its description is as follows:
|
|
Columns:
|
|
- receipt_id: INTEGER
|
|
- customer_name: VARCHAR(16)
|
|
- price: FLOAT
|
|
- tip: FLOAT
|
|
|
|
Args:
|
|
query: The query to perform. This should be correct SQL.
|
|
"""
|
|
output = ""
|
|
with engine.connect() as con:
|
|
rows = con.execute(text(query))
|
|
for row in rows:
|
|
output += "\n" + str(row)
|
|
return output
|
|
|
|
from smolagents import CodeAgent, HfApiModel
|
|
|
|
agent = CodeAgent(
|
|
tools=[sql_engine],
|
|
model=HfApiModel("meta-llama/Meta-Llama-3.1-8B-Instruct"),
|
|
)
|
|
agent.run("Can you give me the name of the client who got the most expensive receipt?") |