Feat: assistant api (#423)

Co-authored-by: skyline2006 <skyline2006@163.com>
Co-authored-by: Zhicheng Zhang <zzhang.purdue@gmail.com>

Author: 3 people

docker/docker-compose.yaml

@@ -7,4 +7,4 @@ services:
       dockerfile: docker/tool_manager.dockerfile
     ports:
       - "31511:31511"
-    entrypoint: uvicorn tool_service.tool_manager.api:app --host 0.0.0.0 --port 31511
+    entrypoint: uvicorn modelscope_agent_servers.tool_manager_server.api:app --host 0.0.0.0 --port 31511

docker/tool_manager.dockerfile

@@ -24,9 +24,9 @@ RUN mkdir -p assets
 RUN mkdir -p workspace
 
 # install dependency
-ENV PYTHONPATH $PYTHONPATH:/app/tool_service
+ENV PYTHONPATH $PYTHONPATH:/app/modelscope_agent_servers
 RUN pip install fastapi pydantic uvicorn docker sqlmodel
 
-COPY tool_service /app/tool_service
+COPY modelscope_agent_servers /app/modelscope_agent_servers
 
-#ENTRYPOINT exec uvicorn tool_service.tool_manager.api:app --host 0.0.0.0 --port 31511
+#ENTRYPOINT exec uvicorn modelscope_agent_servers.tool_manager_server.api:app --host 0.0.0.0 --port 31511

docker/tool_node.dockerfile

@@ -33,11 +33,11 @@ RUN pip install --no-cache-dir -r requirements.txt
 RUN pip install fastapi uvicorn
 
 COPY modelscope_agent /app/modelscope_agent
-ENV PYTHONPATH $PYTHONPATH:/app/modelscope_agent:/app/tool_service
+ENV PYTHONPATH $PYTHONPATH:/app/modelscope_agent:/app/modelscope_agent_servers
 ENV BASE_TOOL_DIR /app/assets
 
 # install tool_node
-COPY tool_service /app/tool_service
+COPY modelscope_agent_servers /app/modelscope_agent_servers
 
 
 #ENTRYPOINT exec uvicorn tool_service.tool_node.api:app --host 0.0.0.0 --port $PORT

modelscope_agent/agent.py

@@ -18,7 +18,7 @@ def __init__(self,
                  name: Optional[str] = None,
                  description: Optional[str] = None,
                  instruction: Union[str, dict] = None,
-                 use_api: bool = False,
+                 use_tool_api: bool = False,
                  **kwargs):
         """
         init tools/llm/instruction for one agent
@@ -34,7 +34,7 @@ def __init__(self,
             name: the name of agent
             description: the description of agent, which is used for multi_agent
             instruction: the system instruction of this agent
-            use_api: whether to use the tool service api, else to use the tool cls instance
+            use_tool_api: whether to use the tool service api, else to use the tool cls instance
             kwargs: other potential parameters
         """
         if isinstance(llm, Dict):
@@ -43,7 +43,7 @@ def __init__(self,
         else:
             self.llm = llm
         self.stream = True
-        self.use_api = use_api
+        self.use_tool_api = use_tool_api
 
         self.function_list = []
         self.function_map = {}
@@ -122,12 +122,12 @@ def _register_tool(self,
             tool_class_with_tenant = TOOL_REGISTRY[tool_name]
 
             # check if the tenant_id of tool instance or tool service are exists
-            # TODO: change from use_api=True to False, to get the tenant_id of the tool changes to
+            # TODO: change from use_tool_api=True to False, to get the tenant_id of the tool changes to
             if tenant_id in tool_class_with_tenant:
                 return
 
             try:
-                if self.use_api:
+                if self.use_tool_api:
                     # get service proxy as tool instance, call method will call remote tool service
                     tool_instance = ToolServiceProxy(tool_name, tool_cfg,
                                                      tenant_id)

modelscope_agent/agents/role_play.py

@@ -5,6 +5,7 @@
 from modelscope_agent import Agent
 from modelscope_agent.agent_env_util import AgentEnvMixin
 from modelscope_agent.llm.base import BaseChatModel
+from modelscope_agent.tools.base import BaseTool
 from modelscope_agent.utils.tokenization_utils import count_tokens
 from modelscope_agent.utils.utils import check_and_limit_input_length
 
@@ -107,6 +108,11 @@
     'en': '. you can use tools: [{tool_names}]',
 }
 
+SPECIAL_PREFIX_TEMPLATE_TOOL_FOR_CHAT = {
+    'zh': '。你必须使用工具中的一个或多个：[{tool_names}]',
+    'en': '. you must use one or more tools: [{tool_names}]',
+}
+
 SPECIAL_PREFIX_TEMPLATE_KNOWLEDGE = {
     'zh': '。请查看前面的知识库',
     'en': '. Please read the knowledge base at the beginning',
@@ -146,10 +152,26 @@ def _run(self,
              lang: str = 'zh',
              **kwargs):
 
-        self.tool_descs = '\n\n'.join(tool.function_plain_text
-                                      for tool in self.function_map.values())
-        self.tool_names = ','.join(tool.name
-                                   for tool in self.function_map.values())
+        chat_mode = kwargs.get('chat_mode', False)
+        tools = kwargs.get('tools', None)
+        tool_choice = kwargs.get('tool_choice', 'auto')
+
+        if tools is not None:
+            self.tool_descs = BaseTool.parser_function(tools)
+            tool_name_list = []
+            for tool in tools:
+                func_info = tool.get('function', {})
+                if func_info == {}:
+                    continue
+                if 'name' in func_info:
+                    tool_name_list.append(func_info['name'])
+            self.tool_names = ','.join(tool_name_list)
+        else:
+            self.tool_descs = '\n\n'.join(
+                tool.function_plain_text
+                for tool in self.function_map.values())
+            self.tool_names = ','.join(tool.name
+                                       for tool in self.function_map.values())
 
         self.system_prompt = ''
         self.query_prefix = ''
@@ -172,7 +194,7 @@ def _run(self,
                 'knowledge'] = SPECIAL_PREFIX_TEMPLATE_KNOWLEDGE[lang]
 
         # concat tools information
-        if self.function_map and not self.llm.support_function_calling():
+        if self.tool_descs and not self.llm.support_function_calling():
             self.system_prompt += TOOL_TEMPLATE[lang].format(
                 tool_descs=self.tool_descs, tool_names=self.tool_names)
             self.query_prefix_dict['tool'] = SPECIAL_PREFIX_TEMPLATE_TOOL[
@@ -215,10 +237,18 @@ def _run(self,
             messages.extend(history)
 
         # concat the new messages
-        messages.append({
-            'role': 'user',
-            'content': self.query_prefix + user_request
-        })
+        if chat_mode and tool_choice == 'required':
+            required_prefix = SPECIAL_PREFIX_TEMPLATE_TOOL_FOR_CHAT[
+                lang].format(tool_names=self.tool_names)
+            messages.append({
+                'role': 'user',
+                'content': required_prefix + user_request
+            })
+        else:
+            messages.append({
+                'role': 'user',
+                'content': self.query_prefix + user_request
+            })
 
         planning_prompt = ''
         if self.llm.support_raw_prompt() and hasattr(self.llm,
@@ -265,6 +295,12 @@ def _run(self,
             else:
                 assert 'llm_result must be an instance of dict or str'
 
+            if chat_mode:
+                if use_tool and tool_choice != 'none':
+                    return f'Action: {action}\nAction Input: {action_input}\nResult: {output}'
+                else:
+                    return f'Result: {output}'
+
             # yield output
             if use_tool:
                 if self.llm.support_function_calling():

modelscope_agent/constants.py

@@ -9,6 +9,7 @@
 AGENT_REGISTRY_NAME = 'agent_center'
 TASK_CENTER_NAME = 'task_center'
 DEFAULT_TOOL_MANAGER_SERVICE_URL = 'http://localhost:31511'
+DEFAULT_ASSISTANT_SERVICE_URL = 'http://localhost:31512'
 
 
 class ApiNames(Enum):

modelscope_agent/memory/memory_with_retrieval_knowledge.py

@@ -29,8 +29,12 @@ def __init__(self,
             description=description)
 
         # allow vector storage to save knowledge
+        embedding = kwargs.get('embedding', None)
         self.store_knowledge = KnowledgeVector(
-            storage_path, name, use_cache=use_knowledge_cache)
+            storage_path,
+            name,
+            use_cache=use_knowledge_cache,
+            embedding=embedding)
 
     def _run(self,
              query: str = None,

tool_service/__init__.py modelscope_agent/rag/__init__.py

@@ -0,0 +1,50 @@
+import os
+from typing import Any, Dict, List, Union
+
+from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
+from llama_index.core.llama_pack.base import BaseLlamaPack
+from llama_index.core.readers.base import BaseReader
+
+
+class Knowledge(BaseLlamaPack):
+    """ rag pipeline.
+
+    从不同的源加载知识，支持：文件夹路径（str），文件路径列表（list），将不同源配置到不同的召回方式（dict）.
+    Automatically select the best file reader given file extensions.
+
+    Args:
+        knowledge_source: Path to the directory，或文件路径列表，或指定召回方式的文件路径。
+        cache_dir: 缓存indexing后的信息。
+    """
+
+    def __init__(self,
+                 knowledge_source: Union[List, str, Dict],
+                 cache_dir: str = './run',
+                 **kwargs) -> None:
+
+        # extra_readers = self.get_extra_readers()
+        self.documents = []
+        if isinstance(knowledge_source, str):
+            if os.path.exists(knowledge_source):
+                self.documents.append(
+                    SimpleDirectoryReader(
+                        input_dir=knowledge_source,
+                        recursive=True).load_data())
+
+        self.documents = SimpleDirectoryReader(
+            input_files=knowledge_source).load_data()
+
+    def get_extra_readers(self) -> Dict[str, BaseReader]:
+        return {}
+
+    def get_modules(self) -> Dict[str, Any]:
+        """Get modules for rewrite."""
+        return {
+            'node_parser': self.node_parser,
+            'recursive_retriever': self.recursive_retriever,
+            'query_engines': self.query_engines,
+            'reader': self.path_reader,
+        }
+
+    def run(self, query: str, **kwargs) -> str:
+        return self.query_engine.query(query, **kwargs)

modelscope_agent/rag/base.py

@@ -0,0 +1,89 @@
+import os
+from enum import Enum
+from http import HTTPStatus
+from typing import Any, List, Optional
+
+import dashscope
+from llama_index.legacy.bridge.pydantic import Field
+from llama_index.legacy.callbacks import CallbackManager
+from llama_index.legacy.core.embeddings.base import (DEFAULT_EMBED_BATCH_SIZE,
+                                                     BaseEmbedding)
+
+# Enums for validation and type safety
+DashscopeModelName = [
+    'text-embedding-v1',
+    'text-embedding-v2',
+]
+
+
+# Assuming BaseEmbedding is a Pydantic model and handles its own initializations
+class DashscopeEmbedding(BaseEmbedding):
+    """DashscopeEmbedding uses the dashscope API to generate embeddings for text."""
+
+    def __init__(
+        self,
+        model_name: str = 'text-embedding-v2',
+        embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
+        callback_manager: Optional[CallbackManager] = None,
+    ):
+        """
+        A class representation for generating embeddings using the dashscope API.
+
+        Args:
+            model_name (str): The name of the model to be used for generating embeddings. The class ensures that
+                          this model is supported and that the input type provided is compatible with the model.
+        """
+
+        assert os.environ.get(
+            'DASHSCOPE_API_KEY',
+            None), 'DASHSCOPE_API_KEY should be set in environ.'
+
+        # Validate model_name and input_type
+        if model_name not in DashscopeModelName:
+            raise ValueError(f'model {model_name} is not supported.')
+
+        super().__init__(
+            model_name=model_name,
+            embed_batch_size=embed_batch_size,
+            callback_manager=callback_manager,
+        )
+
+    @classmethod
+    def class_name(cls) -> str:
+        return 'DashscopeEmbedding'
+
+    def _embed(self,
+               texts: List[str],
+               text_type='document') -> List[List[float]]:
+        """Embed sentences using dashscope."""
+        resp = dashscope.TextEmbedding.call(
+            input=texts,
+            model=self.model_name,
+            text_type=text_type,
+        )
+        if resp.status_code == HTTPStatus.OK:
+            res = resp.output['embeddings']
+        else:
+            raise ValueError(f'call dashscope api failed: {resp}')
+
+        return [list(map(float, e['embedding'])) for e in res]
+
+    def _get_query_embedding(self, query: str) -> List[float]:
+        """Get query embedding."""
+        return self._embed([query], text_type='query')[0]
+
+    async def _aget_query_embedding(self, query: str) -> List[float]:
+        """Get query embedding async."""
+        return self._get_query_embedding(query)
+
+    def _get_text_embedding(self, text: str) -> List[float]:
+        """Get text embedding."""
+        return self._embed([text], text_type='document')[0]
+
+    async def _aget_text_embedding(self, text: str) -> List[float]:
+        """Get text embedding async."""
+        return self._get_text_embedding(text)
+
+    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Get text embeddings."""
+        return self._embed(texts, text_type='document')