Update tools/sql_tool.py

tools/sql_tool.py

@@ -22,8 +22,7 @@ def get_md_connection() -> duckdb.DuckDBPyConnection:
         )
     
     # 2. Connect to the MotherDuck service
-    # Note: Replace 'my_db' with your actual MotherDuck database name if necessary, 
-    # otherwise it connects to the default MotherDuck endpoint.
+    # Note: If you have a specific database name, you can adjust the connection string here.
     conn = duckdb.connect(f'md:?motherduck_token={token}')
     return conn
 
@@ -35,6 +34,7 @@ def run_duckdb_query(query: str) -> str:
     Runs a read-only SQL query against the connected MotherDuck database and returns the results as a string.
     The query must be valid DuckDB SQL. This tool only supports SELECT queries.
     """
+    conn = None
     try:
         conn = get_md_connection()
         
@@ -52,12 +52,14 @@ def run_duckdb_query(query: str) -> str:
         return result_df.to_string(index=False)
         
     except ConnectionError as e:
+        # Re-raise or handle specific connection errors
         return f"Connection Error: {e}"
     except Exception as e:
+        # Catch all other DuckDB or SQL execution errors
         return f"DuckDB Query Error: {e}"
     finally:
-        # Always close the connection
-        if 'conn' in locals() and conn:
+        # Always close the connection if it was successfully opened
+        if conn:
             conn.close()
 
 @tool
@@ -66,15 +68,17 @@ def get_table_schema(table_name: str = "positions") -> str:
     Returns the schema (column names and data types) for the specified table in the MotherDuck database.
     Defaults to the 'positions' table.
     """
+    conn = None
     try:
         conn = get_md_connection()
         
         # Use PRAGMA table_info to get the schema details dynamically
+        # This is a standard DuckDB/SQLite way to get table schema
         query = f"PRAGMA table_info('{table_name}')"
         schema_df = conn.execute(query).fetchdf()
         
         if schema_df.empty:
-            return f"Error: Table '{table_name}' not found in the MotherDuck database."
+            return f"Error: Table '{table_name}' not found in the MotherDuck database. Available tables: {conn.execute('SHOW TABLES;').fetchnames()}"
 
         # Format the schema into a simple string: name TYPE, name TYPE, ...
         schema_parts = [f"{row['name']} {row['type']}" for index, row in schema_df.iterrows()]
@@ -85,6 +89,5 @@ def get_table_schema(table_name: str = "positions") -> str:
     except Exception as e:
         return f"DuckDB Schema Error: {e}"
     finally:
-        # Always close the connection
-        if 'conn' in locals() and conn:
-            conn.close()
+        if conn:
+            conn.close()