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

dao_ai/state.py

@@ -1,41 +1,177 @@
-from langchain_core.messages import AnyMessage
+"""
+State definitions for DAO AI agents.
+
+This module defines the state schemas used by DAO AI agents,
+compatible with both LangChain v1's create_agent and LangGraph's StateGraph.
+
+State Schema:
+- AgentState: Primary state schema for all agent operations
+- Context: Runtime context passed via ToolRuntime[Context] or Runtime[Context]
+- GenieSpaceState: Per-space state for Genie conversations
+- SessionState: Accumulated state that flows between requests
+"""
+
+from datetime import datetime
+from typing import Any, Optional
+
 from langgraph.graph import MessagesState
-from langgraph.managed import RemainingSteps
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import NotRequired
 
 
-class IncomingState(MessagesState): ...
+class GenieSpaceState(BaseModel):
+    """State for a single Genie space/conversation.
 
+    This tracks the conversation state and metadata for a Genie space,
+    allowing multi-turn conversations and caching information to be preserved.
+    """
+
+    conversation_id: str = Field(description="Genie conversation ID for this space")
+    cache_hit: bool = Field(
+        default=False, description="Whether the last query was a cache hit"
+    )
+    cache_key: Optional[str] = Field(default=None, description="Cache key if cached")
+    follow_up_questions: list[str] = Field(
+        default_factory=list, description="Suggested follow-up questions from Genie"
+    )
+    last_query: Optional[str] = Field(
+        default=None, description="The last query sent to Genie"
+    )
+    last_query_time: Optional[datetime] = Field(
+        default=None, description="When the last query was made"
+    )
 
-class OutgoingState(MessagesState):
-    is_valid: bool
-    message_error: str
 
+class GenieState(BaseModel):
+    """State for all Genie spaces.
 
-class SharedState(MessagesState):
+    Maps space_id to GenieSpaceState for each Genie space the user has interacted with.
     """
-    State representation for the DAO AI agent conversation workflow.
 
-    Extends LangGraph's MessagesState to maintain the conversation history while
-    adding additional state fields specific to the DAO domain. This state is
-    passed between nodes in the agent graph and modified during execution.
+    spaces: dict[str, GenieSpaceState] = Field(
+        default_factory=dict, description="Map of space_id to space state"
+    )
+
+    def get_conversation_id(self, space_id: str) -> Optional[str]:
+        """Get conversation ID for a space, if it exists."""
+        if space_id in self.spaces:
+            return self.spaces[space_id].conversation_id
+        return None
+
+    def update_space(
+        self,
+        space_id: str,
+        conversation_id: str,
+        cache_hit: bool = False,
+        cache_key: Optional[str] = None,
+        follow_up_questions: Optional[list[str]] = None,
+        last_query: Optional[str] = None,
+    ) -> None:
+        """Update or create state for a Genie space."""
+        self.spaces[space_id] = GenieSpaceState(
+            conversation_id=conversation_id,
+            cache_hit=cache_hit,
+            cache_key=cache_key,
+            follow_up_questions=follow_up_questions or [],
+            last_query=last_query,
+            last_query_time=datetime.now() if last_query else None,
+        )
+
+
+class SessionState(BaseModel):
+    """Accumulated state that flows between requests.
+
+    This is the "paste from previous output" portion of the request.
+    Users can copy the session from custom_outputs and paste it back
+    as custom_inputs.session to restore state.
     """
 
-    context: str  # short term/long term memory
+    genie: GenieState = Field(
+        default_factory=GenieState, description="Genie conversation state per space"
+    )
+
+    # Future: Add other stateful tool state here
+    # other_tool_state: OtherToolState = Field(default_factory=OtherToolState)
+
 
-    active_agent: str  # langgraph-swarm
-    remaining_steps: RemainingSteps  # langgraph-supervisor
+class AgentState(MessagesState, total=False):
+    """
+    Primary state schema for DAO AI agents.
+
+    Extends MessagesState to include the messages channel with proper
+    add_messages reducer, plus additional fields for DAO AI functionality.
 
-    summarized_messages: list[AnyMessage]
+    Used for:
+    - state_schema in create_agent calls
+    - state_schema in StateGraph for orchestration
+    - Type parameter in ToolRuntime[Context, AgentState]
+    - Type parameter in AgentMiddleware[AgentState, Context]
+    - API input/output contracts
 
-    is_valid: bool  # message validation node
-    message_error: str
+    Fields:
+        messages: Conversation history with add_messages reducer (from MessagesState)
+        context: Short/long term memory context
+        active_agent: Name of currently active agent in multi-agent workflows
+        is_valid: Message validation status
+        message_error: Error message if validation failed
+        session: Accumulated session state (genie conversations, etc.)
+        structured_response: Structured output from response_format (populated by LangChain)
+    """
 
-    # A mapping of genie space_id to conversation_id
-    genie_conversation_ids: dict[str, str]  # Genie
+    context: NotRequired[str]
+    active_agent: NotRequired[str]
+    is_valid: NotRequired[bool]
+    message_error: NotRequired[str]
+    session: NotRequired[SessionState]
+    structured_response: NotRequired[Any]
 
 
 class Context(BaseModel):
+    """
+    Runtime context for DAO AI agents.
+
+    This is passed to tools and middleware via the runtime parameter.
+    Access via ToolRuntime[Context] in tools or Runtime[Context] in middleware.
+
+    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.store_num  # Direct attribute access
+            return f"Hello, {user_id} at store {store_num}!"
+
+        class MyMiddleware(AgentMiddleware[AgentState, Context]):
+            def before_model(
+                self,
+                state: AgentState,
+                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
-    store_num: int | None = None
+
+    @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

@@ -1,13 +1,15 @@
+from dao_ai.genie.cache import LRUCacheService, SemanticCacheService
 from dao_ai.hooks.core import create_hooks
 from dao_ai.tools.agent import create_agent_endpoint_tool
-from dao_ai.tools.core import (
-    create_tools,
-    search_tool,
-)
+from dao_ai.tools.core import create_tools, say_hello_tool
+from dao_ai.tools.email import create_send_email_tool
 from dao_ai.tools.genie import create_genie_tool
 from dao_ai.tools.mcp import create_mcp_tools
+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,
@@ -23,11 +25,15 @@ 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",
     "create_mcp_tools",
     "create_python_tool",
+    "create_search_memory_tool",
+    "create_search_tool",
+    "create_send_email_tool",
     "create_send_slack_message_tool",
     "create_tools",
     "create_uc_tools",
@@ -35,7 +41,9 @@ __all__ = [
     "current_time_tool",
     "format_time_tool",
     "is_business_hours_tool",
-    "search_tool",
+    "LRUCacheService",
+    "say_hello_tool",
+    "SemanticCacheService",
     "time_difference_tool",
     "time_in_timezone_tool",
     "time_until_tool",

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

@@ -1,7 +1,19 @@
+"""
+Core tool creation infrastructure for DAO AI.
+
+This module provides the foundational tool creation and registration system:
+- Tool registry for caching created tools
+- Factory function for creating tools from configuration
+- Example tools demonstrating runtime context usage
+
+This is "core" because it contains the essential infrastructure that all
+tool creation flows through, not because it contains all tools.
+"""
+
 from collections import OrderedDict
 from typing import Sequence
 
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.tools import ToolRuntime, tool
 from langchain_core.runnables.base import RunnableLike
 from loguru import logger
 
@@ -10,7 +22,9 @@ from dao_ai.config import (
     ToolModel,
 )
 from dao_ai.hooks.core import create_hooks
+from dao_ai.state import Context
 
+# Module-level tool registry for caching created tools
 tool_registry: dict[str, Sequence[RunnableLike]] = {}
 
 
@@ -22,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
@@ -33,27 +47,66 @@ 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] = tool_registry.get(name)
+        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
 
 
-def search_tool() -> RunnableLike:
-    logger.debug("search_tool")
-    return DuckDuckGoSearchRun(output_format="list")
+# =============================================================================
+# Example Tools
+# =============================================================================
+# The following tools serve as examples and are included here because they
+# demonstrate core patterns (like ToolRuntime usage) rather than because they
+# are fundamental infrastructure. They're simple enough to colocate with the
+# core tool creation logic.
+
+
+@tool
+def say_hello_tool(
+    name: str | None = None,
+    runtime: ToolRuntime[Context] = None,
+) -> str:
+    """
+    Say hello to someone by name.
+
+    This is an example tool demonstrating how to use ToolRuntime to access
+    runtime context (like user_id) within a tool.
+
+    If no name is provided, uses the user_id from the runtime context.
+
+    Args:
+        name: Optional name of the person to greet. If not provided,
+              uses user_id from context.
+        runtime: Runtime context (automatically injected, not provided by user)
+
+    Returns:
+        A greeting string
+    """
+    # Use provided name, or fall back to user_id from context
+    if name is None:
+        if runtime and runtime.context:
+            user_id: str | None = runtime.context.user_id
+            if user_id:
+                name = user_id
+            else:
+                name = "there"  # Default fallback
+        else:
+            name = "there"  # Default fallback
+
+    return f"Hello, {name}!"

dao_ai/tools/email.py

@@ -0,0 +1,232 @@
+"""Email tool for sending emails via SMTP."""
+
+import smtplib
+from email.mime.multipart import MIMEMultipart
+from email.mime.text import MIMEText
+from typing import Any, Callable, Optional
+
+from langchain_core.tools import tool
+from loguru import logger
+from pydantic import BaseModel, ConfigDict, Field
+
+from dao_ai.config import AnyVariable, value_of
+
+
+class SMTPConfigModel(BaseModel):
+    """Configuration model for SMTP email settings."""
+
+    model_config = ConfigDict(use_enum_values=True, extra="forbid")
+
+    host: AnyVariable = Field(
+        default="smtp.gmail.com",
+        description="SMTP server hostname",
+    )
+    port: AnyVariable = Field(
+        default=587,
+        description="SMTP server port",
+    )
+    username: AnyVariable = Field(
+        description="SMTP username for authentication",
+    )
+    password: AnyVariable = Field(
+        description="SMTP password for authentication",
+    )
+    sender_email: Optional[AnyVariable] = Field(
+        default=None,
+        description="Email address to use as sender (defaults to username)",
+    )
+    use_tls: bool = Field(
+        default=True,
+        description="Whether to use TLS encryption",
+    )
+
+
+def create_send_email_tool(
+    smtp_config: SMTPConfigModel | dict[str, Any],
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+) -> Callable[[str, str, str, Optional[str]], str]:
+    """
+    Create a tool that sends emails via SMTP.
+
+    This factory function creates a tool for sending emails with configurable SMTP settings.
+    All configuration values support AnyVariable types, allowing use of environment variables,
+    secrets, and composite variables.
+
+    Args:
+        smtp_config: SMTP configuration (SMTPConfigModel or dict). Supports:
+            - host: SMTP server hostname (supports variables/secrets)
+            - port: SMTP server port (supports variables/secrets)
+            - username: SMTP username (supports variables/secrets)
+            - password: SMTP password (supports variables/secrets)
+            - sender_email: Sender email address, defaults to username (supports variables/secrets)
+            - use_tls: Whether to use TLS encryption (default: True)
+        name: Custom tool name (default: 'send_email')
+        description: Custom tool description
+
+    Returns:
+        A tool function that sends emails via SMTP
+
+    Example:
+        Basic usage with environment variables:
+        ```yaml
+        tools:
+          send_email:
+            name: send_email
+            function:
+              type: factory
+              name: dao_ai.tools.email.create_send_email_tool
+              args:
+                smtp_config:
+                  host: smtp.gmail.com
+                  port: 587
+                  username: ${SMTP_USER}
+                  password: ${SMTP_PASSWORD}
+                  sender_email: bot@example.com
+        ```
+
+        With secrets:
+        ```yaml
+        tools:
+          send_email:
+            name: send_email
+            function:
+              type: factory
+              name: dao_ai.tools.email.create_send_email_tool
+              args:
+                smtp_config:
+                  host: smtp.gmail.com
+                  port: 587
+                  username:
+                    type: secret
+                    scope: email
+                    key: smtp_user
+                  password:
+                    type: secret
+                    scope: email
+                    key: smtp_password
+        ```
+    """
+    logger.debug(
+        "Creating send_email_tool",
+        config_type=type(smtp_config).__name__,
+        tool_name=name,
+    )
+
+    # Convert dict to SMTPConfigModel if needed
+    if isinstance(smtp_config, dict):
+        smtp_config = SMTPConfigModel(**smtp_config)
+
+    # Resolve all variable values
+    host: str = value_of(smtp_config.host)
+    port: int = int(value_of(smtp_config.port))
+    username: str = value_of(smtp_config.username)
+    password: str = value_of(smtp_config.password)
+    sender_email: str = (
+        value_of(smtp_config.sender_email) if smtp_config.sender_email else username
+    )
+    use_tls: bool = smtp_config.use_tls
+
+    logger.info(
+        "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"
+    if description is None:
+        description = "Send an email to a recipient with subject and body content"
+
+    logger.debug("Creating email tool with decorator", tool_name=name)
+
+    @tool(
+        name_or_callable=name,
+        description=description,
+    )
+    def send_email(
+        to: str,
+        subject: str,
+        body: str,
+        cc: Optional[str] = None,
+    ) -> str:
+        """
+        Send an email via SMTP.
+
+        Args:
+            to: Recipient email address
+            subject: Email subject line
+            body: Email body content (plain text)
+            cc: Optional CC recipients (comma-separated email addresses)
+
+        Returns:
+            str: Success or error message
+        """
+        logger.info(
+            "Sending email", to=to, subject=subject, body_length=len(body), cc=cc
+        )
+
+        try:
+            # Create message
+            msg = MIMEMultipart()
+            msg["From"] = sender_email
+            msg["To"] = to
+            msg["Subject"] = subject
+
+            if cc:
+                msg["Cc"] = cc
+
+            # Attach body as plain text
+            msg.attach(MIMEText(body, "plain"))
+
+            # Send email
+            logger.debug("Connecting to SMTP server", host=host, port=port)
+            with smtplib.SMTP(host, port) as server:
+                if use_tls:
+                    logger.trace("Upgrading to TLS")
+                    server.starttls()
+
+                logger.trace("Authenticating", username=username)
+                server.login(username, password)
+
+                # Build recipient list
+                recipients = [to]
+                if cc:
+                    cc_addresses = [addr.strip() for addr in cc.split(",")]
+                    recipients.extend(cc_addresses)
+
+                logger.debug("Sending message", recipients_count=len(recipients))
+                server.send_message(msg)
+
+            success_msg = f"✓ Email sent successfully to {to}"
+            if cc:
+                success_msg += f" (cc: {cc})"
+
+            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(
+                "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("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(
+                "Failed to send email", error_type=type(e).__name__, error=str(e)
+            )
+            return error_msg
+
+    logger.success("Email tool created", tool_name=name)
+    return send_email