dao-ai 0.0.25__py3-none-any.whl → 0.1.2__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,12 @@
-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")
 
 mlflow.langchain.autolog()
 
@@ -14,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,37 @@ 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
 
 
 env_path: str = find_dotenv()
@@ -240,9 +267,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 +291,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 +342,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 +385,217 @@ 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):
+                                logger.debug(f"AI message content: {msg.content}")
+                                logger.debug(
+                                    f"AI message has tool_calls: {hasattr(msg, 'tool_calls')}"
+                                )
+                                if hasattr(msg, "tool_calls"):
+                                    logger.debug(f"Tool calls: {msg.tool_calls}")
+
+                                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 +609,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 +654,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 +662,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: