feat: sql agent

app.py

@@ -77,8 +77,8 @@ def prepare_payout_table(engine, metadata):
 engine = create_engine("sqlite:///:memory:")
 metadata = MetaData()
 
-prepare_payment_table(engine, metadata)
-prepare_payout_table(engine, metadata)
+payment_table=prepare_payment_table(engine, metadata)
+payout_table=prepare_payout_table(engine, metadata)
 
 from smolagents import tool, CodeAgent, InferenceClientModel
 
@@ -115,7 +115,7 @@ for table in ["payments", "payouts"]:
     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)
+# print("SQL tool description", tool_description)
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
@@ -197,16 +197,16 @@ def extract_transaction_info(response):
             resp_json = json.loads(json_match.group())
             found = resp_json.get("found", False)
             if found == False:
-                return None
+                return None, None
             
             transaction_id = resp_json.get("transaction_id")
-            if transaction_id:
-                # Ensure the transaction ID is a string
-                return str(transaction_id).strip()
+            transaction_type = resp_json.get("transaction_type")
+            if transaction_id and transaction_type:
+                return str(transaction_id).strip(), transaction_type.strip()
             else:
-                return None
+                return None, None
     except Exception as e:
-        return None
+        return None, None
 
 def predict(message, history):
     # Always inject the user message into the system prompt's {input} placeholder
@@ -242,14 +242,31 @@ def predict(message, history):
         # Fallback: just return the decoded output
         analysis_response = decoded.strip()
 
-    transaction_id = extract_transaction_info(analysis_response)
+    transaction_id, transaction_type = extract_transaction_info(analysis_response)
     if transaction_id == None:
         return analysis_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")
+    # payment is required to use the SQL agent, error: Subscribe to PRO to get 20x more monthly
+    # sql_prompt = f"""
+    #     Given the '{transaction_type}', find the corresponding table name in the database, the find a record with given id '{transaction_id}' from the table, return the record as JSON. If there is no record, return null.
 
-    print(f"SQL response: {sql_response}\n")
+    #     Example:
+    #     Input: transaction_type: payout, transaction_id: payout-abc456
+    #     Record in payouts: {{"id": "payout-abc456", "amount": 50.0, "created_at": "2021-01-01 00:00:00"}}
+    #     Response: {{"id": "payout-abc456", "amount": 50.0, "created_at": "2021-01-01 00:00:00"}}
+
+    #     Input: transaction_type: payout, transaction_id: null
+    #     Record in payouts: None
+    #     Response: null
+    # """
+
+    # sql_response = agent.run(sql_prompt)
+
+    # Directly use the sql_engine tool without CodeAgent to avoid payment
+    table = payment_table if transaction_type == "payment" else payout_table
+    sql_response = table.select().where(table.c.id == '{transaction_id}')
+
+    # print(f"SQL response: {sql_response}\n")
 
     try:
         analysis_json = json.loads(analysis_response)