Optimize DataJuicer Agent doc & linter (#30)

data_juicer_agent/tools/mcp_helpers.py

@@ -1,11 +1,16 @@
+# -*- coding: utf-8 -*-
 import json
 import os
 import logging
-from typing import Optional, List
+from typing import Optional
 import string
 
 from agentscope.tool import Toolkit
-from agentscope.mcp import HttpStatefulClient, HttpStatelessClient, StdIOStatefulClient
+from agentscope.mcp import (
+    HttpStatefulClient,
+    HttpStatelessClient,
+    StdIOStatefulClient,
+)
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -13,6 +18,7 @@ logger = logging.getLogger(__name__)
 
 root_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
 
+
 def _load_config(config_path: str) -> dict:
     """Load MCP configuration from file"""
     try:
@@ -23,13 +29,15 @@ def _load_config(config_path: str) -> dict:
             return config
         else:
             logger.warning(
-                f"Configuration file {config_path} not found, using default settings"
+                f"Configuration file {config_path} not found, "
+                "using default settings",
             )
             return _create_default_config()
     except Exception as e:
         logger.error(f"Error loading configuration: {e}")
         return _create_default_config()
 
+
 def _create_default_config() -> dict:
     """Create default configuration"""
     return {
@@ -38,10 +46,11 @@ def _create_default_config() -> dict:
                 "command": "python",
                 "args": ["/home/test/data_juicer/tools/DJ_mcp_recipe_flow.py"],
                 "env": {"SERVER_TRANSPORT": "stdio"},
-            }
-        }
+            },
+        },
     }
 
+
 def _expand_env_vars(value: str) -> str:
     """Expand environment variables in configuration values"""
     if isinstance(value, str):
@@ -53,6 +62,7 @@ def _expand_env_vars(value: str) -> str:
             return value
     return value
 
+
 async def _create_clients(config: dict, toolkit: Toolkit):
     """Create MCP clients based on configuration"""
     server_configs = config.get("mcpServers", {})
@@ -88,33 +98,38 @@ async def _create_clients(config: dict, toolkit: Toolkit):
 
                 if stateful:
                     client = HttpStatefulClient(
-                        name=server_name, transport=transport, url=url
+                        name=server_name,
+                        transport=transport,
+                        url=url,
                     )
                     await client.connect()
                     await toolkit.register_mcp_client(client)
                 else:
                     client = HttpStatelessClient(
-                        name=server_name, transport=transport, url=url
+                        name=server_name,
+                        transport=transport,
+                        url=url,
                     )
                     await toolkit.register_mcp_client(client)
 
             else:
                 raise ValueError("Invalid server configuration")
-            
+
             clients.append(client)
         except Exception as e:
             if "Invalid server configuration" in str(e):
                 raise e
             logger.error(f"Failed to create client {server_name}: {e}")
-        
+
     return clients
 
+
 async def get_mcp_toolkit(config_path: Optional[str] = None) -> Toolkit:
     """Get toolkit with all MCP tools registered"""
     config_path = config_path or root_path + "/configs/mcp_config.json"
     config = _load_config(config_path)
     toolkit = Toolkit()
-    
+
     clients = await _create_clients(config, toolkit)
-    
+
     return toolkit, clients