casual-mcp 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

casual_mcp/mcp_tool_chat.py

@@ -14,12 +14,18 @@ from fastmcp import Client
 
 from casual_mcp.convert_tools import tools_from_mcp
 from casual_mcp.logging import get_logger
+from casual_mcp.models.chat_stats import ChatStats
+from casual_mcp.models.toolset_config import ToolSetConfig
 from casual_mcp.tool_cache import ToolCache
+from casual_mcp.tool_filter import filter_tools_by_toolset
 from casual_mcp.utils import format_tool_call_result
 
 logger = get_logger("mcp_tool_chat")
 sessions: dict[str, list[ChatMessage]] = {}
 
+# Type alias for metadata dictionary
+MetaDict = dict[str, Any]
+
 
 def get_session_messages(session_id: str) -> list[ChatMessage]:
     global sessions
@@ -44,19 +50,65 @@ class McpToolChat:
         provider: LLMProvider,
         system: str | None = None,
         tool_cache: ToolCache | None = None,
+        server_names: set[str] | None = None,
     ):
         self.provider = provider
         self.mcp_client = mcp_client
         self.system = system
         self.tool_cache = tool_cache or ToolCache(mcp_client)
+        self.server_names = server_names or set()
         self._tool_cache_version = -1
+        self._last_stats: ChatStats | None = None
 
     @staticmethod
     def get_session(session_id: str) -> list[ChatMessage] | None:
         global sessions
         return sessions.get(session_id)
 
-    async def generate(self, prompt: str, session_id: str | None = None) -> list[ChatMessage]:
+    def get_stats(self) -> ChatStats | None:
+        """
+        Get usage statistics from the last chat() or generate() call.
+
+        Returns None if no calls have been made yet.
+        Stats are reset at the start of each new chat()/generate() call.
+        """
+        return self._last_stats
+
+    def _extract_server_from_tool_name(self, tool_name: str) -> str:
+        """
+        Extract server name from a tool name.
+
+        With multiple servers, fastmcp prefixes tools as "serverName_toolName".
+        With a single server, tools are not prefixed.
+
+        Returns the server name or "default" if it cannot be determined.
+        """
+        if "_" in tool_name:
+            return tool_name.split("_", 1)[0]
+        return "default"
+
+    async def generate(
+        self,
+        prompt: str,
+        session_id: str | None = None,
+        tool_set: ToolSetConfig | None = None,
+        meta: MetaDict | None = None,
+    ) -> list[ChatMessage]:
+        """
+        Generate a response to a prompt, optionally using session history.
+
+        Args:
+            prompt: The user prompt to respond to
+            session_id: Optional session ID for conversation persistence
+            tool_set: Optional tool set configuration to filter available tools
+            meta: Optional metadata to pass through to MCP tool calls.
+                  Useful for passing context like character_id without
+                  exposing it to the LLM. Servers can access this via
+                  ctx.request_context.meta.
+
+        Returns:
+            List of response messages including any tool calls and results
+        """
         # Fetch the session if we have a session ID
         messages: list[ChatMessage]
         if session_id:
@@ -73,7 +125,7 @@ class McpToolChat:
             add_messages_to_session(session_id, [user_message])
 
         # Perform Chat
-        response = await self.chat(messages=messages)
+        response = await self.chat(messages=messages, tool_set=tool_set, meta=meta)
 
         # Add responses to session
         if session_id:
@@ -81,9 +133,36 @@ class McpToolChat:
 
         return response
 
-    async def chat(self, messages: list[ChatMessage]) -> list[ChatMessage]:
+    async def chat(
+        self,
+        messages: list[ChatMessage],
+        tool_set: ToolSetConfig | None = None,
+        meta: MetaDict | None = None,
+    ) -> list[ChatMessage]:
+        """
+        Process a conversation with tool calling support.
+
+        Args:
+            messages: The conversation messages to process
+            tool_set: Optional tool set configuration to filter available tools
+            meta: Optional metadata to pass through to MCP tool calls.
+                  Useful for passing context like character_id without
+                  exposing it to the LLM. Servers can access this via
+                  ctx.request_context.meta.
+
+        Returns:
+            List of response messages including any tool calls and results
+        """
         tools = await self.tool_cache.get_tools()
 
+        # Filter tools if a toolset is specified
+        if tool_set is not None:
+            tools = filter_tools_by_toolset(tools, tool_set, self.server_names, validate=True)
+            logger.info(f"Filtered to {len(tools)} tools using toolset")
+
+        # Reset stats at the start of each chat
+        self._last_stats = ChatStats()
+
         # Add a system message if required
         has_system_message = any(message.role == "system" for message in messages)
         if self.system and not has_system_message:
@@ -97,6 +176,15 @@ class McpToolChat:
             logger.info("Calling the LLM")
             ai_message = await self.provider.chat(messages=messages, tools=tools_from_mcp(tools))
 
+            # Accumulate token usage stats
+            self._last_stats.llm_calls += 1
+            usage = self.provider.get_usage()
+            if usage:
+                prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
+                completion_tokens = getattr(usage, "completion_tokens", 0) or 0
+                self._last_stats.tokens.prompt_tokens += prompt_tokens
+                self._last_stats.tokens.completion_tokens += completion_tokens
+
             # Add the assistant's message
             response_messages.append(ai_message)
             messages.append(ai_message)
@@ -108,14 +196,29 @@ class McpToolChat:
             logger.info(f"Executing {len(ai_message.tool_calls)} tool calls")
             result_count = 0
             for tool_call in ai_message.tool_calls:
+                # Track tool call stats
+                tool_name = tool_call.function.name
+                self._last_stats.tool_calls.by_tool[tool_name] = (
+                    self._last_stats.tool_calls.by_tool.get(tool_name, 0) + 1
+                )
+                server_name = self._extract_server_from_tool_name(tool_name)
+                self._last_stats.tool_calls.by_server[server_name] = (
+                    self._last_stats.tool_calls.by_server.get(server_name, 0) + 1
+                )
+
                 try:
-                    result = await self.execute(tool_call)
+                    result = await self.execute(tool_call, meta=meta)
                 except Exception as e:
                     logger.error(
                         f"Failed to execute tool '{tool_call.function.name}' "
                         f"(id={tool_call.id}): {e}"
                     )
-                    continue
+                    # Surface the failure to the LLM so it knows the tool failed
+                    result = ToolResultMessage(
+                        name=tool_call.function.name,
+                        tool_call_id=tool_call.id,
+                        content=f"Error executing tool: {e}",
+                    )
                 if result:
                     messages.append(result)
                     response_messages.append(result)
@@ -127,12 +230,27 @@ class McpToolChat:
 
         return response_messages
 
-    async def execute(self, tool_call: AssistantToolCall) -> ToolResultMessage:
+    async def execute(
+        self,
+        tool_call: AssistantToolCall,
+        meta: MetaDict | None = None,
+    ) -> ToolResultMessage:
+        """
+        Execute a single tool call.
+
+        Args:
+            tool_call: The tool call to execute
+            meta: Optional metadata to pass to the MCP server.
+                  Servers can access this via ctx.request_context.meta.
+
+        Returns:
+            ToolResultMessage with the tool execution result
+        """
         tool_name = tool_call.function.name
         tool_args = json.loads(tool_call.function.arguments)
         try:
             async with self.mcp_client:
-                result = await self.mcp_client.call_tool(tool_name, tool_args)
+                result = await self.mcp_client.call_tool(tool_name, tool_args, meta=meta)
         except Exception as e:
             if isinstance(e, ValueError):
                 logger.warning(e)
@@ -148,16 +266,35 @@ class McpToolChat:
         logger.debug(f"Tool Call Result: {result}")
 
         result_format = os.getenv("TOOL_RESULT_FORMAT", "result")
-        # Extract text content from result (handle both TextContent and other content types)
-        if not result.content:
+
+        # Prefer structuredContent when available (machine-readable format)
+        # Note: MCP types use camelCase (structuredContent), mypy stubs may differ
+        structured = getattr(result, "structuredContent", None)
+        if structured is not None:
+            try:
+                content_text = json.dumps(structured)
+            except (TypeError, ValueError):
+                content_text = str(structured)
+        elif not result.content:
             content_text = "[No content returned]"
         else:
-            content_item = result.content[0]
-            if hasattr(content_item, "text"):
-                content_text = content_item.text
-            else:
-                # Handle non-text content (e.g., ImageContent)
-                content_text = f"[Non-text content: {type(content_item).__name__}]"
+            # Fall back to processing content items
+            content_parts: list[Any] = []
+            for content_item in result.content:
+                if content_item.type == "text":
+                    try:
+                        parsed = json.loads(content_item.text)
+                        content_parts.append(parsed)
+                    except json.JSONDecodeError:
+                        content_parts.append(content_item.text)
+                elif hasattr(content_item, "mimeType"):
+                    # Image or audio content
+                    content_parts.append(f"[{content_item.type}: {content_item.mimeType}]")
+                else:
+                    content_parts.append(str(content_item))
+
+            content_text = json.dumps(content_parts)
+
         content = format_tool_call_result(tool_call, content_text, style=result_format)
 
         return ToolResultMessage(

casual_mcp/models/__init__.py

@@ -7,6 +7,11 @@ from casual_llm import (
     UserMessage,
 )
 
+from .chat_stats import (
+    ChatStats,
+    TokenUsageStats,
+    ToolCallStats,
+)
 from .mcp_server_config import (
     McpServerConfig,
     RemoteServerConfig,
@@ -17,6 +22,11 @@ from .model_config import (
     OllamaModelConfig,
     OpenAIModelConfig,
 )
+from .toolset_config import (
+    ExcludeSpec,
+    ToolSetConfig,
+    ToolSpec,
+)
 
 __all__ = [
     "UserMessage",
@@ -25,10 +35,16 @@ __all__ = [
     "ToolResultMessage",
     "SystemMessage",
     "ChatMessage",
+    "ChatStats",
+    "TokenUsageStats",
+    "ToolCallStats",
     "McpModelConfig",
     "OllamaModelConfig",
     "OpenAIModelConfig",
     "McpServerConfig",
     "StdioServerConfig",
     "RemoteServerConfig",
+    "ExcludeSpec",
+    "ToolSetConfig",
+    "ToolSpec",
 ]

casual_mcp/models/chat_stats.py

@@ -0,0 +1,37 @@
+"""Usage statistics models for chat sessions."""
+
+from pydantic import BaseModel, Field, computed_field
+
+
+class TokenUsageStats(BaseModel):
+    """Token usage statistics accumulated across all LLM calls."""
+
+    prompt_tokens: int = Field(default=0, ge=0)
+    completion_tokens: int = Field(default=0, ge=0)
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens (prompt + completion)."""
+        return self.prompt_tokens + self.completion_tokens
+
+
+class ToolCallStats(BaseModel):
+    """Statistics about tool calls during a chat session."""
+
+    by_tool: dict[str, int] = Field(default_factory=dict)
+    by_server: dict[str, int] = Field(default_factory=dict)
+
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def total(self) -> int:
+        """Total number of tool calls made."""
+        return sum(self.by_tool.values())
+
+
+class ChatStats(BaseModel):
+    """Combined statistics from a chat session."""
+
+    tokens: TokenUsageStats = Field(default_factory=TokenUsageStats)
+    tool_calls: ToolCallStats = Field(default_factory=ToolCallStats)
+    llm_calls: int = Field(default=0, ge=0, description="Number of LLM calls made")

casual_mcp/models/config.py

@@ -1,10 +1,12 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from casual_mcp.models.mcp_server_config import McpServerConfig
 from casual_mcp.models.model_config import McpModelConfig
+from casual_mcp.models.toolset_config import ToolSetConfig
 
 
 class Config(BaseModel):
     namespace_tools: bool | None = False
     models: dict[str, McpModelConfig]
     servers: dict[str, McpServerConfig]
+    tool_sets: dict[str, ToolSetConfig] = Field(default_factory=dict)

casual_mcp/models/toolset_config.py

@@ -0,0 +1,40 @@
+"""Toolset configuration models for filtering available tools."""
+
+from pydantic import BaseModel, Field
+
+
+class ExcludeSpec(BaseModel):
+    """Specification for excluding specific tools from a server."""
+
+    exclude: list[str] = Field(description="List of tool names to exclude")
+
+
+# Tool specification: true (all), list (include specific), or exclude object
+ToolSpec = bool | list[str] | ExcludeSpec
+
+
+class ToolSetConfig(BaseModel):
+    """Configuration for a named toolset.
+
+    A toolset defines which tools from which servers should be available
+    during a chat session. Each server can be configured to:
+    - Include all tools (True)
+    - Include specific tools (list of tool names)
+    - Include all except specific tools (ExcludeSpec)
+
+    Example:
+        {
+            "description": "Research tools",
+            "servers": {
+                "wikimedia": True,
+                "search": ["brave_web_search"],
+                "fetch": {"exclude": ["fetch_dangerous"]}
+            }
+        }
+    """
+
+    description: str = Field(default="", description="Human-readable description")
+    servers: dict[str, ToolSpec] = Field(
+        default_factory=dict,
+        description="Mapping of server name to tool specification",
+    )

casual_mcp/tool_filter.py

@@ -0,0 +1,171 @@
+"""Tool filtering logic for toolsets.
+
+This module provides functionality to filter MCP tools based on toolset
+configurations, including validation to ensure referenced servers and
+tools actually exist.
+"""
+
+import mcp
+
+from casual_mcp.logging import get_logger
+from casual_mcp.models.toolset_config import ExcludeSpec, ToolSetConfig
+
+logger = get_logger("tool_filter")
+
+
+class ToolSetValidationError(Exception):
+    """Raised when a toolset references invalid servers or tools."""
+
+    pass
+
+
+def extract_server_and_tool(tool_name: str, server_names: set[str]) -> tuple[str, str]:
+    """Extract server name and base tool name from a potentially prefixed tool name.
+
+    When multiple servers are configured, fastmcp prefixes tools as "serverName_toolName".
+    When a single server is configured, tools are not prefixed.
+
+    Args:
+        tool_name: The full tool name (possibly prefixed)
+        server_names: Set of configured server names
+
+    Returns:
+        Tuple of (server_name, base_tool_name)
+    """
+    if "_" in tool_name:
+        prefix = tool_name.split("_", 1)[0]
+        if prefix in server_names:
+            return prefix, tool_name.split("_", 1)[1]
+
+    # Single server case - return the single server name
+    if len(server_names) == 1:
+        return next(iter(server_names)), tool_name
+
+    # Fallback - can't determine server
+    return "default", tool_name
+
+
+def _build_server_tool_map(tools: list[mcp.Tool], server_names: set[str]) -> dict[str, set[str]]:
+    """Build a mapping of server names to their available tool names.
+
+    Args:
+        tools: List of MCP tools
+        server_names: Set of configured server names
+
+    Returns:
+        Dict mapping server name to set of base tool names
+    """
+    server_tool_map: dict[str, set[str]] = {name: set() for name in server_names}
+
+    for tool in tools:
+        server_name, base_name = extract_server_and_tool(tool.name, server_names)
+        if server_name in server_tool_map:
+            server_tool_map[server_name].add(base_name)
+
+    return server_tool_map
+
+
+def validate_toolset(
+    toolset: ToolSetConfig,
+    tools: list[mcp.Tool],
+    server_names: set[str],
+) -> None:
+    """Validate that a toolset references only valid servers and tools.
+
+    Args:
+        toolset: The toolset configuration to validate
+        tools: List of available MCP tools
+        server_names: Set of configured server names
+
+    Raises:
+        ToolSetValidationError: If the toolset references non-existent servers or tools
+    """
+    server_tool_map = _build_server_tool_map(tools, server_names)
+    errors: list[str] = []
+
+    for server_name, tool_spec in toolset.servers.items():
+        # Check server exists
+        if server_name not in server_names:
+            errors.append(f"Server '{server_name}' not found in configuration")
+            continue
+
+        available = server_tool_map.get(server_name, set())
+
+        # Validate tool names in include list
+        if isinstance(tool_spec, list):
+            for tool_name in tool_spec:
+                if tool_name not in available:
+                    errors.append(
+                        f"Tool '{tool_name}' not found in server '{server_name}'. "
+                        f"Available: {sorted(available)}"
+                    )
+
+        # Validate tool names in exclude list
+        elif isinstance(tool_spec, ExcludeSpec):
+            for tool_name in tool_spec.exclude:
+                if tool_name not in available:
+                    errors.append(
+                        f"Tool '{tool_name}' not found in server '{server_name}' "
+                        f"(specified in exclude list). Available: {sorted(available)}"
+                    )
+
+    if errors:
+        raise ToolSetValidationError("\n".join(errors))
+
+
+def filter_tools_by_toolset(
+    tools: list[mcp.Tool],
+    toolset: ToolSetConfig,
+    server_names: set[str],
+    validate: bool = True,
+) -> list[mcp.Tool]:
+    """Filter a list of MCP tools based on a toolset configuration.
+
+    Args:
+        tools: Full list of available MCP tools
+        toolset: The toolset configuration to apply
+        server_names: Set of configured server names
+        validate: Whether to validate the toolset first (raises on invalid)
+
+    Returns:
+        Filtered list of tools matching the toolset
+
+    Raises:
+        ToolSetValidationError: If validate=True and toolset is invalid
+    """
+    if validate:
+        validate_toolset(toolset, tools, server_names)
+
+    filtered: list[mcp.Tool] = []
+
+    for tool in tools:
+        server_name, base_name = extract_server_and_tool(tool.name, server_names)
+
+        # Check if this server is in the toolset
+        if server_name not in toolset.servers:
+            continue
+
+        tool_spec = toolset.servers[server_name]
+
+        # Determine if tool should be included
+        include = False
+
+        if tool_spec is True:
+            # All tools from this server
+            include = True
+        elif isinstance(tool_spec, list):
+            # Only specific tools
+            include = base_name in tool_spec
+        elif isinstance(tool_spec, ExcludeSpec):
+            # All except excluded tools
+            include = base_name not in tool_spec.exclude
+
+        if include:
+            filtered.append(tool)
+
+    logger.debug(
+        f"Filtered {len(tools)} tools to {len(filtered)} using toolset "
+        f"with {len(toolset.servers)} servers"
+    )
+
+    return filtered