Fix pyright type errors and AWS Bedrock semantic tag bug (#47)

## Summary

- Resolved all 74 pyright type errors across the codebase
- Fixed bug where AWS Bedrock LLM assigns semantic tags to numeric
columns (LONG, BIGINT, etc.), causing Stitch jobs to fail with
"Semantics on fields of type LONG are not supported"

## Changes

**Type fixes:**
- Aligned `ModelInfo` TypedDict across LLM providers
- Fixed `WizardState.models` type hint
- Added proper None handling and type narrowing throughout
- Fixed except clause ordering and Callable imports

**Stitch/PII fix:**
- Added `NUMERIC_TYPES` filter in `stitch_tools.py` to strip semantics
from numeric columns
- Updated PII detection prompt to instruct LLM not to tag numeric
columns

## Test plan

- [x] All 626 unit tests pass
- [x] Pyright reports 0 errors
- [x] Ruff linting passes

Author: pragyan-amp

chuck_data/agent/manager.py

@@ -210,29 +210,36 @@ def process_with_tools(self, tools, max_iterations: int = 20):
             if response_message.tool_calls:
                 # Add the assistant's response (requesting tool calls) to history
                 # Convert ChatCompletionMessage to dict format for consistency
+                tool_calls_list = []
+                for tc in response_message.tool_calls:
+                    func = getattr(tc, "function", None)
+                    if func is not None:
+                        tool_calls_list.append(
+                            {
+                                "id": tc.id,
+                                "type": getattr(tc, "type", "function"),
+                                "function": {
+                                    "name": getattr(func, "name", ""),
+                                    "arguments": getattr(func, "arguments", "{}"),
+                                },
+                            }
+                        )
                 assistant_msg = {
                     "role": "assistant",
                     "content": response_message.content,
-                    "tool_calls": [
-                        {
-                            "id": tc.id,
-                            "type": getattr(tc, "type", "function"),
-                            "function": {
-                                "name": tc.function.name,
-                                "arguments": tc.function.arguments,
-                            },
-                        }
-                        for tc in response_message.tool_calls
-                    ],
+                    "tool_calls": tool_calls_list,
                 }
                 self.conversation_history.append(assistant_msg)
 
                 # Execute each tool call
                 for tool_call in response_message.tool_calls:
-                    tool_name = tool_call.function.name
+                    func = getattr(tool_call, "function", None)
+                    if func is None:
+                        continue
+                    tool_name = getattr(func, "name", "")
                     tool_id = tool_call.id
                     try:
-                        tool_args = json.loads(tool_call.function.arguments)
+                        tool_args = json.loads(getattr(func, "arguments", "{}"))
                         tool_result = execute_tool(
                             self.api_client,
                             tool_name,
@@ -276,7 +283,7 @@ def process_with_tools(self, tools, max_iterations: int = 20):
                 continue
             else:
                 # No tool calls, this is the final response
-                final_content = response_message.content
+                final_content = response_message.content or ""
                 # remove all lines with any <function> tags
                 final_content = "\n".join(
                     line

chuck_data/agent/tool_executor.py

@@ -23,7 +23,8 @@
 from chuck_data.clients.databricks import (
     DatabricksAPIClient,
 )  # For type hinting api_client
-from typing import Dict, Any, Optional, List
+from typing import Dict, Any, Optional, List, Callable
+from jsonschema.exceptions import ValidationError
 
 
 # The display_to_user utility and individual tool implementation functions
@@ -48,7 +49,7 @@ def execute_tool(
     api_client: Optional[DatabricksAPIClient],
     tool_name: str,
     tool_args: Dict[str, Any],
-    output_callback: Optional[callable] = None,
+    output_callback: Optional[Callable[..., Any]] = None,
 ) -> Dict[str, Any]:
     """Execute a tool (command) by its name with the provided arguments.
 
@@ -87,7 +88,7 @@ def execute_tool(
     try:
         jsonschema.validate(instance=tool_args, schema=schema_to_validate)
         logging.debug(f"Tool arguments for '{tool_name}' validated successfully.")
-    except jsonschema.exceptions.ValidationError as ve:
+    except ValidationError as ve:
         logging.error(
             f"Validation error for tool '{tool_name}' args {tool_args}: {ve.message}"
         )

chuck_data/api_client.py

@@ -135,6 +135,8 @@ def upload_file(self, path, file_path=None, content=None, overwrite=False):
                 binary_data = f.read()
         else:
             # Convert string content to bytes
+            # content is guaranteed non-None by the validation above
+            assert content is not None
             binary_data = content.encode("utf-8")
 
         try:

chuck_data/clients/amperity.py

@@ -8,6 +8,7 @@
 import webbrowser
 import readchar
 import json
+from typing import Optional
 from rich.console import Console
 
 from chuck_data.config import set_amperity_token
@@ -113,7 +114,7 @@ def get_auth_status(self) -> dict:
         return {"state": self.state, "nonce": self.nonce, "has_token": bool(self.token)}
 
     def wait_for_auth_completion(
-        self, poll_interval: int = 1, timeout: int = None
+        self, poll_interval: int = 1, timeout: Optional[int] = None
     ) -> tuple[bool, str]:
         """Wait for authentication to complete in a blocking manner."""
         if not self.nonce:

chuck_data/clients/databricks.py

@@ -724,6 +724,8 @@ def upload_file(self, path, file_path=None, content=None, overwrite=False):
                 binary_data = f.read()
         else:
             # Convert string content to bytes
+            # content is guaranteed non-None by the validation above
+            assert content is not None
             binary_data = content.encode("utf-8")
 
         try:

chuck_data/command_output.py

@@ -10,7 +10,7 @@
 
 from chuck_data.ui.table_formatter import display_table
 
-from chuck_data.command_result import CommandResult
+from chuck_data.commands.base import CommandResult
 from chuck_data.ui.theme import (
     SUCCESS,
     WARNING,
@@ -132,7 +132,7 @@ def format_for_agent(result: CommandResult) -> Dict[str, Any]:
             }
 
         # Start with a base response
-        response = {"success": True}
+        response: Dict[str, Any] = {"success": True}
 
         # Add the message if available
         if result.message:

chuck_data/commands/catalog_selection.py

@@ -64,7 +64,7 @@ def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandRe
         client: API client instance
         **kwargs: catalog (str) - catalog name, tool_output_callback (optional)
     """
-    catalog: str = kwargs.get("catalog")
+    catalog = kwargs.get("catalog")
     tool_output_callback = kwargs.get("tool_output_callback")
 
     if not catalog:

chuck_data/commands/job_status.py

@@ -639,7 +639,7 @@ def handle_list_jobs(client=None, **kwargs) -> CommandResult:
         cached_job_data = job_entry.get("job_data")
 
         # If we have cached data for a terminal state, use it
-        if cached_job_data:
+        if cached_job_data and isinstance(cached_job_data, dict):
             state = (cached_job_data.get("state") or "").lower().replace(":", "")
             # Only use cache for terminal states (succeeded, failed, unknown)
             if state in ["succeeded", "success", "failed", "error", "unknown"]:

chuck_data/commands/jobs.py

@@ -13,8 +13,8 @@ def handle_launch_job(client: Optional[DatabricksAPIClient], **kwargs) -> Comman
         client: API client instance
         **kwargs: config_path (str), init_script_path (str), run_name (str, optional), tool_output_callback (callable, optional)
     """
-    config_path: str = kwargs.get("config_path")
-    init_script_path: str = kwargs.get("init_script_path")
+    config_path = kwargs.get("config_path")
+    init_script_path = kwargs.get("init_script_path")
     run_name: Optional[str] = kwargs.get("run_name")
     tool_output_callback = kwargs.get("tool_output_callback")
     policy_id: Optional[str] = kwargs.get("policy_id")

chuck_data/commands/model_selection.py

@@ -24,7 +24,7 @@ def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandRe
         client: API client instance (used for Databricks provider)
         **kwargs: model_name (str)
     """
-    model_name: str = kwargs.get("model_name")
+    model_name = kwargs.get("model_name")
     if not model_name:
         return CommandResult(False, message="model_name parameter is required.")
 
@@ -45,7 +45,7 @@ def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandRe
         models_list = provider.list_models(tool_calling_only=False)
 
         # Extract model IDs (field name varies by provider)
-        model_ids = [m.get("model_id") or m.get("name") for m in models_list]
+        model_ids = [m.get("model_id") or m.get("name") or "" for m in models_list]
 
         # Validate model exists
         if model_name not in model_ids: