wehub-resource-sync
139 lines
5.2 KiB
Python
139 lines
5.2 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||
import json
|
||
import re
|
||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||
|
||
from swift.infer_engine import Function
|
||
from swift.template import Prompt
|
||
from .base import BaseAgentTemplate
|
||
|
||
DSML_TOKEN = '｜DSML｜'
|
||
|
||
TOOLS_TEMPLATE = """## Tools
|
||
|
||
You have access to a set of tools to help answer the user's question. \
|
||
You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following:
|
||
|
||
<{dsml_token}tool_calls>
|
||
<{dsml_token}invoke name="$TOOL_NAME">
|
||
<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
|
||
...
|
||
</{dsml_token}invoke>
|
||
<{dsml_token}invoke name="$TOOL_NAME2">
|
||
...
|
||
</{dsml_token}invoke>
|
||
</{dsml_token}tool_calls>
|
||
|
||
String parameters should be specified as is and set `string="true"`. \
|
||
For all other types (numbers, booleans, arrays, objects), \
|
||
pass the value in JSON format and set `string="false"`.
|
||
|
||
If thinking_mode is enabled (triggered by <think>), \
|
||
you MUST output your complete reasoning inside <think>...</think> BEFORE any tool calls or final response.
|
||
|
||
Otherwise, output directly after </think> with tool calls or final response.
|
||
|
||
### Available Tool Schemas
|
||
|
||
{tool_schemas}
|
||
|
||
You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.
|
||
"""
|
||
|
||
|
||
def _to_json(value: Any) -> str:
|
||
try:
|
||
return json.dumps(value, ensure_ascii=False)
|
||
except Exception:
|
||
return json.dumps(value, ensure_ascii=True)
|
||
|
||
|
||
def _encode_arguments_to_dsml(arguments: Dict[str, Any]) -> str:
|
||
"""Encode tool call arguments dict into DSML parameter lines."""
|
||
lines = []
|
||
for k, v in arguments.items():
|
||
is_str = 'true' if isinstance(v, str) else 'false'
|
||
val = v if isinstance(v, str) else _to_json(v)
|
||
lines.append(f'<{DSML_TOKEN}parameter name="{k}" string="{is_str}">{val}</{DSML_TOKEN}parameter>')
|
||
return '\n'.join(lines)
|
||
|
||
|
||
class DeepSeekV4AgentTemplate(BaseAgentTemplate):
|
||
|
||
def get_toolcall(self, response: str) -> List[Function]:
|
||
# Parse DSML tool calls from model output
|
||
# Pattern: <｜DSML｜invoke name="tool_name">...params...</｜DSML｜invoke>
|
||
invoke_pattern = re.compile(
|
||
rf'<{re.escape(DSML_TOKEN)}invoke\s+name="([^"]+)">\s*(.*?)\s*</{re.escape(DSML_TOKEN)}invoke>', re.DOTALL)
|
||
param_pattern = re.compile(
|
||
rf'<{re.escape(DSML_TOKEN)}parameter\s+name="([^"]+)"\s+string="(true|false)">'
|
||
rf'(.*?)</{re.escape(DSML_TOKEN)}parameter>', re.DOTALL)
|
||
|
||
functions = []
|
||
for match in invoke_pattern.finditer(response):
|
||
tool_name = match.group(1)
|
||
params_block = match.group(2)
|
||
arguments = {}
|
||
for pm in param_pattern.finditer(params_block):
|
||
param_name = pm.group(1)
|
||
is_string = pm.group(2)
|
||
param_value = pm.group(3)
|
||
if is_string == 'false':
|
||
try:
|
||
param_value = json.loads(param_value)
|
||
except json.JSONDecodeError:
|
||
pass
|
||
arguments[param_name] = param_value
|
||
functions.append(Function(name=tool_name, arguments=json.dumps(arguments, ensure_ascii=False)))
|
||
|
||
if len(functions) == 0:
|
||
# Fallback to ReAct format
|
||
return super().get_toolcall(response)
|
||
return functions
|
||
|
||
def _get_tool_responses(self, tool_messages):
|
||
return ''.join(f'<tool_result>{tool_message["content"]}</tool_result>' for tool_message in tool_messages)
|
||
|
||
def _format_tool_responses(
|
||
self,
|
||
assistant_content: str,
|
||
tool_messages,
|
||
) -> Tuple[str, 'Prompt']:
|
||
with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content
|
||
if with_action:
|
||
return super()._format_tool_responses(assistant_content, tool_messages)
|
||
res = [
|
||
'<｜end▁of▁sentence｜><｜User｜>',
|
||
self._get_tool_responses(tool_messages),
|
||
'<｜Assistant｜>',
|
||
]
|
||
return assistant_content, res
|
||
|
||
def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str:
|
||
tool_schemas = []
|
||
for tool in tools:
|
||
tool = self.unwrap_tool(tool)
|
||
tool_schemas.append(_to_json(tool))
|
||
|
||
tools_section = TOOLS_TEMPLATE.format(
|
||
tool_schemas='\n'.join(tool_schemas),
|
||
dsml_token=DSML_TOKEN,
|
||
)
|
||
|
||
system = system or ''
|
||
return f'{system}\n\n{tools_section}' if system else tools_section
|
||
|
||
def _format_tool_calls(self, tool_call_messages) -> str:
|
||
invocations = []
|
||
for message in tool_call_messages:
|
||
tool_call = self._parse_tool_call(message['content'])
|
||
name = tool_call['name']
|
||
arguments = tool_call['arguments']
|
||
if isinstance(arguments, str):
|
||
arguments = json.loads(arguments)
|
||
dsml_args = _encode_arguments_to_dsml(arguments)
|
||
invocations.append(f'<{DSML_TOKEN}invoke name="{name}">\n{dsml_args}\n</{DSML_TOKEN}invoke>')
|
||
|
||
tool_calls_str = '\n'.join(invocations)
|
||
return f'<{DSML_TOKEN}tool_calls>\n{tool_calls_str}\n</{DSML_TOKEN}tool_calls>'
|