[tool] Fix tool calling output format

deepjavalibrary · Feb 5, 2025 · 6fe111a · 6fe111a
1 parent 390fc69
commit 6fe111a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/engines/python/setup/djl_python/chat_completions/vllm_chat_utils.py b/engines/python/setup/djl_python/chat_completions/vllm_chat_utils.py
@@ -88,7 +88,9 @@ def parse_chat_completions_request_vllm(
         "sampling_params": sampling_params,
         "conversation": conversation,
         "request_prompts": request_prompt,
-        "engine_prompt": engine_prompt
+        "engine_prompt": engine_prompt,
+        "tool_parser": tool_parser,
+        "chat_params": chat_params,
     }
     return input_text, params
 
@@ -136,7 +138,7 @@ def _preprocess_chat(
 
     should_parse_tools = tool_parser is not None and request.tool_choice != "none"
     if should_parse_tools:
-        request = tool_parser(tokenizer).adjust_request(request=request)
+        request = tool_parser.adjust_request(request=request)
 
     if isinstance(request_prompt, str):
         # Hf tokenizer case

diff --git a/engines/python/setup/djl_python/output_formatter.py b/engines/python/setup/djl_python/output_formatter.py
@@ -322,9 +322,10 @@ def _json_chat_output_formatter(request_output: TextGenerationOutput):
             "logprobs": None,
             "finish_reason": best_sequence.finish_reason,
         }
-    elif parameters.get("tools") and (parameters.get("tool_choice") == "auto"
-                                      or parameters.get("tool_choice") is None
-                                      ) and parameters.get("tool_parser"):
+    elif chat_params and chat_params.tools and (
+            parameters.get("tool_choice") == "auto"
+            or parameters.get("tool_choice") is None
+    ) and parameters.get("tool_parser"):
         tool_call_info = tool_parser.extract_tool_calls(generated_text,
                                                         request=chat_params)
         auto_tools_called = tool_call_info.tools_called

diff --git a/engines/python/setup/djl_python/rolling_batch/vllm_rolling_batch.py b/engines/python/setup/djl_python/rolling_batch/vllm_rolling_batch.py
@@ -57,6 +57,8 @@ def __init__(self, model_id_or_path: str, properties: dict,
             try:
                 self.tool_parser = ToolParserManager.get_tool_parser(
                     self.vllm_configs.tool_call_parser)
+                self.tool_parser = self.tool_parser(
+                    self.engine.tokenizer.tokenizer)
             except Exception as e:
                 raise TypeError("Error in tool parser creation.") from e