feat: sql tool

app.py

@@ -4,6 +4,119 @@ import os
 from huggingface_hub import login
 login(token=os.getenv("HUGGINGFACEHUB_API_KEY"))
 
+
+# %pip install gradio smolagents sqlalchemy
+
+# "meta-llama/Llama-3.2-3B does not work
+model_name="meta-llama/Llama-3.2-3B-Instruct"   # 6.5G
+
+from sqlalchemy import (
+    create_engine,
+    MetaData,
+    Table,
+    Column,
+    DECIMAL,
+    TEXT,
+    insert,
+    inspect,
+    text,
+)
+
+def insert_rows_into_table(rows, table, engine):
+    for row in rows:
+        stmt = insert(table).values(**row)
+        with engine.begin() as connection:
+            connection.execute(stmt)
+
+def prepare_payment_table(engine, metadata):
+    inspector = inspect(engine)
+    table_name = "payments"
+    if not inspector.has_table(table_name):
+        table = Table(
+            table_name,
+            metadata,
+            Column("id", TEXT, primary_key=True),
+            Column("amount", DECIMAL),
+            Column("created_at", TEXT),
+        )
+        metadata.create_all(engine)
+    else:
+        table = Table(table_name, metadata, autoload_with=engine)
+
+    rows = [
+        {"id": "payment-123", "amount": 100.0, "created_at": "2021-01-01 00:00:00"},
+        {"id": "payment-abc-12", "amount": 200.0, "created_at": "2021-01-02 00:00:00"},
+    ]
+    insert_rows_into_table(rows, table, engine)
+
+    return table
+
+def prepare_payout_table(engine, metadata):
+    inspector = inspect(engine)
+    table_name = "payouts"
+    if not inspector.has_table(table_name):
+        table = Table(
+            table_name,
+            metadata,
+            Column("id", TEXT, primary_key=True),
+            Column("amount", DECIMAL),
+            Column("created_at", TEXT),
+        )
+        metadata.create_all(engine)
+    else:
+        table = Table(table_name, metadata, autoload_with=engine)
+
+    rows = [
+        {"id": "payout-abc456", "amount": 50.0, "created_at": "2021-01-01 00:00:00"},
+        {"id": "payout-b2c2", "amount": 100.0, "created_at": "2021-01-02 00:00:00"},
+    ]
+    insert_rows_into_table(rows, table, engine)
+
+    return table
+
+engine = create_engine("sqlite:///:memory:")
+metadata = MetaData()
+
+prepare_payment_table(engine, metadata)
+prepare_payout_table(engine, metadata)
+
+from smolagents import tool, CodeAgent, InferenceClientModel
+
+@tool
+def sql_engine(query: str) -> str:
+    """
+    Allows you to perform SQL queries on the table. Returns a string representation of the result.
+
+    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
+
+
+agent = CodeAgent(
+    tools=[sql_engine],
+    model=InferenceClientModel(model_id=model_name),
+)
+
+tool_description = """Allows you to perform SQL queries on the table. Beware that this tool's output is a string representation of the execution output.
+It can use the following tables:"""
+
+inspector = inspect(engine)
+for table in ["payments", "payouts"]:
+    columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]
+
+    table_description = f"Table '{table}':\n"
+
+    table_description += "Columns:\n" + "\n".join([f"  - {name}: {col_type}" for name, col_type in columns_info])
+    tool_description += "\n\n" + table_description
+
+print("SQL tool description", tool_description)
+
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import torch
@@ -18,12 +131,6 @@ def get_device_type() -> str:
     else:
         return "cpu"
 
-# # HuggingFaceTB/SmolLM2-135M-Instruct
-# model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" # 15G
-
-# "meta-llama/Llama-3.2-3B does not work
-model_name="meta-llama/Llama-3.2-3B-Instruct"   # 6.5G
-
 device = get_device_type()
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name).to(device, dtype=torch.float16)
@@ -81,6 +188,26 @@ examples = [
     "I am having trouble with my transaction",
 ]
 
+import json, re
+
+def extract_transaction_info(response):
+    try:
+        json_match = re.search(r'\{.*\}', response, re.DOTALL)
+        if json_match:
+            resp_json = json.loads(json_match.group())
+            found = resp_json.get("found", False)
+            if found == False:
+                return None
+            
+            transaction_id = resp_json.get("transaction_id")
+            if transaction_id:
+                # Ensure the transaction ID is a string
+                return str(transaction_id).strip()
+            else:
+                return None
+    except Exception as e:
+        return None
+
 def predict(message, history):
     # Always inject the user message into the system prompt's {input} placeholder
     sys_prompt = system_prompt.replace("{input}", message)
@@ -105,17 +232,32 @@ def predict(message, history):
     
     # Extract only the assistant's message (after the last user message)
     if "<|start_header_id|>assistant<|end_header_id|>" in decoded:
-        response = decoded.split("<|start_header_id|>assistant<|end_header_id|>")[-1]
-        response = response.replace("<|eot_id|>", "").strip()
+        analysis_response = decoded.split("<|start_header_id|>assistant<|end_header_id|>")[-1]
+        analysis_response = analysis_response.replace("<|eot_id|>", "").strip()
     elif "<|im_start|>assistant" in decoded:
         # This works for most chat templates that append the assistant's reply at the end
-        response = decoded.split("<|im_start|>assistant")[-1]
-        response = response.replace("<|im_end|>", "").strip()
+        analysis_response = decoded.split("<|im_start|>assistant")[-1]
+        analysis_response = analysis_response.replace("<|im_end|>", "").strip()
     else:
         # Fallback: just return the decoded output
-        response = decoded.strip()
+        analysis_response = decoded.strip()
+
+    transaction_id = extract_transaction_info(analysis_response)
+    if transaction_id == None:
+        return analysis_response
         
-    return response
+    # If we successfully extracted a transaction ID, we can invoke the SQL tool
+    sql_response = agent.run(f"find a record with id {transaction_id} from correct table in the database")
+
+    print(f"SQL response: {sql_response}\n")
+
+    try:
+        analysis_json = json.loads(analysis_response)
+        analysis_json["database_result"] = sql_response
+        return json.dumps(analysis_json, ensure_ascii=False)
+    except Exception:
+        # If analysis_response is not valid JSON, return both as plain text
+        return analysis_response
 
 demo = gr.ChatInterface(predict, type="messages", examples=examples)
 

requirements.txt

@@ -5,4 +5,6 @@ safetensors>=0.4.5
 transformers==4.49.0
 gradio>=5.23.0
 datasets
-bitsandbytes
+bitsandbytes
+smolagents
+sqlalchemy