dao-ai 0.0.28__py3-none-any.whl → 0.1.5__py3-none-any.whl

dao_ai/__init__.py

@@ -0,0 +1,29 @@
+"""
+DAO AI - A framework for building AI agents with Databricks.
+
+This module configures package-level settings including warning filters
+for expected runtime warnings that don't indicate actual problems.
+"""
+
+import warnings
+
+# Suppress Pydantic serialization warnings for Context objects during checkpointing.
+# This warning occurs because LangGraph's checkpointer serializes the context_schema
+# and Pydantic reports that serialization may not be as expected. This is benign
+# since Context is only used at runtime and doesn't need to be persisted.
+#
+# The warning looks like:
+# PydanticSerializationUnexpectedValue(Expected `none` - serialized value may not
+# be as expected [field_name='context', input_value=Context(...), input_type=Context])
+warnings.filterwarnings(
+    "ignore",
+    message=r".*Pydantic serializer warnings.*",
+    category=UserWarning,
+)
+
+# Also filter the specific PydanticSerializationUnexpectedValue warning
+warnings.filterwarnings(
+    "ignore",
+    message=r".*PydanticSerializationUnexpectedValue.*",
+    category=UserWarning,
+)

dao_ai/agent_as_code.py

@@ -1,11 +1,9 @@
-import sys
-
 import mlflow
-from loguru import logger
 from mlflow.models import ModelConfig
 from mlflow.pyfunc import ResponsesAgent
 
 from dao_ai.config import AppConfig
+from dao_ai.logging import configure_logging
 
 mlflow.set_registry_uri("databricks-uc")
 mlflow.set_tracking_uri("databricks")
@@ -17,8 +15,7 @@ config: AppConfig = AppConfig(**model_config.to_dict())
 
 log_level: str = config.app.log_level
 
-logger.remove()
-logger.add(sys.stderr, level=log_level)
+configure_logging(level=log_level)
 
 app: ResponsesAgent = config.as_responses_agent()
 

dao_ai/cli.py

@@ -1,4 +1,5 @@
 import argparse
+import getpass
 import json
 import os
 import subprocess
@@ -13,11 +14,88 @@ from loguru import logger
 
 from dao_ai.config import AppConfig
 from dao_ai.graph import create_dao_ai_graph
+from dao_ai.logging import configure_logging
 from dao_ai.models import save_image
 from dao_ai.utils import normalize_name
 
-logger.remove()
-logger.add(sys.stderr, level="ERROR")
+configure_logging(level="ERROR")
+
+
+def get_default_user_id() -> str:
+    """
+    Get the default user ID for the CLI session.
+
+    Tries to get the current user from Databricks, falls back to local user.
+
+    Returns:
+        User ID string (Databricks username or local username)
+    """
+    try:
+        # Try to get current user from Databricks SDK
+        from databricks.sdk import WorkspaceClient
+
+        w = WorkspaceClient()
+        current_user = w.current_user.me()
+        user_id = current_user.user_name
+        logger.debug(f"Using Databricks user: {user_id}")
+        return user_id
+    except Exception as e:
+        # Fall back to local system user
+        logger.debug(f"Could not get Databricks user, using local user: {e}")
+        local_user = getpass.getuser()
+        logger.debug(f"Using local user: {local_user}")
+        return local_user
+
+
+def detect_cloud_provider(profile: Optional[str] = None) -> Optional[str]:
+    """
+    Detect the cloud provider from the Databricks workspace URL.
+
+    The cloud provider is determined by the workspace URL pattern:
+    - Azure: *.azuredatabricks.net
+    - AWS: *.cloud.databricks.com (without gcp subdomain)
+    - GCP: *.gcp.databricks.com
+
+    Args:
+        profile: Optional Databricks CLI profile name
+
+    Returns:
+        Cloud provider string ('azure', 'aws', 'gcp') or None if detection fails
+    """
+    try:
+        from databricks.sdk import WorkspaceClient
+
+        # Create workspace client with optional profile
+        if profile:
+            w = WorkspaceClient(profile=profile)
+        else:
+            w = WorkspaceClient()
+
+        # Get the workspace URL from config
+        host = w.config.host
+        if not host:
+            logger.warning("Could not determine workspace URL for cloud detection")
+            return None
+
+        host_lower = host.lower()
+
+        if "azuredatabricks.net" in host_lower:
+            logger.debug(f"Detected Azure cloud from workspace URL: {host}")
+            return "azure"
+        elif ".gcp.databricks.com" in host_lower:
+            logger.debug(f"Detected GCP cloud from workspace URL: {host}")
+            return "gcp"
+        elif ".cloud.databricks.com" in host_lower or "databricks.com" in host_lower:
+            # AWS uses *.cloud.databricks.com or regional patterns
+            logger.debug(f"Detected AWS cloud from workspace URL: {host}")
+            return "aws"
+        else:
+            logger.warning(f"Could not determine cloud provider from URL: {host}")
+            return None
+
+    except Exception as e:
+        logger.warning(f"Could not detect cloud provider: {e}")
+        return None
 
 
 env_path: str = find_dotenv()
@@ -193,6 +271,13 @@ Examples:
         "-t",
         "--target",
         type=str,
+        help="Bundle target name (default: auto-generated from app name and cloud)",
+    )
+    bundle_parser.add_argument(
+        "--cloud",
+        type=str,
+        choices=["azure", "aws", "gcp"],
+        help="Cloud provider (auto-detected from workspace URL if not specified)",
     )
     bundle_parser.add_argument(
         "--dry-run",
@@ -240,9 +325,9 @@ Use Ctrl-C to interrupt and exit immediately.
         """,
         epilog="""
 Examples:
-  dao-ai chat -c config/model_config.yaml                              # Start chat with default settings
+  dao-ai chat -c config/model_config.yaml                              # Start chat (auto-detects user)
   dao-ai chat -c config/retail.yaml --custom-input store_num=87887     # Chat with custom store number
-  dao-ai chat -c config/prod.yaml --user-id john123                    # Chat with specific user ID
+  dao-ai chat -c config/prod.yaml --user-id john.doe@company.com       # Chat with specific user ID
   dao-ai chat -c config/retail.yaml --custom-input store_num=123 --custom-input region=west  # Multiple custom inputs
         """,
         formatter_class=argparse.RawDescriptionHelpFormatter,
@@ -264,28 +349,38 @@ Examples:
     chat_parser.add_argument(
         "--user-id",
         type=str,
-        default="my_user_id",
+        default=None,  # Will be set to actual user in handle_chat_command
         metavar="ID",
-        help="User ID for the chat session (default: my_user_id)",
+        help="User ID for the chat session (default: current Databricks user or local username)",
     )
     chat_parser.add_argument(
         "--thread-id",
         type=str,
-        default="1",
+        default=None,
         metavar="ID",
-        help="Thread ID for the chat session (default: 1)",
+        help="Thread ID for the chat session (default: auto-generated UUID)",
     )
 
     options = parser.parse_args(args)
 
+    # Generate a new thread_id UUID if not provided (only for chat command)
+    if hasattr(options, "thread_id") and options.thread_id is None:
+        import uuid
+
+        options.thread_id = str(uuid.uuid4())
+
     return options
 
 
 def handle_chat_command(options: Namespace) -> None:
-    """Interactive chat REPL with the DAO AI system."""
+    """Interactive chat REPL with the DAO AI system with Human-in-the-Loop support."""
     logger.debug("Starting chat session with DAO AI system...")
 
     try:
+        # Set default user_id if not provided
+        if options.user_id is None:
+            options.user_id = get_default_user_id()
+
         config: AppConfig = AppConfig.from_file(options.config)
         app = create_dao_ai_graph(config)
 
@@ -305,9 +400,7 @@ def handle_chat_command(options: Namespace) -> None:
         print("-" * 50)
 
         # Import streaming function and interrupt handling
-        from langchain_core.messages import HumanMessage
-
-        from dao_ai.models import process_messages_stream
+        from langchain_core.messages import AIMessage, HumanMessage
 
         # Conversation history
         messages = []
@@ -350,47 +443,210 @@ def handle_chat_command(options: Namespace) -> None:
                             )
                             continue
 
-                # Prepare custom inputs for the agent
-                custom_inputs = {"configurable": configurable}
+                # Create Context object from configurable dict
+                from dao_ai.state import Context
+
+                context = Context(**configurable)
+
+                # Prepare config with thread_id for checkpointer
+                # Note: thread_id is needed in config for checkpointer/memory
+                config = {"configurable": {"thread_id": options.thread_id}}
+
+                # Invoke the graph and handle interrupts (HITL)
+                # Wrap in async function to maintain connection pool throughout
+                logger.debug(f"Invoking graph with {len(messages)} messages")
+                logger.debug(f"Context: {context}")
+                logger.debug(f"Config: {config}")
+
+                import asyncio
+
+                from langgraph.types import Command
+
+                async def _invoke_with_hitl():
+                    """Invoke graph and handle HITL interrupts in single async context."""
+                    result = await app.ainvoke(
+                        {"messages": messages},
+                        config=config,
+                        context=context,  # Pass context as separate parameter
+                    )
+
+                    # Check for interrupts (Human-in-the-Loop) using __interrupt__
+                    # This is the modern LangChain pattern
+                    while "__interrupt__" in result:
+                        interrupts = result["__interrupt__"]
+                        logger.info(f"HITL: {len(interrupts)} interrupt(s) detected")
+
+                        # Collect decisions for all interrupts
+                        decisions = []
+
+                        for interrupt in interrupts:
+                            interrupt_value = interrupt.value
+                            action_requests = interrupt_value.get("action_requests", [])
+
+                            for action_request in action_requests:
+                                # Display interrupt information
+                                print("\n⚠️  Human in the Loop - Tool Approval Required")
+                                print(f"{'=' * 60}")
+
+                                tool_name = action_request.get("name", "unknown")
+                                tool_args = action_request.get("arguments", {})
+                                description = action_request.get("description", "")
+
+                                print(f"Tool: {tool_name}")
+                                if description:
+                                    print(f"\n{description}\n")
+
+                                print("Arguments:")
+                                for arg_name, arg_value in tool_args.items():
+                                    # Truncate long values
+                                    arg_str = str(arg_value)
+                                    if len(arg_str) > 100:
+                                        arg_str = arg_str[:97] + "..."
+                                    print(f"  - {arg_name}: {arg_str}")
+
+                                print(f"{'=' * 60}")
+
+                                # Prompt user for decision
+                                while True:
+                                    decision_input = (
+                                        input(
+                                            "\nAction? (a)pprove / (r)eject / (e)dit / (h)elp: "
+                                        )
+                                        .strip()
+                                        .lower()
+                                    )
+
+                                    if decision_input in ["a", "approve"]:
+                                        logger.info("User approved tool call")
+                                        print("✅ Approved - continuing execution...")
+                                        decisions.append({"type": "approve"})
+                                        break
+                                    elif decision_input in ["r", "reject"]:
+                                        logger.info("User rejected tool call")
+                                        feedback = input(
+                                            "   Feedback for agent (optional): "
+                                        ).strip()
+                                        if feedback:
+                                            decisions.append(
+                                                {"type": "reject", "message": feedback}
+                                            )
+                                        else:
+                                            decisions.append(
+                                                {
+                                                    "type": "reject",
+                                                    "message": "Tool call rejected by user",
+                                                }
+                                            )
+                                        print(
+                                            "❌ Rejected - agent will receive feedback..."
+                                        )
+                                        break
+                                    elif decision_input in ["e", "edit"]:
+                                        print(
+                                            "ℹ️  Edit functionality not yet implemented in CLI"
+                                        )
+                                        print("   Please approve or reject.")
+                                        continue
+                                    elif decision_input in ["h", "help"]:
+                                        print("\nAvailable actions:")
+                                        print(
+                                            "  (a)pprove - Execute the tool call as shown"
+                                        )
+                                        print(
+                                            "  (r)eject  - Cancel the tool call with optional feedback"
+                                        )
+                                        print(
+                                            "  (e)dit    - Modify arguments (not yet implemented)"
+                                        )
+                                        print("  (h)elp    - Show this help message")
+                                        continue
+                                    else:
+                                        print("Invalid option. Type 'h' for help.")
+                                        continue
+
+                        # Resume execution with decisions using Command
+                        # This is the modern LangChain pattern
+                        logger.debug(f"Resuming with {len(decisions)} decision(s)")
+                        result = await app.ainvoke(
+                            Command(resume={"decisions": decisions}),
+                            config=config,
+                            context=context,
+                        )
+
+                    return result
+
+                try:
+                    # Use async invoke - keep connection pool alive throughout HITL
+                    loop = asyncio.get_event_loop()
+                except RuntimeError:
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+
+                try:
+                    result = loop.run_until_complete(_invoke_with_hitl())
+                except Exception as e:
+                    logger.error(f"Error invoking graph: {e}")
+                    print(f"\n❌ Error: {e}")
+                    continue
 
+                # After all interrupts handled, display the final response
                 print("\n🤖 Assistant: ", end="", flush=True)
 
-                # Stream the response
                 response_content = ""
+                structured_response = None
                 try:
-                    for chunk in process_messages_stream(app, messages, custom_inputs):
-                        # Handle different chunk types
-                        if hasattr(chunk, "content") and chunk.content:
-                            content = chunk.content
-                            print(content, end="", flush=True)
-                            response_content += content
-                        elif hasattr(chunk, "choices") and chunk.choices:
-                            # Handle ChatCompletionChunk format
-                            for choice in chunk.choices:
-                                if (
-                                    hasattr(choice, "delta")
-                                    and choice.delta
-                                    and choice.delta.content
-                                ):
-                                    content = choice.delta.content
-                                    print(content, end="", flush=True)
-                                    response_content += content
-
-                    print()  # New line after streaming
+                    # Debug: Log what's in the result
+                    logger.debug(f"Result keys: {result.keys() if result else 'None'}")
+                    if result:
+                        for key in result.keys():
+                            logger.debug(f"Result['{key}'] type: {type(result[key])}")
+
+                    # Get the latest messages from the result
+                    if result and "messages" in result:
+                        latest_messages = result["messages"]
+                        # Find the last AI message
+                        for msg in reversed(latest_messages):
+                            if isinstance(msg, AIMessage):
+                                if hasattr(msg, "content") and msg.content:
+                                    response_content = msg.content
+                                    print(response_content, end="", flush=True)
+                                    break
+
+                    # Check for structured output and display it separately
+                    if result and "structured_response" in result:
+                        structured_response = result["structured_response"]
+                        import json
+
+                        structured_json = json.dumps(
+                            structured_response.model_dump()
+                            if hasattr(structured_response, "model_dump")
+                            else structured_response,
+                            indent=2,
+                        )
+
+                        # If there was message content, add separator
+                        if response_content.strip():
+                            print("\n\n📊 Structured Output:")
+                            print(structured_json)
+                        else:
+                            # No message content, just show structured output
+                            print(structured_json, end="", flush=True)
+
+                        response_content = response_content or structured_json
+
+                    print()  # New line after response
 
                     # Add assistant response to history if we got content
                     if response_content.strip():
-                        from langchain_core.messages import AIMessage
-
                         assistant_message = AIMessage(content=response_content)
                         messages.append(assistant_message)
                     else:
                         print("(No response content generated)")
 
                 except Exception as e:
-                    print(f"\n❌ Error during streaming: {e}")
+                    print(f"\n❌ Error processing response: {e}")
                     print(f"Stack trace:\n{traceback.format_exc()}")
-                    logger.error(f"Streaming error: {e}")
+                    logger.error(f"Response processing error: {e}")
                     logger.error(f"Stack trace: {traceback.format_exc()}")
 
             except EOFError:
@@ -404,6 +660,7 @@ def handle_chat_command(options: Namespace) -> None:
             except Exception as e:
                 print(f"\n❌ Error: {e}")
                 logger.error(f"Chat error: {e}")
+                traceback.print_exc()
 
     except Exception as e:
         logger.error(f"Failed to initialize chat session: {e}")
@@ -448,7 +705,6 @@ def handle_validate_command(options: Namespace) -> None:
 
 
 def setup_logging(verbosity: int) -> None:
-    logger.remove()
     levels: dict[int, str] = {
         0: "ERROR",
         1: "WARNING",
@@ -457,7 +713,7 @@ def setup_logging(verbosity: int) -> None:
         4: "TRACE",
     }
     level: str = levels.get(verbosity, "TRACE")
-    logger.add(sys.stderr, level=level)
+    configure_logging(level=level)
 
 
 def generate_bundle_from_template(config_path: Path, app_name: str) -> Path:
@@ -471,7 +727,7 @@ def generate_bundle_from_template(config_path: Path, app_name: str) -> Path:
     4. Returns the path to the generated file
 
     The generated databricks.yaml is overwritten on each deployment and is not tracked in git.
-    Schema reference remains pointing to ./schemas/bundle_config_schema.json.
+    The template contains cloud-specific targets (azure, aws, gcp) with appropriate node types.
 
     Args:
         config_path: Path to the app config file
@@ -508,39 +764,59 @@ def run_databricks_command(
     profile: Optional[str] = None,
     config: Optional[str] = None,
     target: Optional[str] = None,
+    cloud: Optional[str] = None,
     dry_run: bool = False,
 ) -> None:
-    """Execute a databricks CLI command with optional profile and target."""
+    """Execute a databricks CLI command with optional profile, target, and cloud.
+
+    Args:
+        command: The databricks CLI command to execute (e.g., ["bundle", "deploy"])
+        profile: Optional Databricks CLI profile name
+        config: Optional path to the configuration file
+        target: Optional bundle target name (if not provided, auto-generated from app name and cloud)
+        cloud: Optional cloud provider ('azure', 'aws', 'gcp'). Auto-detected if not specified.
+        dry_run: If True, print the command without executing
+    """
     config_path = Path(config) if config else None
 
     if config_path and not config_path.exists():
         logger.error(f"Configuration file {config_path} does not exist.")
         sys.exit(1)
 
-    # Load app config and generate bundle from template
+    # Load app config
     app_config: AppConfig = AppConfig.from_file(config_path) if config_path else None
     normalized_name: str = normalize_name(app_config.app.name) if app_config else None
 
+    # Auto-detect cloud provider if not specified
+    if not cloud:
+        cloud = detect_cloud_provider(profile)
+        if cloud:
+            logger.info(f"Auto-detected cloud provider: {cloud}")
+        else:
+            logger.warning("Could not detect cloud provider. Defaulting to 'azure'.")
+            cloud = "azure"
+
     # Generate app-specific bundle from template (overwrites databricks.yaml temporarily)
     if config_path and app_config:
         generate_bundle_from_template(config_path, normalized_name)
 
-    # Use app name as target if not explicitly provided
-    # This ensures each app gets its own Terraform state in .databricks/bundle/<app-name>/
-    if not target and normalized_name:
-        target = normalized_name
-        logger.debug(f"Using app-specific target: {target}")
+    # Use cloud as target (azure, aws, gcp) - can be overridden with explicit --target
+    if not target:
+        target = cloud
+        logger.debug(f"Using cloud-based target: {target}")
 
-    # Build databricks command (no -c flag needed, uses databricks.yaml in current dir)
+    # Build databricks command
+    # --profile is a global flag, --target is a subcommand flag for 'bundle'
     cmd = ["databricks"]
     if profile:
         cmd.extend(["--profile", profile])
 
+    cmd.extend(command)
+
+    # --target must come after the bundle subcommand (it's a subcommand-specific flag)
     if target:
         cmd.extend(["--target", target])
 
-    cmd.extend(command)
-
     # Add config_path variable for notebooks
     if config_path and app_config:
         # Calculate relative path from notebooks directory to config file
@@ -595,30 +871,38 @@ def handle_bundle_command(options: Namespace) -> None:
     profile: Optional[str] = options.profile
     config: Optional[str] = options.config
     target: Optional[str] = options.target
+    cloud: Optional[str] = options.cloud
     dry_run: bool = options.dry_run
 
     if options.deploy:
         logger.info("Deploying DAO AI asset bundle...")
         run_databricks_command(
-            ["bundle", "deploy"], profile, config, target, dry_run=dry_run
+            ["bundle", "deploy"],
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
+            dry_run=dry_run,
         )
     if options.run:
         logger.info("Running DAO AI system with current configuration...")
         # Use static job resource key that matches databricks.yaml (resources.jobs.deploy_job)
         run_databricks_command(
             ["bundle", "run", "deploy_job"],
-            profile,
-            config,
-            target,
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
             dry_run=dry_run,
         )
     if options.destroy:
         logger.info("Destroying DAO AI system with current configuration...")
         run_databricks_command(
             ["bundle", "destroy", "--auto-approve"],
-            profile,
-            config,
-            target,
+            profile=profile,
+            config=config,
+            target=target,
+            cloud=cloud,
             dry_run=dry_run,
         )
     else: