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

dao_ai/tools/mcp.py

@@ -1,21 +1,127 @@
+"""
+MCP (Model Context Protocol) tool creation for LangChain agents.
+
+This module provides tools for connecting to MCP servers using the
+MCP SDK and langchain-mcp-adapters library.
+
+For compatibility with Databricks APIs, we use manual tool wrappers
+that give us full control over the response format.
+
+Reference: https://docs.langchain.com/oss/python/langchain/mcp
+"""
+
 import asyncio
 from typing import Any, Sequence
 
-from databricks_mcp import DatabricksOAuthClientProvider
 from langchain_core.runnables.base import RunnableLike
 from langchain_core.tools import tool as create_tool
 from langchain_mcp_adapters.client import MultiServerMCPClient
-from langchain_mcp_adapters.tools import load_mcp_tools
 from loguru import logger
-from mcp import ClientSession
-from mcp.client.streamable_http import streamablehttp_client
-from mcp.types import ListToolsResult, Tool
+from mcp.types import CallToolResult, TextContent, Tool
 
 from dao_ai.config import (
     McpFunctionModel,
     TransportType,
+    value_of,
 )
-from dao_ai.tools.human_in_the_loop import as_human_in_the_loop
+
+
+def _build_connection_config(
+    function: McpFunctionModel,
+) -> dict[str, Any]:
+    """
+    Build the connection configuration dictionary for MultiServerMCPClient.
+
+    Args:
+        function: The MCP function model configuration.
+
+    Returns:
+        A dictionary containing the transport-specific connection settings.
+    """
+    if function.transport == TransportType.STDIO:
+        return {
+            "command": function.command,
+            "args": function.args,
+            "transport": function.transport.value,
+        }
+
+    # For HTTP transport with UC Connection, use DatabricksOAuthClientProvider
+    if function.connection:
+        from databricks_mcp import DatabricksOAuthClientProvider
+
+        workspace_client = function.connection.workspace_client
+        auth_provider = DatabricksOAuthClientProvider(workspace_client)
+
+        logger.trace(
+            "Using DatabricksOAuthClientProvider for authentication",
+            connection_name=function.connection.name,
+        )
+
+        return {
+            "url": function.mcp_url,
+            "transport": "http",
+            "auth": auth_provider,
+        }
+
+    # For HTTP transport with headers-based authentication
+    headers: dict[str, str] = {
+        key: str(value_of(val)) for key, val in function.headers.items()
+    }
+
+    if "Authorization" not in headers:
+        logger.trace("Generating fresh authentication token")
+
+        from dao_ai.providers.databricks import DatabricksProvider
+
+        try:
+            provider = DatabricksProvider(
+                workspace_host=value_of(function.workspace_host),
+                client_id=value_of(function.client_id),
+                client_secret=value_of(function.client_secret),
+                pat=value_of(function.pat),
+            )
+            headers["Authorization"] = f"Bearer {provider.create_token()}"
+            logger.trace("Generated fresh authentication token")
+        except Exception as e:
+            logger.error("Failed to create fresh token", error=str(e))
+    else:
+        logger.trace("Using existing authentication token")
+
+    return {
+        "url": function.mcp_url,
+        "transport": "http",
+        "headers": headers,
+    }
+
+
+def _extract_text_content(result: CallToolResult) -> str:
+    """
+    Extract text content from an MCP CallToolResult.
+
+    Converts the MCP result content to a plain string format that is
+    compatible with all LLM APIs (avoiding extra fields like 'id').
+
+    Args:
+        result: The MCP tool call result.
+
+    Returns:
+        A string containing the concatenated text content.
+    """
+    if not result.content:
+        return ""
+
+    text_parts: list[str] = []
+    for item in result.content:
+        if isinstance(item, TextContent):
+            text_parts.append(item.text)
+        elif hasattr(item, "text"):
+            # Handle other content types that have text
+            text_parts.append(str(item.text))
+        else:
+            # Fallback: convert to string representation
+            text_parts.append(str(item))
+
+    return "\n".join(text_parts)
 
 
 def create_mcp_tools(
@@ -25,152 +131,133 @@ def create_mcp_tools(
     Create tools for invoking Databricks MCP functions.
 
     Supports both direct MCP connections and UC Connection-based MCP access.
-    Uses session-based approach to handle authentication token expiration properly.
+    Uses manual tool wrappers to ensure response format compatibility with
+    Databricks APIs (which reject extra fields in tool results).
 
     Based on: https://docs.databricks.com/aws/en/generative-ai/mcp/external-mcp
+
+    Args:
+        function: The MCP function model configuration.
+
+    Returns:
+        A sequence of LangChain tools that can be used by agents.
     """
-    logger.debug(f"create_mcp_tools: {function}")
+    mcp_url = function.mcp_url
+    logger.debug("Creating MCP tools", mcp_url=mcp_url)
+
+    connection_config = _build_connection_config(function)
 
-    # Check if using UC Connection or direct MCP connection
     if function.connection:
-        # Use UC Connection approach with DatabricksOAuthClientProvider
-        logger.debug(f"Using UC Connection for MCP: {function.connection.name}")
+        logger.debug(
+            "Using UC Connection for MCP",
+            connection_name=function.connection.name,
+            mcp_url=mcp_url,
+        )
+    else:
+        logger.debug(
+            "Using direct connection for MCP",
+            transport=function.transport,
+            mcp_url=mcp_url,
+        )
 
-        # Construct URL if not provided
-        if function.url:
-            mcp_url = function.url
-            logger.debug(f"Using provided MCP URL: {mcp_url}")
-        else:
-            # Construct URL from workspace host and connection name
-            # Pattern: https://{workspace_host}/api/2.0/mcp/external/{connection_name}
-            workspace_client = function.connection.workspace_client
-            workspace_host = workspace_client.config.host
-            connection_name = function.connection.name
-            mcp_url = f"{workspace_host}/api/2.0/mcp/external/{connection_name}"
-            logger.debug(f"Constructed MCP URL from connection: {mcp_url}")
-
-        async def _get_tools_with_connection():
-            """Get tools using DatabricksOAuthClientProvider."""
-            workspace_client = function.connection.workspace_client
-
-            async with streamablehttp_client(
-                mcp_url, auth=DatabricksOAuthClientProvider(workspace_client)
-            ) as (read_stream, write_stream, _):
-                async with ClientSession(read_stream, write_stream) as session:
-                    # Initialize and list tools
-                    await session.initialize()
-                    tools = await load_mcp_tools(session)
-                    return tools
+    # Create client to list available tools
+    client = MultiServerMCPClient({"mcp_function": connection_config})
 
-        try:
-            langchain_tools = asyncio.run(_get_tools_with_connection())
+    async def _list_tools() -> list[Tool]:
+        """List available MCP tools from the server."""
+        async with client.session("mcp_function") as session:
+            result = await session.list_tools()
+            return result.tools if hasattr(result, "tools") else list(result)
+
+    try:
+        mcp_tools: list[Tool] = asyncio.run(_list_tools())
+
+        # Log discovered tools
+        logger.info(
+            "Discovered MCP tools",
+            tools_count=len(mcp_tools),
+            mcp_url=mcp_url,
+        )
+        for mcp_tool in mcp_tools:
             logger.debug(
-                f"Retrieved {len(langchain_tools)} MCP tools via UC Connection"
+                "MCP tool discovered",
+                tool_name=mcp_tool.name,
+                tool_description=(
+                    mcp_tool.description[:100] if mcp_tool.description else None
+                ),
             )
 
-            # Wrap tools with human-in-the-loop if needed
-            wrapped_tools = [
-                as_human_in_the_loop(tool, function) for tool in langchain_tools
-            ]
-            return wrapped_tools
-
-        except Exception as e:
-            logger.error(f"Failed to get tools from MCP server via UC Connection: {e}")
+    except Exception as e:
+        if function.connection:
+            logger.error(
+                "Failed to get tools from MCP server via UC Connection",
+                connection_name=function.connection.name,
+                error=str(e),
+            )
             raise RuntimeError(
-                f"Failed to list MCP tools for function '{function.name}' via UC Connection '{function.connection.name}': {e}"
+                f"Failed to list MCP tools via UC Connection "
+                f"'{function.connection.name}': {e}"
+            ) from e
+        else:
+            logger.error(
+                "Failed to get tools from MCP server",
+                transport=function.transport,
+                url=function.url,
+                error=str(e),
             )
+            raise RuntimeError(
+                f"Failed to list MCP tools with transport '{function.transport}' "
+                f"and URL '{function.url}': {e}"
+            ) from e
 
-    else:
-        # Use direct MCP connection with MultiServerMCPClient
-        logger.debug("Using direct MCP connection with MultiServerMCPClient")
-
-        def _create_fresh_connection() -> dict[str, Any]:
-            """Create connection config with fresh authentication headers."""
-            logger.debug("Creating fresh connection...")
-
-            if function.transport == TransportType.STDIO:
-                return {
-                    "command": function.command,
-                    "args": function.args,
-                    "transport": function.transport,
-                }
-
-            # For HTTP transport, generate fresh headers
-            headers = function.headers.copy() if function.headers else {}
-
-            if "Authorization" not in headers:
-                logger.debug("Generating fresh authentication token for MCP function")
-
-                from dao_ai.config import value_of
-                from dao_ai.providers.databricks import DatabricksProvider
-
-                try:
-                    provider = DatabricksProvider(
-                        workspace_host=value_of(function.workspace_host),
-                        client_id=value_of(function.client_id),
-                        client_secret=value_of(function.client_secret),
-                        pat=value_of(function.pat),
-                    )
-                    headers["Authorization"] = f"Bearer {provider.create_token()}"
-                    logger.debug("Generated fresh authentication token")
-                except Exception as e:
-                    logger.error(f"Failed to create fresh token: {e}")
-            else:
-                logger.debug("Using existing authentication token")
-
-            return {
-                "url": function.url,
-                "transport": function.transport,
-                "headers": headers,
-            }
-
-        # Get available tools from MCP server
-        async def _list_mcp_tools():
-            connection = _create_fresh_connection()
-            client = MultiServerMCPClient({function.name: connection})
+    def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
+        """
+        Create a LangChain tool wrapper for an MCP tool.
 
-            try:
-                async with client.session(function.name) as session:
-                    return await session.list_tools()
-            except Exception as e:
-                logger.error(f"Failed to list MCP tools: {e}")
-                return []
+        This wrapper handles:
+        - Fresh session creation per invocation (stateless)
+        - Content extraction to plain text (avoiding extra fields)
+        """
 
-        # Note: This still needs to run sync during tool creation/registration
-        # The actual tool execution will be async
-        try:
-            mcp_tools: list[Tool] | ListToolsResult = asyncio.run(_list_mcp_tools())
-            if isinstance(mcp_tools, ListToolsResult):
-                mcp_tools = mcp_tools.tools
+        @create_tool(
+            mcp_tool.name,
+            description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
+            args_schema=mcp_tool.inputSchema,
+        )
+        async def tool_wrapper(**kwargs: Any) -> str:
+            """Execute MCP tool with fresh session."""
+            logger.trace("Invoking MCP tool", tool_name=mcp_tool.name, args=kwargs)
 
-            logger.debug(f"Retrieved {len(mcp_tools)} MCP tools")
-        except Exception as e:
-            logger.error(f"Failed to get tools from MCP server: {e}")
-            raise RuntimeError(
-                f"Failed to list MCP tools for function '{function.name}' with transport '{function.transport}' and URL '{function.url}': {e}"
+            # Create a fresh client/session for each invocation
+            invocation_client = MultiServerMCPClient(
+                {"mcp_function": _build_connection_config(function)}
             )
 
-        # Create wrapper tools with fresh session per invocation
-        def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
-            @create_tool(
-                mcp_tool.name,
-                description=mcp_tool.description or f"MCP tool: {mcp_tool.name}",
-                args_schema=mcp_tool.inputSchema,
-            )
-            async def tool_wrapper(**kwargs):
-                """Execute MCP tool with fresh session and authentication."""
-                logger.debug(f"Invoking MCP tool {mcp_tool.name} with fresh session")
+            try:
+                async with invocation_client.session("mcp_function") as session:
+                    result: CallToolResult = await session.call_tool(
+                        mcp_tool.name, kwargs
+                    )
 
-                connection = _create_fresh_connection()
-                client = MultiServerMCPClient({function.name: connection})
+                    # Extract text content, avoiding extra fields
+                    text_result = _extract_text_content(result)
 
-                try:
-                    async with client.session(function.name) as session:
-                        return await session.call_tool(mcp_tool.name, kwargs)
-                except Exception as e:
-                    logger.error(f"MCP tool {mcp_tool.name} failed: {e}")
-                    raise
+                    logger.trace(
+                        "MCP tool completed",
+                        tool_name=mcp_tool.name,
+                        result_length=len(text_result),
+                    )
+
+                    return text_result
+
+            except Exception as e:
+                logger.error(
+                    "MCP tool failed",
+                    tool_name=mcp_tool.name,
+                    error=str(e),
+                )
+                raise
 
-            return as_human_in_the_loop(tool_wrapper, function)
+        return tool_wrapper
 
-        return [_create_tool_wrapper(tool) for tool in mcp_tools]
+    return [_create_tool_wrapper(tool) for tool in mcp_tools]

dao_ai/tools/memory.py

@@ -0,0 +1,50 @@
+"""Memory tools for DAO AI."""
+
+from typing import Any
+
+from langchain_core.tools import BaseTool, StructuredTool
+from langmem import create_search_memory_tool as langmem_create_search_memory_tool
+from pydantic import BaseModel, Field
+
+
+def create_search_memory_tool(namespace: tuple[str, ...]) -> BaseTool:
+    """
+    Create a Databricks-compatible search_memory tool.
+
+    The langmem search_memory tool has a 'filter' field with additionalProperties: true
+    in its schema, which Databricks LLM endpoints reject. This function creates a
+    wrapper tool that omits the problematic filter field.
+
+    Args:
+        namespace: The memory namespace tuple
+
+    Returns:
+        A StructuredTool compatible with Databricks
+    """
+    # Get the original tool
+    original_tool = langmem_create_search_memory_tool(namespace=namespace)
+
+    # Create a schema without the problematic filter field
+    class SearchMemoryInput(BaseModel):
+        """Input for search_memory tool."""
+
+        query: str = Field(..., description="The search query")
+        limit: int = Field(default=10, description="Maximum number of results")
+        offset: int = Field(default=0, description="Offset for pagination")
+
+    # Create a wrapper function
+    async def search_memory_wrapper(
+        query: str, limit: int = 10, offset: int = 0
+    ) -> Any:
+        """Search your long-term memories for information relevant to your current context."""
+        return await original_tool.ainvoke(
+            {"query": query, "limit": limit, "offset": offset}
+        )
+
+    # Create the new tool
+    return StructuredTool.from_function(
+        coroutine=search_memory_wrapper,
+        name="search_memory",
+        description="Search your long-term memories for information relevant to your current context.",
+        args_schema=SearchMemoryInput,
+    )

dao_ai/tools/python.py

@@ -7,7 +7,6 @@ from dao_ai.config import (
     FactoryFunctionModel,
     PythonFunctionModel,
 )
-from dao_ai.tools.human_in_the_loop import as_human_in_the_loop
 from dao_ai.utils import load_function
 
 
@@ -22,14 +21,11 @@ def create_factory_tool(
     Returns:
         A callable tool function that wraps the specified factory function
     """
-    logger.debug(f"create_factory_tool: {function}")
+    logger.trace("Creating factory tool", function=function.full_name)
 
     factory: Callable[..., Any] = load_function(function_name=function.full_name)
-    tool: Callable[..., Any] = factory(**function.args)
-    tool = as_human_in_the_loop(
-        tool=tool,
-        function=function,
-    )
+    tool: RunnableLike = factory(**function.args)
+    # HITL is now handled at middleware level via HumanInTheLoopMiddleware
     return tool
 
 
@@ -45,16 +41,15 @@ def create_python_tool(
     Returns:
         A callable tool function that wraps the specified Python function
     """
-    logger.debug(f"create_python_tool: {function}")
+    function_name = (
+        function.full_name if isinstance(function, PythonFunctionModel) else function
+    )
+    logger.trace("Creating Python tool", function=function_name)
 
     if isinstance(function, PythonFunctionModel):
         function = function.full_name
 
     # Load the Python function dynamically
-    tool: Callable[..., Any] = load_function(function_name=function)
-
-    tool = as_human_in_the_loop(
-        tool=tool,
-        function=function,
-    )
+    tool: RunnableLike = load_function(function_name=function)
+    # HITL is now handled at middleware level via HumanInTheLoopMiddleware
     return tool

dao_ai/tools/search.py

@@ -0,0 +1,14 @@
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.runnables.base import RunnableLike
+from loguru import logger
+
+
+def create_search_tool() -> RunnableLike:
+    """
+    Create a DuckDuckGo search tool.
+
+    Returns:
+        RunnableLike: A DuckDuckGo search tool that returns results as a list
+    """
+    logger.trace("Creating DuckDuckGo search tool")
+    return DuckDuckGoSearchRun(output_format="list")

dao_ai/tools/slack.py

@@ -26,7 +26,7 @@ def _find_channel_id_by_name(
     # Remove '#' prefix if present
     clean_name = channel_name.lstrip("#")
 
-    logger.debug(f"Looking up Slack channel ID for channel name: {clean_name}")
+    logger.trace("Looking up Slack channel ID", channel_name=clean_name)
 
     try:
         # Call Slack API to list conversations
@@ -37,14 +37,18 @@ def _find_channel_id_by_name(
         )
 
         if response.status_code != 200:
-            logger.error(f"Failed to list Slack channels: {response.text}")
+            logger.error(
+                "Failed to list Slack channels",
+                status_code=response.status_code,
+                response=response.text,
+            )
             return None
 
         # Parse response
         data = response.json()
 
         if not data.get("ok"):
-            logger.error(f"Slack API returned error: {data.get('error')}")
+            logger.error("Slack API returned error", error=data.get("error"))
             return None
 
         # Search for channel by name
@@ -53,15 +57,19 @@ def _find_channel_id_by_name(
             if channel.get("name") == clean_name:
                 channel_id = channel.get("id")
                 logger.debug(
-                    f"Found channel ID '{channel_id}' for channel name '{clean_name}'"
+                    "Found Slack channel ID",
+                    channel_id=channel_id,
+                    channel_name=clean_name,
                 )
                 return channel_id
 
-        logger.warning(f"Channel '{clean_name}' not found in Slack workspace")
+        logger.warning("Slack channel not found", channel_name=clean_name)
         return None
 
     except Exception as e:
-        logger.error(f"Error looking up Slack channel: {e}")
+        logger.error(
+            "Error looking up Slack channel", channel_name=clean_name, error=str(e)
+        )
         return None
 
 
@@ -71,7 +79,7 @@ def create_send_slack_message_tool(
     channel_name: Optional[str] = None,
     name: Optional[str] = None,
     description: Optional[str] = None,
-) -> Callable[[str], Any]:
+) -> Callable[[str], str]:
     """
     Create a tool that sends a message to a Slack channel.
 
@@ -87,7 +95,7 @@ def create_send_slack_message_tool(
 
     Based on: https://docs.databricks.com/aws/en/generative-ai/agent-framework/slack-agent
     """
-    logger.debug("create_send_slack_message_tool")
+    logger.trace("Creating send Slack message tool")
 
     # Validate inputs
     if channel_id is None and channel_name is None:
@@ -99,12 +107,16 @@ def create_send_slack_message_tool(
 
     # Look up channel_id from channel_name if needed
     if channel_id is None and channel_name is not None:
-        logger.debug(f"Looking up channel_id for channel_name: {channel_name}")
+        logger.trace(
+            "Looking up channel ID for channel name", channel_name=channel_name
+        )
         channel_id = _find_channel_id_by_name(connection, channel_name)
         if channel_id is None:
             raise ValueError(f"Could not find Slack channel with name '{channel_name}'")
         logger.debug(
-            f"Resolved channel_name '{channel_name}' to channel_id '{channel_id}'"
+            "Resolved channel name to ID",
+            channel_name=channel_name,
+            channel_id=channel_id,
         )
 
     if name is None: