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

dao_ai/tools/mcp.py

@@ -1,196 +1,263 @@
+"""
+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 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,
 )
 
 
-def create_mcp_tools(
+def _build_connection_config(
     function: McpFunctionModel,
-) -> Sequence[RunnableLike]:
+) -> dict[str, Any]:
     """
-    Create tools for invoking Databricks MCP functions.
+    Build the connection configuration dictionary for MultiServerMCPClient.
 
-    Supports both direct MCP connections and UC Connection-based MCP access.
-    Uses session-based approach to handle authentication token expiration properly.
+    Args:
+        function: The MCP function model configuration.
 
-    Based on: https://docs.databricks.com/aws/en/generative-ai/mcp/external-mcp
+    Returns:
+        A dictionary containing the transport-specific connection settings.
     """
-    logger.debug(f"create_mcp_tools: {function}")
+    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
 
-    # Get MCP URL - handles all convenience objects (connection, genie_room, warehouse, etc.)
-    mcp_url = function.mcp_url
-    logger.debug(f"Using MCP URL: {mcp_url}")
+        workspace_client = function.connection.workspace_client
+        auth_provider = DatabricksOAuthClientProvider(workspace_client)
 
-    # 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.trace(
+            "Using DatabricksOAuthClientProvider for authentication",
+            connection_name=function.connection.name,
+        )
 
-        async def _list_tools_with_connection():
-            """List available tools using DatabricksOAuthClientProvider."""
-            workspace_client = function.connection.workspace_client
+        return {
+            "url": function.mcp_url,
+            "transport": "http",
+            "auth": auth_provider,
+        }
 
-            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()
-                    return await session.list_tools()
+    # 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:
-            mcp_tools: list[Tool] | ListToolsResult = asyncio.run(
-                _list_tools_with_connection()
+            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),
             )
-            if isinstance(mcp_tools, ListToolsResult):
-                mcp_tools = mcp_tools.tools
+            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")
 
-            logger.debug(f"Retrieved {len(mcp_tools)} MCP tools via UC Connection")
+    return {
+        "url": function.mcp_url,
+        "transport": "http",
+        "headers": headers,
+    }
 
-        except Exception as e:
-            logger.error(f"Failed to get tools from MCP server via UC Connection: {e}")
-            raise RuntimeError(
-                f"Failed to list MCP tools for function '{function.name}' via UC Connection '{function.connection.name}': {e}"
-            )
 
-        # Create wrapper tools with fresh session per invocation
-        def _create_tool_wrapper_with_connection(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 UC Connection session."""
-                logger.debug(
-                    f"Invoking MCP tool {mcp_tool.name} with fresh UC Connection session"
-                )
-                workspace_client = function.connection.workspace_client
-
-                try:
-                    async with streamablehttp_client(
-                        mcp_url, auth=DatabricksOAuthClientProvider(workspace_client)
-                    ) as (read_stream, write_stream, _):
-                        async with ClientSession(read_stream, write_stream) as session:
-                            await session.initialize()
-                            result = await session.call_tool(mcp_tool.name, kwargs)
-                            logger.debug(
-                                f"MCP tool {mcp_tool.name} completed successfully"
-                            )
-                            return result
-                except Exception as e:
-                    logger.error(f"MCP tool {mcp_tool.name} failed: {e}")
-                    raise
-
-            # HITL is now handled at middleware level via HumanInTheLoopMiddleware
-            return tool_wrapper
-
-        return [_create_tool_wrapper_with_connection(tool) for tool in mcp_tools]
+def _extract_text_content(result: CallToolResult) -> str:
+    """
+    Extract text content from an MCP CallToolResult.
 
-    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": mcp_url,  # Use the resolved MCP 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})
+    Converts the MCP result content to a plain string format that is
+    compatible with all LLM APIs (avoiding extra fields like 'id').
 
-            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 []
+    Args:
+        result: The MCP tool call result.
 
-        # 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
+    Returns:
+        A string containing the concatenated text content.
+    """
+    if not result.content:
+        return ""
 
-            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}"
+    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(
+    function: McpFunctionModel,
+) -> Sequence[RunnableLike]:
+    """
+    Create tools for invoking Databricks MCP functions.
+
+    Supports both direct MCP connections and UC Connection-based MCP access.
+    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.
+    """
+    mcp_url = function.mcp_url
+    logger.debug("Creating MCP tools", mcp_url=mcp_url)
+
+    connection_config = _build_connection_config(function)
+
+    if function.connection:
+        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,
+        )
+
+    # Create client to list available tools
+    client = MultiServerMCPClient({"mcp_function": connection_config})
+
+    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(
+                "MCP tool discovered",
+                tool_name=mcp_tool.name,
+                tool_description=(
+                    mcp_tool.description[:100] if mcp_tool.description else None
+                ),
             )
 
-        # 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,
+    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 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
+
+    def _create_tool_wrapper(mcp_tool: Tool) -> RunnableLike:
+        """
+        Create a LangChain tool wrapper for an MCP tool.
+
+        This wrapper handles:
+        - Fresh session creation per invocation (stateless)
+        - Content extraction to plain text (avoiding extra fields)
+        """
+
+        @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)
+
+            # Create a fresh client/session for each invocation
+            invocation_client = MultiServerMCPClient(
+                {"mcp_function": _build_connection_config(function)}
             )
-            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")
 
-                connection = _create_fresh_connection()
-                client = MultiServerMCPClient({function.name: connection})
+            try:
+                async with invocation_client.session("mcp_function") as session:
+                    result: CallToolResult = await session.call_tool(
+                        mcp_tool.name, kwargs
+                    )
+
+                    # Extract text content, avoiding extra fields
+                    text_result = _extract_text_content(result)
+
+                    logger.trace(
+                        "MCP tool completed",
+                        tool_name=mcp_tool.name,
+                        result_length=len(text_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
+                    return text_result
+
+            except Exception as e:
+                logger.error(
+                    "MCP tool failed",
+                    tool_name=mcp_tool.name,
+                    error=str(e),
+                )
+                raise
 
-            # HITL is now handled at middleware level via HumanInTheLoopMiddleware
-            return tool_wrapper
+        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/python.py

@@ -21,7 +21,7 @@ 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: RunnableLike = factory(**function.args)
@@ -41,7 +41,10 @@ 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

dao_ai/tools/search.py

@@ -10,5 +10,5 @@ def create_search_tool() -> RunnableLike:
     Returns:
         RunnableLike: A DuckDuckGo search tool that returns results as a list
     """
-    logger.debug("Creating DuckDuckGo search tool")
+    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
 
 
@@ -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: