import datetime
import os
from zoneinfo import ZoneInfo
from smolagents import OpenAIServerModel, ToolCallingAgent, tool
from sqlalchemy import (
Column,
Float,
Integer,
MetaData,
String,
Table,
create_engine,
insert,
text,
)
db_path = "receipts_db.sqlite"
if os.path.exists(db_path):
os.remove(db_path)
engine = create_engine(f"sqlite:///{db_path}")
metadata_obj = MetaData()
table_name = "receipts"
receipts = Table(
table_name,
metadata_obj,
Column("receipt_id", Integer, primary_key=True),
Column("customer_name", String(32)),
Column("price", Float),
Column("tip", Float),
Column("city", String(32)),
)
metadata_obj.create_all(engine)
rows = [
{
"receipt_id": 1,
"customer_name": "Alan Payne",
"price": 12.06,
"tip": 1.20,
"city": "New York",
},
# ... more rows ...
]
with engine.begin() as connection:
for row in rows:
connection.execute(insert(receipts).values(**row))
@tool
def sql_engine(query: str) -> str:
"""Execute SQL queries against the receipts table."""
with engine.connect() as con:
rows = con.execute(text(query))
return "\n".join(str(row) for row in rows)
@tool
def current_time(city: str) -> str:
"""Return the current time for a known city."""
mapping = {
"new york": "America/New_York",
"san francisco": "America/Los_Angeles",
"london": "Europe/London",
}
tz_identifier = mapping.get(city.lower())
if not tz_identifier:
return f"Timezone for {city} not available."
now = datetime.datetime.now(ZoneInfo(tz_identifier))
return now.strftime("%Y-%m-%d %H:%M:%S %Z%z")
sql_agent = ToolCallingAgent(
tools=[sql_engine, current_time],
model=OpenAIServerModel("gpt-4o"),
)
