chore: code agent

app.py

@@ -247,24 +247,30 @@ def predict(message, history):
         return analysis_response
         
     # 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.
+    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.
 
-    #     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"}}
+        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)
+        Input: transaction_type: payout, transaction_id: null
+        Record in payouts: None
+        Response: null
+    """
 
-    # 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}')
+    try:
+        sql_response = agent.run(sql_prompt)
+    except Exception as e:
+        print(f"An error occurred while running the SQL agent: {e}")
+        try:
+            analysis_json = json.loads(analysis_response)
+            analysis_json["database_result"] = "Error running SQL agent: " + str(e)
+            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
 
     # print(f"SQL response: {sql_response}\n")