wip

Author: tastelikefeet

modelscope_agent/cli/arg_parser.py

@@ -2,15 +2,13 @@
 import asyncio
 import os
 
-from modelscope_agent.cli.arg_parser import parse_args
 from modelscope_agent.config import Config
 from modelscope_agent.engine.simple_engine import SimpleEngine
 
 if __name__ == '__main__':
     cur_file = __file__
     cur_dir = os.path.dirname(cur_file)
     config = Config.from_task(os.path.join(cur_dir, 'search.yaml'))
-    parse_args(config)
     engine = SimpleEngine(config=config)
     query = input('>>>Please input the query')
     asyncio.run(engine.run(query))

modelscope_agent/cli/search.py

@@ -18,23 +18,27 @@ servers:
   crawl4ai:
     command: "uv"
     args:
+      - run
+      - <crawl4ai_script>
 
   web-search:
     command: "npx"
     args:
       - "-y"
       - "tavily-mcp@0.1.4"
     env:
-      - TAVILY_API_KEY:
+      TAVILY_API_KEY:
     disabled: false
-    autoApprove:
+    autoApprove: []
     exclude:
       - tavily-extract
 
   edgeone-pages-mcp-server:
     command: "npx"
     args:
       - "edgeone-pages-mcp"
+    exclude:
+      - deploy_folder_or_zip
 
 help: |
   This search agent uses the following mcp servers:
@@ -48,9 +52,9 @@ help: |
     https://app.tavily.com/home
   3. Add `modelscope_api_key` to your env variable, you can create one from:
     https://www.modelscope.cn/my/myaccesstoken
-  4. Run `modelscope_api_key=xxx TAVILY_API_KEY=xxx modelscope-agent search --servers.crawl4ai.args /your/path/to/mcp-central/mcp_central/crawl4ai/server.py`
+  4. Run `modelscope_api_key=xxx TAVILY_API_KEY=xxx modelscope-agent search --crawl4ai_script /your/path/to/mcp-central/mcp_central/crawl4ai/server.py`
   5. If you want to change models, you can add more arguments:
-    - --llm.model xxx 
+    - --model xxx 
   
 
 

modelscope_agent/cli/search.yaml

@@ -1,4 +1,6 @@
+import argparse
 import os.path
+from copy import deepcopy
 from typing import Dict, Union, Any
 
 from modelscope import snapshot_download
@@ -42,13 +44,28 @@ def from_task(cls, task_dir_or_id: str, env: Dict[str, str] = None) -> Union[Dic
         assert config is not None, (f'Cannot find any config file in {task_dir_or_id} named `config.json`, '
                                     f'`config.yml` or `config.yaml`')
         envs = Env.load_env(env)
-        cls._update_envs(config, envs)
+        cls._update_config(config, envs)
+        _dict_config = cls.parse_args()
+        cls._update_config(config, _dict_config)
         config.local_dir = task_dir_or_id
         return config
 
     @staticmethod
-    def _update_envs(config: Union[DictConfig, ListConfig], envs: Dict[str, str]=None):
-        if not envs:
+    def parse_args():
+        arg_parser = argparse.ArgumentParser()
+        args, unknown = arg_parser.parse_known_args()
+        _dict_config = {}
+        if unknown:
+            for idx in range(0, len(unknown), 2):
+                key = unknown[idx]
+                value = unknown[idx + 1]
+                assert key.startswith('--'), f'Parameter not correct: {unknown}'
+                _dict_config[key[2:]] = value
+        return _dict_config
+
+    @staticmethod
+    def _update_config(config: Union[DictConfig, ListConfig], extra: Dict[str, str]=None):
+        if not extra:
             return config
 
         def traverse_config(_config: Union[DictConfig, ListConfig, Any]):
@@ -57,13 +74,25 @@ def traverse_config(_config: Union[DictConfig, ListConfig, Any]):
                     if isinstance(value, BaseContainer):
                         traverse_config(value)
                     else:
-                        if name in envs:
-                            logger.info(f'Replacing {name} with the value in your environment variables.')
-                            setattr(_config, name, envs[name])
+                        if name in extra:
+                            logger.info(f'Replacing {name} with extra value.')
+                            setattr(_config, name, extra[name])
+                        if (isinstance(value, str) and value.startswith('<') and
+                                value.endswith('>') and value[1:-1] in extra):
+                            logger.info(f'Replacing {value} with extra value.')
+                            setattr(_config, name, extra[name])
+
             elif isinstance(_config, ListConfig):
-                for value in _config:
+                for idx in range(len(_config)):
+                    value = _config[idx]
                     if isinstance(value, BaseContainer):
                         traverse_config(value)
+                    else:
+                        if (isinstance(value, str) and value.startswith('<') and
+                                value.endswith('>') and value[1:-1] in extra):
+                            logger.info(f'Replacing {value} with extra value.')
+                            _config[idx] = extra[value[1:-1]]
+
         traverse_config(config)
         return None
 
@@ -77,5 +106,5 @@ def convert_mcp_servers_to_json(config: Union[DictConfig, ListConfig]):
         }
         if config.servers:
             for server, server_config in config.servers.items():
-                servers['mcpServers'][server] = server_config.to_json()
+                servers['mcpServers'][server] = deepcopy(server_config)
         return servers

modelscope_agent/config/config.py

@@ -74,25 +74,26 @@ def register_callback(self, callback: Callback):
         self.callbacks.append(callback)
 
     def _register_callback_from_config(self):
-        local_dir = self.config.local_dir
-        for _callback in self.config.callbacks:
-            if _callback.endswith('.py'):
-                if not self.trust_remote_code:
-                    raise AssertionError(f'Your config file contains external code, '
-                                         f'instantiate the code may be UNSAFE, if you trust the code, '
-                                         f'please pass `trust_remote_code=True` or `--trust_remote_code true`')
-                if sys.path[0] != local_dir:
-                    assert local_dir is not None, 'Using external py files, but local_dir cannot be found.'
-                    sys.path.insert(0, local_dir)
-                callback_file = importlib.import_module(_callback[:-3])
-                module_classes = {name: cls for name, cls in inspect.getmembers(callback_file, inspect.isclass)}
-                for name, cls in module_classes.items():
-                    # Find cls which base class is `Callback`
-                    if cls.__base__[0] is Callback:
-                        self.callbacks.append(cls())
-            else:
-                assert _callback in callbacks_mapping
-                self.callbacks.append(callbacks_mapping[_callback]())
+        local_dir = self.config.local_dir if hasattr(self.config, 'local_dir') else None
+        if hasattr(self.config, 'callbacks'):
+            for _callback in self.config.callbacks:
+                if _callback.endswith('.py'):
+                    if not self.trust_remote_code:
+                        raise AssertionError(f'Your config file contains external code, '
+                                             f'instantiate the code may be UNSAFE, if you trust the code, '
+                                             f'please pass `trust_remote_code=True` or `--trust_remote_code true`')
+                    if sys.path[0] != local_dir:
+                        assert local_dir is not None, 'Using external py files, but local_dir cannot be found.'
+                        sys.path.insert(0, local_dir)
+                    callback_file = importlib.import_module(_callback[:-3])
+                    module_classes = {name: cls for name, cls in inspect.getmembers(callback_file, inspect.isclass)}
+                    for name, cls in module_classes.items():
+                        # Find cls which base class is `Callback`
+                        if cls.__base__[0] is Callback:
+                            self.callbacks.append(cls())
+                else:
+                    assert _callback in callbacks_mapping
+                    self.callbacks.append(callbacks_mapping[_callback]())
 
     def _loop_callback(self, point, messages: List[Message]):
         for callback in self.callbacks:
@@ -120,18 +121,18 @@ def _prepare_messages(self, prompt):
             {'role': 'system', 'content': self.config.prompt.system or self.DEFAULT_SYSTEM_EN},
             {'role': 'user', 'content': prompt or self.config.prompt.query},
         ]
-        messages['query'] = self._query_documents(messages[1]['content'])
+        messages[1]['content'] = self._query_documents(messages[1]['content'])
         return messages
 
     def _prepare_memory(self):
-        if self.config.memory:
+        if hasattr(self.config, 'memory') and self.config.memory:
             for _memory in self.config.memory:
                 assert _memory in memory_mapping, (f'{_memory} not in memory_mapping, '
                                                    f'which supports: {list(memory_mapping.keys())}')
                 self.memory_tools.append(memory_mapping[_memory]())
 
     def _prepare_rag(self):
-        if self.config.rag:
+        if hasattr(self.config, 'rag') and self.config.rag:
             assert self.config.rag in rag_mapping
             self.rag: Rag = rag_mapping(self.config.rag)()
 

modelscope_agent/engine/simple_engine.py

@@ -1,15 +1,13 @@
 import inspect
-
 from typing import Any, List, Dict, Optional, Generator
 
-from grpc.framework.interfaces.base.utilities import completion
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
+from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 
-from modelscope_agent.utils.utils import assert_package_exist
 from modelscope_agent.llm.llm import LLM
 from modelscope_agent.llm.utils import Message, Tool, ToolCall
+from modelscope_agent.utils.utils import assert_package_exist
 
-from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall, Function
 
 def _stream_generator() -> Generator[Message, None, None]:
     pass

modelscope_agent/llm/openai_llm.py

@@ -1,6 +1,7 @@
 from omegaconf import DictConfig
 
 from modelscope_agent.tools.base import ToolBase
+from modelscope_agent.tools.utils import ToolField
 
 
 class LoopTool(ToolBase):
@@ -16,13 +17,15 @@ async def cleanup(self):
 
     async def get_tools(self):
         return {
-            'split_complex_task': {
-                'tool_name': 'split_to_sub_task',
-                'tool_args': {
-                    'system': 'The system prompt of this sub task',
-                    'query': 'The query to solve of this sub task',
-                }
-            }
+            'split_complex_task': [ToolField(
+                name='split_to_sub_task',
+                description='Split complex task into sub tasks, for example, split a website generation task into sub tasks, '
+                               'you plan the framework, include code files and classes and functions, and give the detail '
+                               'information to the system and query field of the subtask, then '
+                               'let each subtask to write a single file',
+                input_schema='system: str, The system prompt to use in this sub task. \n'
+                                'query: str, The specific query in this sub task. \n',
+        )]
         }
 
 

modelscope_agent/tools/loop_tool.py

@@ -1,10 +1,11 @@
 from contextlib import AsyncExitStack
-from typing import Any, Dict, Literal, Optional
+from typing import Any, Dict, Literal, Optional, List
 
-from mcp import ClientSession, StdioServerParameters
+from mcp import ClientSession, StdioServerParameters, ListToolsResult
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
 
+from modelscope_agent.config import Config
 from modelscope_agent.config.env import Env
 from modelscope_agent.tools.base import ToolBase
 from modelscope_agent.utils import get_logger
@@ -26,7 +27,10 @@ def __init__(self, config, mcp_config: Optional[Dict[str, Any]] = None):
         super().__init__(config)
         self.sessions: Dict[str, ClientSession] = {}
         self.exit_stack = AsyncExitStack()
-        self.mcp_config = mcp_config
+        self.mcp_config = Config.convert_mcp_servers_to_json(config)
+        self._exclude_functions = {}
+        if mcp_config is not None:
+            self.mcp_config.update(mcp_config)
 
     async def call_tool(self, server_name: str, tool_name: str,
                         tool_args: dict):
@@ -45,11 +49,17 @@ async def get_tools(self) -> Dict:
         for key, session in self.sessions.items():
             tools[key] = []
             response = await session.list_tools()
-            tools[key].extend(response.tools)
+            _session_tools = response.tools
+            exclude = []
+            if key in self._exclude_functions:
+                exclude = self._exclude_functions[key]
+            _session_tools = [t for t in _session_tools if t.name not in exclude]
+            tools[key].extend(_session_tools)
         return tools
 
     @staticmethod
-    def print_tools(server_name: str, tools: Dict):
+    def print_tools(server_name: str, tools: ListToolsResult):
+        tools = tools.tools
         if len(tools) > 10:
             tools = [tool.name for tool in tools][:10]
             logger.info(f'\nConnected to server "{server_name}" '
@@ -116,14 +126,16 @@ async def connect_to_server(self, server_name: str, **kwargs):
     async def connect(self):
         assert self.mcp_config, 'MCP config is required'
         envs = Env.load_env()
-        for name, server in self.mcp_config.items():
-            cmd = server['cmd']
-            env_dict = cmd.pop('env', {})
+        mcp_config = self.mcp_config['mcpServers']
+        for name, server in mcp_config.items():
+            env_dict = server.pop('env', {})
             env_dict = {
                 key: value if value else envs.get(key, '')
                 for key, value in env_dict.items()
             }
-            await self.connect_to_server(server_name=name, env=env_dict, **cmd)
+            if 'exclude' in server:
+                self._exclude_functions[name] = server.pop('exclude')
+            await self.connect_to_server(server_name=name, env=env_dict, **server)
 
     async def cleanup(self):
         """Clean up resources"""

modelscope_agent/tools/mcp_client.py

@@ -14,13 +14,13 @@ def __init__(self, config):
         self.loop_tool = LoopTool(config)
         self.extra_tools: List[ToolBase] = []
         self._tool_index = {}
-        asyncio.run(self.reindex_tool())
 
     async def connect(self):
         await self.mcp_client.connect()
         await self.loop_tool.connect()
         for tool in self.extra_tools:
             await tool.connect()
+        asyncio.run(self.reindex_tool())
 
     async def cleanup(self):
         await self.mcp_client.cleanup()
@@ -32,8 +32,8 @@ async def reindex_tool(self):
 
         def extend_tool(tool_ins: ToolBase, server_name: str, tool_list: List):
             for tool in tool_list:
-                assert tool['name'] not in self._tool_index, f'Tool name duplicated {tool["name"]}'
-                self._tool_index[tool['name']] = (tool_ins, server_name, tool)
+                assert tool.name not in self._tool_index, f'Tool name duplicated {tool.name}'
+                self._tool_index[tool.name] = (tool_ins, server_name, tool)
 
         mcps = await self.mcp_client.get_tools()
         for server_name, tool_list in mcps.items():

modelscope_agent/tools/tool_manager.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class ToolField:
+
+    name: str = None
+    description: str = None
+    input_schema: str = None