diff --git a/.gitignore b/.gitignore index b3581ef..720c91e 100644 --- a/.gitignore +++ b/.gitignore @@ -2,5 +2,6 @@ __pycache__/ *.py[cod] outputs/_test_*/ outputs/tmp*/ +outputs/* models/* !models/.gitkeep diff --git a/code/b4_local_agent_llm.py b/code/b4_local_agent_llm.py index ce6d2dd..94eb4a9 100644 --- a/code/b4_local_agent_llm.py +++ b/code/b4_local_agent_llm.py @@ -256,7 +256,7 @@ def _load_model_bundle( return tokenizer, model -def _build_prompt_messages(messages: list[dict], tools_schema: list[dict]) -> list[dict]: +def _build_prompt_messages(messages: list[dict], tools_schema: list[dict],include_tools_schema: bool = True) -> list[dict]: prompt_messages = deepcopy(messages) format_instruction = ( "IMPORTANT OUTPUT FORMAT:\n" @@ -284,12 +284,17 @@ def _build_prompt_messages(messages: list[dict], tools_schema: list[dict]) -> li "Choose exactly one schema: final content with an empty tool_calls array, or empty content with tool calls. " 'Never put tool_calls inside content. Never output {"content":"tool_calls": ...}.' ) - system_instruction = ( - "\n\nAvailable tools JSON schema:\n" - + json.dumps(tools_schema, ensure_ascii=False) - + "\n" - + format_instruction - ) + + system_instruction = "\n\n" + format_instruction + + if include_tools_schema: + system_instruction = ( + "\n\nAvailable tools JSON schema:\n" + + json.dumps(tools_schema, ensure_ascii=False) + + "\n" + + format_instruction + ) + if prompt_messages and prompt_messages[0].get("role") == "system": prompt_messages[0]["content"] += system_instruction else: @@ -344,7 +349,42 @@ def _prompt_json_generate(config_path: Path, config: dict, messages: list[dict], model_config.get("device_map", "auto"), model_config.get("max_memory"), ) - prompt_messages = _build_prompt_messages(messages, tools_schema) + tool_calling_config = config.get("tool_calling", {}) + schema_injection = tool_calling_config.get("schema_injection", "prompt") + + if schema_injection not in {"prompt", "native"}: + raise ValueError( + "tool_calling.schema_injection must be prompt or native" + ) + + template_kwargs = {} + + if schema_injection == "prompt": + # 原有方案:schema 转成文字后拼接到 system prompt。 + prompt_messages = _build_prompt_messages( + messages, + tools_schema, + include_tools_schema=True, + ) + + else: + # 新方案:schema 不写入 prompt,通过 tokenizer 的 tools 参数注入。 + prompt_messages = _build_prompt_messages( + messages, + tools_schema, + include_tools_schema=False, + ) + + # 内部临时移除B3 的 x-returns 字段, + native_tools = deepcopy(tools_schema) + + for tool in native_tools: + function = tool.get("function") + if isinstance(function, dict): + function.pop("x-returns", None) + + template_kwargs["tools"] = native_tools + inputs = tokenizer.apply_chat_template( prompt_messages, tokenize=True, @@ -352,6 +392,7 @@ def _prompt_json_generate(config_path: Path, config: dict, messages: list[dict], return_tensors="pt", return_dict=True, enable_thinking=False, + **template_kwargs, ) device = next(model.parameters()).device inputs = inputs.to(device) diff --git a/configs/model.yaml b/configs/model.yaml index d8d1182..2c1db8d 100644 --- a/configs/model.yaml +++ b/configs/model.yaml @@ -21,7 +21,8 @@ context: tool_calling: mode: prompt_json - + schema_injection: prompt + outputs: save_raw_output: true save_ai_message: true