wehub-resource-sync
68 lines
2.9 KiB
Python
68 lines
2.9 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
import json
|
|
import re
|
|
from typing import List, Optional, Tuple, Union
|
|
|
|
from swift.infer_engine import Function
|
|
from swift.template import Prompt
|
|
from .base import BaseAgentTemplate
|
|
|
|
|
|
class MistralAgentTemplate(BaseAgentTemplate):
|
|
|
|
def get_toolcall(self, response: str) -> List[Function]:
|
|
res_list = re.findall(r'\[TOOL_CALLS\]\[(.*?)\]</s>', response, re.DOTALL)
|
|
if not res_list:
|
|
return []
|
|
res_list = res_list[0].strip().split('\n')
|
|
functions = []
|
|
for res_str in res_list:
|
|
parsed_res = self._parse_json(res_str)
|
|
if isinstance(parsed_res, dict):
|
|
parsed_res = [parsed_res] # Handle single tool call
|
|
if isinstance(parsed_res, list):
|
|
for tool_call in parsed_res:
|
|
if isinstance(tool_call, dict) and 'name' in tool_call and 'arguments' in tool_call:
|
|
functions.append(Function(name=tool_call['name'], arguments=tool_call['arguments']))
|
|
if len(functions) == 0:
|
|
# compat react_en
|
|
return super().get_toolcall(response)
|
|
return functions
|
|
|
|
def _format_tool_responses(
|
|
self,
|
|
assistant_content: str,
|
|
tool_messages,
|
|
) -> Tuple[str, 'Prompt']:
|
|
if not hasattr(self, 'template_meta'):
|
|
raise ValueError('MistralAgentTemplate requires template_meta to be registered')
|
|
prompt = self.template_meta.prompt
|
|
chat_sep = self.template_meta.chat_sep
|
|
|
|
res = chat_sep.copy()
|
|
res_tool = []
|
|
for tool_message in tool_messages:
|
|
tool_content = tool_message['content']
|
|
# append `[TOOL_RESULTS]{"content": {{ .Content }}}[/TOOL_RESULTS]` to res_tool
|
|
res_tool.append(f'[TOOL_RESULTS]{json.dumps({"content": tool_content}, ensure_ascii=False)}[/TOOL_RESULTS]')
|
|
total_tool = '\n'.join(res_tool)
|
|
for context in prompt:
|
|
if isinstance(context, str):
|
|
context = context.replace('{{QUERY}}', total_tool)
|
|
res.append(context)
|
|
return assistant_content, res
|
|
|
|
def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
|
|
tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools]
|
|
system = system or ''
|
|
return f"""{system}[AVAILABLE_TOOLS]{' '.join(tool_descs)}[/AVAILABLE_TOOLS]"""
|
|
|
|
def _format_tool_calls(self, tool_call_messages):
|
|
tool_calls = []
|
|
for message in tool_call_messages:
|
|
# needs `{'name': name, 'arguments': arguments}`, which self._parse_tool_call
|
|
# satisfies
|
|
tool_call = self._parse_tool_call(message['content'])
|
|
tool_calls.append(json.dumps(tool_call, ensure_ascii=False))
|
|
return f'[TOOL_CALLS][\n{chr(10).join(tool_calls)}\n]</s>' # check if need `</s>` at end
|