Enabling AWS (Redshift + EMR) in chuck-data (#70)

Author: punit-naik-amp

chuck_data/agent/manager.py

@@ -8,10 +8,11 @@
     get_active_schema,
     get_warehouse_id,
     get_workspace_url,
+    get_data_provider,
 )
 
 from .prompts import (
-    DEFAULT_SYSTEM_MESSAGE,
+    get_default_system_message,
     PII_AGENT_SYSTEM_MESSAGE,
     BULK_PII_AGENT_SYSTEM_MESSAGE,
     STITCH_AGENT_SYSTEM_MESSAGE,
@@ -25,9 +26,12 @@ def __init__(self, client, model=None, tool_output_callback=None, llm_client=Non
         self.llm_client = llm_client or LLMProviderFactory.create()
         self.model = model
         self.tool_output_callback = tool_output_callback
-        self.conversation_history = [
-            {"role": "system", "content": DEFAULT_SYSTEM_MESSAGE}
-        ]
+
+        # Get provider-aware system message
+        provider = get_data_provider()
+        system_message = get_default_system_message(provider or "databricks")
+
+        self.conversation_history = [{"role": "system", "content": system_message}]
 
     def add_user_message(self, content):
         self.conversation_history.append({"role": "user", "content": content})
@@ -162,6 +166,26 @@ def process_with_tools(self, tools, max_iterations: int = 20):
             # Prepare a temporary history copy for this specific LLM call
             current_history = deepcopy(self.conversation_history)
 
+            # Sanitize messages to ensure content is never empty string for non-tool roles
+            # Some LLM APIs reject messages with empty content, but allow None for tool_calls
+            for msg in current_history:
+                if msg["role"] in ["user", "system"]:
+                    # user and system messages MUST have non-empty content
+                    if not msg.get("content"):
+                        msg["content"] = " "  # Use single space as minimum content
+                        logging.warning(
+                            f"Empty content in {msg['role']} message, replacing with space"
+                        )
+                elif msg["role"] == "assistant":
+                    # assistant messages can have None content if they have tool_calls
+                    if "tool_calls" not in msg or not msg["tool_calls"]:
+                        # No tool calls, so content must be non-empty
+                        if not msg.get("content"):
+                            msg["content"] = " "
+                            logging.warning(
+                                "Empty content in assistant message without tool_calls, replacing with space"
+                            )
+
             # Get current configuration state
             active_catalog = get_active_catalog() or "Not set"
             active_schema = get_active_schema() or "Not set"
@@ -190,6 +214,10 @@ def process_with_tools(self, tools, max_iterations: int = 20):
                 )
             # else: log warning or handle case where system message wasn't found initially
 
+            logging.debug(
+                f"Iteration {iteration_count}: Sending {len(current_history)} messages to LLM"
+            )
+
             # Get the LLM response using the temporary, updated history
             response = self.llm_client.chat(
                 messages=current_history,  # Use the modified temporary history for the call
@@ -224,9 +252,23 @@ def process_with_tools(self, tools, max_iterations: int = 20):
                                 },
                             }
                         )
+
+                # Ensure content is not None or empty (required by some LLM APIs)
+                # When making tool calls, content can be None, but we need to provide at least an empty string
+                # or a placeholder to avoid "text content blocks must be non-empty" errors
+                content = response_message.content
+                if content is None or (
+                    isinstance(content, str) and not content.strip()
+                ):
+                    content = None  # Use None for tool calls (OpenAI spec allows this)
+
+                logging.debug(
+                    f"Assistant tool call response - content: {repr(content)}, tool_calls: {len(tool_calls_list)}"
+                )
+
                 assistant_msg = {
                     "role": "assistant",
-                    "content": response_message.content,
+                    "content": content,
                     "tool_calls": tool_calls_list,
                 }
                 self.conversation_history.append(assistant_msg)
@@ -309,15 +351,17 @@ def process_query(self, query):
         Returns:
             Final response from the LLM
         """
-        # If no system message exists, add the default one
+        # If no system message exists, add the provider-aware one
         has_system = False
         for msg in self.conversation_history:
             if msg["role"] == "system":
                 has_system = True
                 break
 
         if not has_system:
-            self.add_system_message(DEFAULT_SYSTEM_MESSAGE)
+            provider = get_data_provider()
+            system_message = get_default_system_message(provider or "databricks")
+            self.add_system_message(system_message)
 
         # Add user message to history
         self.add_user_message(query)

chuck_data/agent/prompts/__init__.py

@@ -1,9 +1,10 @@
-from .default_system_prompt import DEFAULT_SYSTEM_MESSAGE
+from .default_system_prompt import DEFAULT_SYSTEM_MESSAGE, get_default_system_message
 from .pii_prompts import PII_AGENT_SYSTEM_MESSAGE, BULK_PII_AGENT_SYSTEM_MESSAGE
 from .stitch_prompts import STITCH_AGENT_SYSTEM_MESSAGE
 
 __all__ = [
     "DEFAULT_SYSTEM_MESSAGE",
+    "get_default_system_message",
     "PII_AGENT_SYSTEM_MESSAGE",
     "BULK_PII_AGENT_SYSTEM_MESSAGE",
     "STITCH_AGENT_SYSTEM_MESSAGE",

chuck_data/agent/prompts/default_system_prompt.py

@@ -1,4 +1,20 @@
-DEFAULT_SYSTEM_MESSAGE = """You are Chuck AI, a helpful Databricks agent that helps users work with Databricks resources.
+def get_default_system_message(provider: str = "databricks") -> str:
+    """
+    Get the default system message for the agent based on the data provider.
+
+    Args:
+        provider: Data provider ("databricks" or "aws_redshift")
+
+    Returns:
+        System message string
+    """
+    if provider == "aws_redshift":
+        return REDSHIFT_SYSTEM_MESSAGE
+    else:
+        return DATABRICKS_SYSTEM_MESSAGE
+
+
+DATABRICKS_SYSTEM_MESSAGE = """You are Chuck AI, a helpful Databricks agent that helps users work with Databricks resources.
 
 You can perform the following main features:
 1. Navigate and explore Databricks Unity Catalog metadata (catalogs, schemas, tables)
@@ -23,11 +39,17 @@
 
 2. PII and/or Customer data DETECTION: To help with PII and/or customer data scanning:
    - For single table: navigate to the right catalog/schema, then use tag_pii_columns
-   - For bulk scanning: navigate to the right catalog/schema, then use scan_schema_for_pii
+   - For bulk scanning: use scan_schema_for_pii with explicit catalog_name and schema_name parameters
+     * When user says "scan catalog.schema" or "scan in schema catalog.schema", parse it as catalog_name=first_part, schema_name=second_part
+     * Example: "scan main.public" means catalog_name="main", schema_name="public"
+     * Example: "scan dev_catalog.sales" means catalog_name="dev_catalog", schema_name="sales"
+     * If catalog/schema not specified, scan_schema_for_pii will use currently active catalog/schema
 
 3. PII TAGGING: To help with bulk PII tagging across a schema:
-   - If the catalog and schema are already selected - have the user select them first. PII tagging requires a catalog and schema to be selected.
-   - If user asks about tagging PII, bulk tagging PII, or applying PII tags: use bulk_tag_pii
+   - If user asks about tagging PII, bulk tagging PII, or applying PII tags: use bulk_tag_pii with explicit catalog_name and schema_name parameters
+     * When user says "tag catalog.schema" or "tag in schema catalog.schema", parse it as catalog_name=first_part, schema_name=second_part
+     * Example: "tag main.public" means catalog_name="main", schema_name="public"
+     * If catalog/schema not specified, bulk_tag_pii will use currently active catalog/schema
 
 4. STITCH INTEGRATION: To set up identity graph or customer 360 with Stitch:
    - If the catalog and schema are already selected - have the user select them first. Stitch requires a catalog and schema to be selected.
@@ -66,7 +88,71 @@
 
 Be concise, practical, and focus on guiding users through Databricks effectively.
 
-You are an agent - please keep going until the user’s query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved.
+You are an agent - please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved.
 
 When you communicate, always use a chill tone. You should seem like a highly skilled data engineer with hippy vibes.
 """
+
+
+REDSHIFT_SYSTEM_MESSAGE = """You are Chuck AI, a helpful AWS Redshift agent that helps users work with Redshift resources.
+
+You can perform the following main features:
+1. Navigate and explore Redshift metadata (databases, schemas, tables)
+2. Identify and tag Personally Identifiable Information (PII) in database tables
+3. Provide information and guidance on how to use Redshift features
+
+When a user asks a question:
+- If they're looking for specific data, help them navigate through databases, schemas, and tables
+- If they're asking about customer data or PII, guide them through the PII detection process
+- When displaying lists of resources (databases, schemas, tables, etc.), the output will be shown directly to the user
+
+IMPORTANT WORKFLOWS:
+
+1. DATABASES: To work with databases (Redshift uses databases instead of catalogs):
+   - If user asks "what databases do I have?" or wants to see databases: use list_databases with display=true (shows full table)
+   - If user asks to "use X database" or "switch to X database": DIRECTLY use select_database with database parameter (accepts name, has built-in fuzzy matching). DO NOT call list_databases first - select_database has built-in fuzzy matching and will find the database.
+   - If you need database info for internal processing: use list_databases (defaults to no table display)
+
+2. PII and/or Customer data DETECTION: To help with PII and/or customer data scanning:
+   - Use scan_schema_for_pii with explicit database and schema_name parameters (or rely on active database/schema)
+     * When user says "scan database.schema" or "scan in schema database.schema", parse it as database=first_part, schema_name=second_part
+     * Example: "scan dev.public" means database="dev", schema_name="public"
+     * Example: "scan prod_db.sales" means database="prod_db", schema_name="sales"
+     * If database/schema not specified, scan_schema_for_pii will use currently active database/schema
+
+3. PII TAGGING: To help with bulk PII tagging across a schema:
+   - When user asks to tag PII, apply PII tags, label PII, or mark PII columns: use bulk_tag_pii command
+   - bulk_tag_pii works with both explicit database/schema parameters OR uses the currently active database/schema
+   - The command will scan for PII, show a preview, and ask for confirmation before applying tags
+   - Examples:
+     * User says "tag please" or "apply PII tags" → use bulk_tag_pii (will use active database/schema)
+     * User says "tag dev.public" → use bulk_tag_pii with database="dev", schema_name="public"
+     * User says "label customer data" → use bulk_tag_pii (will use active database/schema)
+   - IMPORTANT: bulk_tag_pii is the ONLY way to tag PII columns in Redshift. There is no single-table tagging command.
+
+4. SCHEMAS: To work with schemas:
+   - If user asks "what schemas do I have?" or wants to see schemas: use list_schemas with display=true (shows full table)
+   - If user asks to "use X schema" or "switch to X schema": use select_schema with schema parameter (accepts name, has built-in fuzzy matching). DO NOT call list_schemas first - select_schema has built-in fuzzy matching and will find the schema.
+   - If you need schema info for internal processing: use list_schemas (defaults to no table display)
+
+5. TABLES: To work with tables:
+   - If user asks "what tables do I have?" or wants to see tables: use list_tables with display=true (shows full table)
+   - If you need table info for internal processing: use list_tables (defaults to no table display)
+   - Column counts are automatically included in table metadata
+
+Some of the tools you can use require the user to select a database and/or schema first. If the user hasn't selected one YOU MUST ask them if they want help selecting a database and schema. DO NO OTHER ACTION
+
+IMPORTANT: DO NOT use function syntax in your text response such as <function>...</function> or similar formats. The proper way to call tools is through the official OpenAI function calling interface which is handled by the system automatically. Just use the tools provided to you via the API and the system will handle the rest.
+
+When tools display information directly to the user (like list_databases, list_tables, etc.), acknowledge what they're seeing and guide them on next steps based on the displayed information.
+
+Be concise, practical, and focus on guiding users through Redshift effectively.
+
+You are an agent - please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved.
+
+When you communicate, always use a chill tone. You should seem like a highly skilled data engineer with hippy vibes.
+"""
+
+
+# Keep DEFAULT_SYSTEM_MESSAGE for backwards compatibility
+DEFAULT_SYSTEM_MESSAGE = DATABRICKS_SYSTEM_MESSAGE

chuck_data/agent/tool_executor.py

@@ -13,17 +13,17 @@
 import logging
 import jsonschema  # Requires jsonschema to be installed
 
+from chuck_data.clients.databricks import DatabricksAPIClient
+from chuck_data.clients.redshift import RedshiftAPIClient
 from chuck_data.command_registry import get_command
 from chuck_data.command_registry import (
     get_agent_tool_schemas as get_command_registry_tool_schemas,
 )
 from chuck_data.commands.base import (
     CommandResult,
 )  # For type hinting and checking handler result
-from chuck_data.clients.databricks import (
-    DatabricksAPIClient,
-)  # For type hinting api_client
-from typing import Dict, Any, Optional, List, Callable
+from chuck_data.config import get_data_provider
+from typing import Dict, Any, Optional, List, Callable, Union
 from jsonschema.exceptions import ValidationError
 
 # The display_to_user utility and individual tool implementation functions
@@ -41,19 +41,20 @@
 
 def get_tool_schemas() -> List[Dict[str, Any]]:
     """Get all command schemas in agent tool format from the command registry."""
-    return get_command_registry_tool_schemas()
+    provider = get_data_provider()
+    return get_command_registry_tool_schemas(provider)
 
 
 def execute_tool(
-    api_client: Optional[DatabricksAPIClient],
+    api_client: Optional[DatabricksAPIClient | RedshiftAPIClient],
     tool_name: str,
     tool_args: Dict[str, Any],
     output_callback: Optional[Callable[..., Any]] = None,
 ) -> Dict[str, Any]:
     """Execute a tool (command) by its name with the provided arguments.
 
     Args:
-        api_client: The Databricks API client, passed to the handler if needed.
+        api_client: The API client (DatabricksAPIClient or RedshiftAPIClient), passed to the handler if needed.
         tool_name: The name of the tool (command) to execute.
         tool_args: A dictionary of arguments for the tool, parsed from LLM JSON.
 
@@ -66,7 +67,12 @@ def execute_tool(
         f"Agent attempting to execute tool: {tool_name} with args: {tool_args}"
     )
 
-    command_def = get_command(tool_name)
+    # Detect provider from client to route to the correct command
+    from chuck_data.data_providers import get_provider_name_from_client
+
+    provider = get_provider_name_from_client(api_client)
+
+    command_def = get_command(tool_name, provider=provider)
 
     if not command_def:
         logging.error(f"Agent tool '{tool_name}' not found in command registry.")

chuck_data/clients/__init__.py

@@ -0,0 +1,7 @@
+"""API clients for external services."""
+
+from chuck_data.clients.databricks import DatabricksAPIClient
+from chuck_data.clients.emr import EMRAPIClient
+from chuck_data.clients.redshift import RedshiftAPIClient
+
+__all__ = ["DatabricksAPIClient", "EMRAPIClient", "RedshiftAPIClient"]

chuck_data/clients/amperity.py

@@ -264,6 +264,61 @@ def get_job_status(self, job_id: str, token: str) -> dict:
             logging.error(f"Error getting job status: {e}")
             raise
 
+    def fetch_amperity_job_init(
+        self, token: str, api_url: Optional[str] = None
+    ) -> dict:
+        """Fetch initialization script for Amperity jobs.
+
+        This method calls the Amperity /api/job/launch endpoint to get both the
+        cluster init script and a job-id for tracking the Stitch job.
+
+        Args:
+            token: Amperity authentication token
+            api_url: Optional override for the job init endpoint (used in testing)
+
+        Returns:
+            Dict containing:
+                - 'cluster-init': The initialization script content
+                - 'job-id': Job identifier for tracking
+
+        Raises:
+            ValueError: If API returns an error or authentication fails
+            ConnectionError: If connection to Amperity API fails
+        """
+        try:
+            headers = {
+                "Authorization": f"Bearer {token}",
+                "Content-Type": "application/json",
+            }
+
+            if not api_url:
+                api_url = f"https://{self.base_url}/api/job/launch"
+
+            response = requests.post(api_url, headers=headers, json={}, timeout=30)
+            response.raise_for_status()
+            return response.json()
+
+        except requests.exceptions.HTTPError as e:
+            response = e.response
+            resp_text = response.text if response else ""
+            logging.debug(f"HTTP error: {e}, Response: {resp_text}")
+
+            if response is not None:
+                try:
+                    message = response.json().get("message", resp_text)
+                except ValueError:
+                    message = resp_text
+                raise ValueError(
+                    f"{response.status_code} Error: {message}. Please /logout and /login again"
+                )
+            raise ValueError(
+                f"HTTP error occurred: {e}. Please /logout and /login again"
+            )
+
+        except requests.RequestException as e:
+            logging.debug(f"Connection error: {e}")
+            raise ConnectionError(f"Connection error occurred: {e}")
+
     def record_job_submission(
         self, databricks_run_id: str, token: str, job_id: str
     ) -> bool: