ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/examples/gaia/mcp_collections/base.py

@@ -0,0 +1,102 @@
+import logging
+import os
+from pathlib import Path
+from typing import Any, Literal
+
+from mcp.server import FastMCP
+from pydantic import BaseModel, Field
+
+from aworld.logs.util import Color
+from examples.gaia.utils import color_log, setup_logger
+
+
+class ActionArguments(BaseModel):
+    r"""Protocol: MCP Action Arguments"""
+
+    name: str = Field(description="The name of the action")
+    transport: Literal["stdio", "sse"] = Field(default="stdio", description="The transport of the action")
+    workspace: str | None = Field(
+        default=None,
+        description="The workspace of the action."
+        " If not specified or invalid, the workspace will be read from the environment variable AWORLD_WORKSPACE.",
+    )
+    unittest: bool = Field(default=False, description="Whether to run in unittest mode")
+
+
+class ActionResponse(BaseModel):
+    r"""Protocol: MCP Action Response"""
+
+    success: bool = Field(default=False, description="Whether the action is successfully executed")
+    message: Any = Field(default=None, description="The execution result of the action")
+    metadata: dict[str, Any] = Field(default={}, description="The metadata of the action")
+
+
+class ActionCollection:
+    r"""Base class for all ActionCollection."""
+
+    server: FastMCP
+    logger: logging.Logger
+
+    def __init__(self, arguments: ActionArguments) -> None:
+        self.unittest = arguments.unittest
+        self.transport = arguments.transport
+        self.supported_extensions = set()
+
+        self.workspace: Path = self._obtain_valid_workspace(arguments.workspace)
+
+        self.logger: logging.Logger = setup_logger(self.__class__.__name__, self.workspace)
+
+        self.server = FastMCP(arguments.name)
+        for tool_name in self.__class__.__dict__:
+            if tool_name.startswith("mcp_") and callable(getattr(self.__class__, tool_name)):
+                tool = getattr(self, tool_name)
+                self.server.add_tool(tool, description=tool.__doc__)
+
+    def run(self) -> None:
+        if not self.unittest:
+            self.server.run(transport=self.transport)
+
+    def _color_log(self, value: str, color: Color = None, level: str = "info"):
+        return color_log(self.logger, value, color, level=level)
+
+    def _obtain_valid_workspace(self, workspace: str | None = None) -> Path:
+        r"""
+        Obtain a valid workspace path.
+        Priority:
+          1. user defined workspace
+          2. environment variable AWORLD_WORKSPACE
+          3. home directory
+        """
+        path = Path(workspace) if workspace else os.getenv("AWORLD_WORKSPACE", "~")
+        if path and path.expanduser().is_dir():
+            return path.expanduser().resolve()
+
+        # self._color_log("Invalid workspace path, using home directory instead.", Color.yellow)
+        return Path.home().expanduser().resolve()
+
+    def _validate_file_path(self, file_path: str) -> Path:
+        """Validate and resolve file path. Rely on the predefined supported_extensions class variable.
+
+        Args:
+            file_path: Path to the document or media file
+
+        Returns:
+            Resolved Path object
+
+        Raises:
+            FileNotFoundError: If file doesn't exist
+            ValueError: If file type is not supported
+        """
+        path = Path(file_path).expanduser()
+        if not path.is_absolute():
+            path = self.workspace / path
+
+        if not path.exists():
+            raise FileNotFoundError(f"File not found: {path}")
+
+        if path.suffix.lower() not in self.supported_extensions:
+            raise ValueError(
+                f"Unsupported file type: {path.suffix}. Supported types: {', '.join(self.supported_extensions)}"
+            )
+
+        return path