dao-ai 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

dao_ai/state.py

@@ -15,7 +15,7 @@ from datetime import datetime
 from typing import Any, Optional
 
 from langgraph.graph import MessagesState
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 from typing_extensions import NotRequired
 
 
@@ -133,15 +133,18 @@ class Context(BaseModel):
     This is passed to tools and middleware via the runtime parameter.
     Access via ToolRuntime[Context] in tools or Runtime[Context] in middleware.
 
-    The `custom` dict allows application-specific context values that can be:
-    - Used as template parameters in prompts (all keys are applied)
-    - Validated by middleware (check for specific keys like "store_num")
+    Additional fields beyond user_id and thread_id can be added dynamically
+    and will be available as top-level attributes on the context object.
+    These fields are:
+    - Used as template parameters in prompts (all fields are applied)
+    - Validated by middleware (check for specific fields like "store_num")
+    - Accessible as direct attributes (e.g., context.store_num)
 
     Example:
         @tool
         def my_tool(runtime: ToolRuntime[Context]) -> str:
             user_id = runtime.context.user_id
-            store_num = runtime.context.custom.get("store_num")
+            store_num = runtime.context.store_num  # Direct attribute access
             return f"Hello, {user_id} at store {store_num}!"
 
         class MyMiddleware(AgentMiddleware[AgentState, Context]):
@@ -151,9 +154,24 @@ class Context(BaseModel):
                 runtime: Runtime[Context]
             ) -> dict[str, Any] | None:
                 user_id = runtime.context.user_id
+                store_num = getattr(runtime.context, "store_num", None)
                 return None
     """
 
+    model_config = ConfigDict(
+        extra="allow"
+    )  # Allow extra fields as top-level attributes
+
     user_id: str | None = None
     thread_id: str | None = None
-    custom: dict[str, Any] = Field(default_factory=dict)
+
+    @classmethod
+    def from_runnable_config(cls, config: dict[str, Any]) -> "Context":
+        """
+        Create Context from LangChain RunnableConfig.
+
+        This method is called by LangChain when context_schema is provided to create_agent.
+        It extracts the 'configurable' dict from the config and uses it to instantiate Context.
+        """
+        configurable = config.get("configurable", {})
+        return cls(**configurable)

dao_ai/tools/__init__.py

@@ -9,6 +9,7 @@ from dao_ai.tools.memory import create_search_memory_tool
 from dao_ai.tools.python import create_factory_tool, create_python_tool
 from dao_ai.tools.search import create_search_tool
 from dao_ai.tools.slack import create_send_slack_message_tool
+from dao_ai.tools.sql import create_execute_statement_tool
 from dao_ai.tools.time import (
     add_time_tool,
     current_time_tool,
@@ -24,6 +25,7 @@ from dao_ai.tools.vector_search import create_vector_search_tool
 __all__ = [
     "add_time_tool",
     "create_agent_endpoint_tool",
+    "create_execute_statement_tool",
     "create_factory_tool",
     "create_genie_tool",
     "create_hooks",

dao_ai/tools/agent.py

@@ -14,9 +14,7 @@ def create_agent_endpoint_tool(
     name: Optional[str] = None,
     description: Optional[str] = None,
 ) -> Callable[..., Any]:
-    logger.debug(
-        f"Creating agent endpoint tool with name: {name} and description: {description}"
-    )
+    logger.debug("Creating agent endpoint tool", name=name, description=description)
 
     default_description: str = dedent("""
     This tool allows you to interact with a language model endpoint to answer questions.

dao_ai/tools/core.py

@@ -36,7 +36,7 @@ def create_tools(tool_models: Sequence[ToolModel]) -> Sequence[RunnableLike]:
     Each tool is created according to its type and parameters defined in the configuration.
 
     Args:
-        tool_configs: A sequence of dictionaries containing tool configurations
+        tool_models: A sequence of ToolModel configurations
 
     Returns:
         A sequence of BaseTool objects created from the provided configurations
@@ -47,24 +47,24 @@ def create_tools(tool_models: Sequence[ToolModel]) -> Sequence[RunnableLike]:
     for tool_config in tool_models:
         name: str = tool_config.name
         if name in tools:
-            logger.warning(f"Tools already registered for: {name}, skipping creation.")
+            logger.warning("Tools already registered, skipping", tool_name=name)
             continue
         registered_tools: Sequence[RunnableLike] | None = tool_registry.get(name)
         if registered_tools is None:
-            logger.debug(f"Creating tools for: {name}...")
+            logger.trace("Creating tools", tool_name=name)
             function: AnyTool = tool_config.function
             registered_tools = create_hooks(function)
-            logger.debug(f"Registering tools for: {tool_config}")
+            logger.trace("Registering tools", tool_name=name)
             tool_registry[name] = registered_tools
         else:
-            logger.debug(f"Tools already registered for: {name}")
+            logger.trace("Tools already registered", tool_name=name)
 
         tools[name] = registered_tools
 
     all_tools: Sequence[RunnableLike] = [
         t for tool_list in tools.values() for t in tool_list
     ]
-    logger.debug(f"Created tools: {all_tools}")
+    logger.debug("Tools created", tools_count=len(all_tools))
     return all_tools
 
 
@@ -101,7 +101,7 @@ def say_hello_tool(
     # Use provided name, or fall back to user_id from context
     if name is None:
         if runtime and runtime.context:
-            user_id = runtime.context.user_id
+            user_id: str | None = runtime.context.user_id
             if user_id:
                 name = user_id
             else:

dao_ai/tools/email.py

@@ -107,70 +107,41 @@ def create_send_email_tool(
                     key: smtp_password
         ```
     """
-    logger.info("=== Creating send_email_tool ===")
     logger.debug(
-        f"Factory called with config type: {type(smtp_config).__name__}, "
-        f"name={name}, description={description}"
+        "Creating send_email_tool",
+        config_type=type(smtp_config).__name__,
+        tool_name=name,
     )
 
     # Convert dict to SMTPConfigModel if needed
     if isinstance(smtp_config, dict):
-        logger.debug("Converting dict config to SMTPConfigModel")
         smtp_config = SMTPConfigModel(**smtp_config)
-    else:
-        logger.debug("Config already is SMTPConfigModel")
 
     # Resolve all variable values
-    logger.debug("Resolving SMTP configuration variables...")
-
-    logger.debug("  - Resolving host")
     host: str = value_of(smtp_config.host)
-    logger.debug(f"    Host resolved: {host}")
-
-    logger.debug("  - Resolving port")
     port: int = int(value_of(smtp_config.port))
-    logger.debug(f"    Port resolved: {port}")
-
-    logger.debug("  - Resolving username")
     username: str = value_of(smtp_config.username)
-    logger.debug(f"    Username resolved: {username}")
-
-    logger.debug("  - Resolving password")
     password: str = value_of(smtp_config.password)
-    logger.debug(
-        f"    Password resolved: {'*' * len(password) if password else 'None'}"
-    )
-
-    logger.debug("  - Resolving sender_email")
     sender_email: str = (
         value_of(smtp_config.sender_email) if smtp_config.sender_email else username
     )
-    logger.debug(
-        f"    Sender email resolved: {sender_email} "
-        f"({'from sender_email' if smtp_config.sender_email else 'defaulted to username'})"
-    )
-
     use_tls: bool = smtp_config.use_tls
-    logger.debug(f"  - TLS enabled: {use_tls}")
 
     logger.info(
-        f"SMTP configuration resolved - host={host}, port={port}, "
-        f"sender={sender_email}, use_tls={use_tls}"
+        "SMTP configuration resolved",
+        host=host,
+        port=port,
+        sender=sender_email,
+        use_tls=use_tls,
+        password_set=bool(password),
     )
 
     if name is None:
         name = "send_email"
-        logger.debug(f"Tool name defaulted to: {name}")
-    else:
-        logger.debug(f"Tool name set to: {name}")
-
     if description is None:
         description = "Send an email to a recipient with subject and body content"
-        logger.debug("Tool description using default")
-    else:
-        logger.debug(f"Tool description set to: {description}")
 
-    logger.info(f"Creating tool '{name}' with @tool decorator")
+    logger.debug("Creating email tool with decorator", tool_name=name)
 
     @tool(
         name_or_callable=name,
@@ -194,87 +165,68 @@ def create_send_email_tool(
         Returns:
             str: Success or error message
         """
-        logger.info("=== send_email tool invoked ===")
-        logger.info(f"  To: {to}")
-        logger.info(f"  Subject: {subject}")
-        logger.info(f"  Body length: {len(body)} characters")
-        logger.info(f"  CC: {cc if cc else 'None'}")
+        logger.info(
+            "Sending email", to=to, subject=subject, body_length=len(body), cc=cc
+        )
 
         try:
-            logger.debug("Constructing email message...")
-
             # Create message
             msg = MIMEMultipart()
             msg["From"] = sender_email
             msg["To"] = to
             msg["Subject"] = subject
-            logger.debug(f"  From: {sender_email}")
-            logger.debug(f"  To: {to}")
-            logger.debug(f"  Subject: {subject}")
 
             if cc:
                 msg["Cc"] = cc
-                logger.debug(f"  CC: {cc}")
 
             # Attach body as plain text
             msg.attach(MIMEText(body, "plain"))
-            logger.debug(f"  Body attached ({len(body)} chars)")
 
             # Send email
-            logger.info(f"Connecting to SMTP server {host}:{port}...")
+            logger.debug("Connecting to SMTP server", host=host, port=port)
             with smtplib.SMTP(host, port) as server:
-                logger.debug("SMTP connection established")
-
                 if use_tls:
-                    logger.debug("Upgrading connection to TLS...")
+                    logger.trace("Upgrading to TLS")
                     server.starttls()
-                    logger.debug("TLS upgrade successful")
 
-                logger.debug(f"Authenticating with username: {username}")
+                logger.trace("Authenticating", username=username)
                 server.login(username, password)
-                logger.info("SMTP authentication successful")
 
                 # Build recipient list
                 recipients = [to]
                 if cc:
                     cc_addresses = [addr.strip() for addr in cc.split(",")]
                     recipients.extend(cc_addresses)
-                    logger.debug(f"Total recipients: {len(recipients)} ({recipients})")
-                else:
-                    logger.debug(f"Single recipient: {to}")
 
-                logger.info(f"Sending message to {len(recipients)} recipient(s)...")
+                logger.debug("Sending message", recipients_count=len(recipients))
                 server.send_message(msg)
-                logger.info("Message sent successfully via SMTP")
 
             success_msg = f"✓ Email sent successfully to {to}"
             if cc:
                 success_msg += f" (cc: {cc})"
 
-            logger.info(success_msg)
-            logger.info("=== send_email completed successfully ===")
+            logger.success("Email sent successfully", to=to, cc=cc)
             return success_msg
 
         except smtplib.SMTPAuthenticationError as e:
             error_msg = f"✗ SMTP authentication failed: {str(e)}"
-            logger.error(error_msg)
-            logger.error(f"  Server: {host}:{port}")
-            logger.error(f"  Username: {username}")
-            logger.error("=== send_email failed (authentication) ===")
+            logger.error(
+                "SMTP authentication failed",
+                server=f"{host}:{port}",
+                username=username,
+                error=str(e),
+            )
             return error_msg
         except smtplib.SMTPException as e:
             error_msg = f"✗ SMTP error: {str(e)}"
-            logger.error(error_msg)
-            logger.error(f"  Server: {host}:{port}")
-            logger.error("=== send_email failed (SMTP error) ===")
+            logger.error("SMTP error", server=f"{host}:{port}", error=str(e))
             return error_msg
         except Exception as e:
             error_msg = f"✗ Failed to send email: {str(e)}"
-            logger.error(error_msg)
-            logger.error(f"  Error type: {type(e).__name__}")
-            logger.error("=== send_email failed (unexpected error) ===")
+            logger.error(
+                "Failed to send email", error_type=type(e).__name__, error=str(e)
+            )
             return error_msg
 
-    logger.info(f"Tool '{name}' created successfully")
-    logger.info("=== send_email_tool creation complete ===")
+    logger.success("Email tool created", tool_name=name)
     return send_email

dao_ai/tools/genie.py

@@ -89,15 +89,15 @@ def create_genie_tool(
     Returns:
         A LangGraph tool that processes natural language queries through Genie
     """
-    logger.debug("create_genie_tool")
-    logger.debug(f"genie_room type: {type(genie_room)}")
-    logger.debug(f"genie_room: {genie_room}")
-    logger.debug(f"persist_conversation: {persist_conversation}")
-    logger.debug(f"truncate_results: {truncate_results}")
-    logger.debug(f"name: {name}")
-    logger.debug(f"description: {description}")
-    logger.debug(f"lru_cache_parameters: {lru_cache_parameters}")
-    logger.debug(f"semantic_cache_parameters: {semantic_cache_parameters}")
+    logger.debug(
+        "Creating Genie tool",
+        genie_room_type=type(genie_room).__name__,
+        persist_conversation=persist_conversation,
+        truncate_results=truncate_results,
+        name=name,
+        has_lru_cache=lru_cache_parameters is not None,
+        has_semantic_cache=semantic_cache_parameters is not None,
+    )
 
     if isinstance(genie_room, dict):
         genie_room = GenieRoomModel(**genie_room)
@@ -188,8 +188,10 @@ GenieResponse: A response object containing the conversation ID and result from
         existing_conversation_id: str | None = session.genie.get_conversation_id(
             space_id_str
         )
-        logger.debug(
-            f"Existing conversation ID for space {space_id_str}: {existing_conversation_id}"
+        logger.trace(
+            "Using existing conversation ID",
+            space_id=space_id_str,
+            conversation_id=existing_conversation_id,
         )
 
         # Call ask_question which always returns CacheResult with cache metadata
@@ -202,8 +204,11 @@ GenieResponse: A response object containing the conversation ID and result from
 
         current_conversation_id: str = genie_response.conversation_id
         logger.debug(
-            f"Current conversation ID for space {space_id_str}: {current_conversation_id}, "
-            f"cache_hit: {cache_hit}, cache_key: {cache_key}"
+            "Genie question answered",
+            space_id=space_id_str,
+            conversation_id=current_conversation_id,
+            cache_hit=cache_hit,
+            cache_key=cache_key,
         )
 
         # Update session state with cache information