dao-ai 0.0.25__py3-none-any.whl → 0.1.2__py3-none-any.whl

dao_ai/tools/sql.py

@@ -0,0 +1,202 @@
+"""
+SQL execution tool for running SQL statements against Databricks SQL warehouses.
+
+This module provides a factory function for creating tools that execute
+pre-configured SQL statements against a Databricks SQL warehouse.
+"""
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.sql import StatementResponse, StatementState
+from langchain.tools import tool
+from loguru import logger
+
+from dao_ai.config import WarehouseModel, value_of
+
+
+def create_execute_statement_tool(
+    warehouse: WarehouseModel,
+    statement: str,
+    name: str = "execute_sql_tool",
+    description: str | None = None,
+) -> tool:
+    """
+    Create a tool for executing a pre-configured SQL statement against a Databricks SQL warehouse.
+
+    This factory function generates a tool that executes a specific SQL statement
+    (defined at configuration time) against a Databricks SQL warehouse. The SQL is
+    fixed and cannot be modified by the LLM at runtime, making this suitable for
+    providing agents with specific, pre-defined queries.
+
+    Args:
+        warehouse: WarehouseModel containing warehouse configuration and credentials
+        sql: The SQL statement to execute (configured at tool creation time)
+        name: Optional custom name for the tool. Defaults to "execute_sql_tool"
+        description: Optional custom description for the tool. If None, uses default description
+
+    Returns:
+        A LangChain tool that executes the pre-configured SQL statement and returns results
+
+    Example:
+        ```python
+        from dao_ai.config import WarehouseModel
+        from dao_ai.tools.sql import create_execute_sql_tool
+
+        # Create warehouse model
+        warehouse = WarehouseModel(
+            name="analytics_warehouse",
+            warehouse_id="abc123def456",
+        )
+
+        # Create SQL execution tool with pre-configured query
+        customer_count_tool = create_execute_sql_tool(
+            warehouse=warehouse,
+            sql="SELECT COUNT(*) as customer_count FROM catalog.schema.customers",
+            name="get_customer_count",
+            description="Get the total number of customers in the database"
+        )
+
+        # Use the tool (no parameters needed - SQL is pre-configured)
+        result = customer_count_tool.invoke({})
+        ```
+    """
+    if description is None:
+        description = f"Execute a pre-configured SQL query against the {warehouse.name} warehouse and return the results."
+
+    warehouse_id: str = value_of(warehouse.warehouse_id)
+    workspace_client: WorkspaceClient = warehouse.workspace_client
+
+    logger.debug(
+        "Creating SQL execution tool",
+        tool_name=name,
+        warehouse_name=warehouse.name,
+        warehouse_id=warehouse_id,
+        sql_preview=statement[:100] + "..." if len(statement) > 100 else statement,
+    )
+
+    @tool(name_or_callable=name, description=description)
+    def execute_statement_tool() -> str:
+        """
+        Execute the pre-configured SQL statement against the Databricks SQL warehouse.
+
+        Returns:
+            A string containing the query results or execution status
+        """
+        logger.info(
+            "Executing SQL statement",
+            tool_name=name,
+            warehouse_id=warehouse_id,
+            sql_preview=statement[:100] + "..." if len(statement) > 100 else statement,
+        )
+
+        try:
+            # Execute the SQL statement
+            statement_response: StatementResponse = (
+                workspace_client.statement_execution.execute_statement(
+                    warehouse_id=warehouse_id,
+                    statement=statement,
+                    wait_timeout="30s",
+                )
+            )
+
+            # Poll for completion if still pending
+            while statement_response.status.state in [
+                StatementState.PENDING,
+                StatementState.RUNNING,
+            ]:
+                logger.trace(
+                    "SQL statement still executing, polling...",
+                    statement_id=statement_response.statement_id,
+                    state=statement_response.status.state,
+                )
+                statement_response = workspace_client.statement_execution.get_statement(
+                    statement_response.statement_id
+                )
+
+            # Check execution status
+            if statement_response.status.state != StatementState.SUCCEEDED:
+                error_msg: str = (
+                    f"SQL execution failed with state {statement_response.status.state}"
+                )
+                if statement_response.status.error:
+                    error_msg += f": {statement_response.status.error.message}"
+                logger.error(
+                    "SQL execution failed",
+                    tool_name=name,
+                    statement_id=statement_response.statement_id,
+                    error=error_msg,
+                )
+                return f"Error: {error_msg}"
+
+            # Extract results
+            result = statement_response.result
+            if result is None:
+                logger.debug(
+                    "SQL statement executed successfully with no results",
+                    tool_name=name,
+                    statement_id=statement_response.statement_id,
+                )
+                return "SQL statement executed successfully (no results returned)"
+
+            # Format results
+            if result.data_array:
+                rows = result.data_array
+                row_count = len(rows)
+
+                # Get column names if available
+                columns = []
+                if (
+                    statement_response.manifest
+                    and statement_response.manifest.schema
+                    and statement_response.manifest.schema.columns
+                ):
+                    columns = [
+                        col.name for col in statement_response.manifest.schema.columns
+                    ]
+
+                logger.info(
+                    "SQL query returned results",
+                    tool_name=name,
+                    row_count=row_count,
+                    column_count=len(columns),
+                )
+
+                # Format as a simple text table
+                result_lines = []
+                if columns:
+                    result_lines.append(" | ".join(columns))
+                    result_lines.append("-" * (len(" | ".join(columns))))
+
+                for row in rows:
+                    result_lines.append(
+                        " | ".join(
+                            str(cell) if cell is not None else "NULL" for cell in row
+                        )
+                    )
+
+                # Add summary
+                result_lines.append("")
+                result_lines.append(
+                    f"({row_count} row{'s' if row_count != 1 else ''} returned)"
+                )
+
+                return "\n".join(result_lines)
+            else:
+                logger.debug(
+                    "SQL statement executed successfully with empty result set",
+                    tool_name=name,
+                    statement_id=statement_response.statement_id,
+                )
+                return "SQL statement executed successfully (empty result set)"
+
+        except Exception as e:
+            error_msg = f"Failed to execute SQL: {str(e)}"
+            logger.error(
+                "SQL execution failed with exception",
+                tool_name=name,
+                warehouse_id=warehouse_id,
+                error=str(e),
+                exc_info=True,
+            )
+            return f"Error: {error_msg}"
+
+    return execute_statement_tool

dao_ai/tools/time.py

@@ -20,7 +20,7 @@ def current_time_tool() -> str:
     """
 
     time_now: str = datetime.now().isoformat()
-    logger.debug(f"Current time: {time_now}")
+    logger.trace("Current time retrieved", time=time_now)
     return time_now
 
 
@@ -38,7 +38,9 @@ def time_in_timezone_tool(timezone_name: str) -> str:
     try:
         tz = pytz.timezone(timezone_name)
         time_in_tz = datetime.now(tz)
-        logger.debug(f"Time in {timezone_name}: {time_in_tz}")
+        logger.trace(
+            "Time in timezone retrieved", timezone=timezone_name, time=str(time_in_tz)
+        )
         return f"{time_in_tz.strftime('%Y-%m-%d %H:%M:%S %Z')} ({timezone_name})"
     except Exception:
         return f"Error: Invalid timezone '{timezone_name}'. Use format like 'US/Eastern' or 'Europe/London'"
@@ -77,7 +79,12 @@ def time_difference_tool(datetime1: str, datetime2: str) -> str:
             parts.append(f"{seconds} second{'s' if seconds != 1 else ''}")
 
         result = ", ".join(parts) if parts else "0 seconds"
-        logger.debug(f"Time difference between {datetime1} and {datetime2}: {result}")
+        logger.trace(
+            "Time difference calculated",
+            datetime1=datetime1,
+            datetime2=datetime2,
+            result=result,
+        )
         return result
 
     except Exception as e:
@@ -112,7 +119,14 @@ def add_time_tool(
         new_dt = base_dt + timedelta(days=days, hours=hours, minutes=minutes)
 
         result = new_dt.isoformat()
-        logger.debug(f"Added {days}d {hours}h {minutes}m to {base_datetime}: {result}")
+        logger.trace(
+            "Time added to datetime",
+            base_datetime=base_datetime,
+            days=days,
+            hours=hours,
+            minutes=minutes,
+            result=result,
+        )
         return result
 
     except Exception as e:
@@ -164,7 +178,7 @@ def is_business_hours_tool(
         elif not is_work_hours:
             result += " (Outside 9 AM - 5 PM)"
 
-        logger.debug(f"Business hours check: {result}")
+        logger.trace("Business hours check completed", result=result)
         return result
 
     except Exception as e:
@@ -199,7 +213,12 @@ def format_time_tool(datetime_str: str, format_type: str = "readable") -> str:
             return f"Error: Unknown format type. Use: {', '.join(formats.keys())}"
 
         result = dt.strftime(formats[format_type])
-        logger.debug(f"Formatted {datetime_str} as {format_type}: {result}")
+        logger.trace(
+            "Datetime formatted",
+            datetime_str=datetime_str,
+            format_type=format_type,
+            result=result,
+        )
         return result
 
     except Exception as e:
@@ -266,7 +285,11 @@ def time_until_tool(target_datetime: str) -> str:
                 else "Less than 1 minute remaining"
             )
 
-        logger.debug(f"Time until {target_datetime}: {result}")
+        logger.trace(
+            "Time until target calculated",
+            target_datetime=target_datetime,
+            result=result,
+        )
         return result
 
     except Exception as e: