fix tool handling in OpenAIServingChat#18996
fix tool handling in OpenAIServingChat#18996JustinTong0323 wants to merge 1 commit intosgl-project:mainfrom
Conversation
…chema behavior Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Summary of ChangesHello @JustinTong0323, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request refines the handling of tool schemas within the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
There was a problem hiding this comment.
Code Review
This pull request correctly refactors the tool handling in OpenAIServingChat to pass the full OpenAI-formatted tool schemas to Jinja templates, which is a great improvement for compatibility with templates expecting this format. The corresponding fix to the fallback logic, which now correctly unwraps tools to a flat format, is also well-implemented. The new unit tests are thorough and validate both the primary path and the fallback mechanism. I have one suggestion to improve the robustness of the exception handling in the fallback logic by catching more specific exceptions.
| @@ -481,11 +481,10 @@ def _apply_jinja_template( | |||
| return_dict=False, | |||
| ) | |||
| except Exception as e: | |||
There was a problem hiding this comment.
Using a broad except Exception can hide underlying issues unrelated to tool formatting, such as OutOfMemoryError or other runtime problems. This could lead to incorrect fallback attempts. To make the error handling more robust, it's better to catch more specific exceptions that are likely to be raised by apply_chat_template due to template incompatibility, like jinja2.TemplateError, TypeError, or ValueError.
| except Exception as e: | |
| except (jinja2.TemplateError, TypeError, ValueError) as e: |
|
This PR may also resolve or is related to #18086. However, using the https://github.com/MoonshotAI/K2-Vendor-Verifier, I cannot observe any obvious accuracy difference, requiring further investigation. |
Summary
{"type": "function", "function": {...}}) to Jinja chat templates instead of flat function-only dicts ({"name": ..., "parameters": ...})Motivation
This PR address #18478
The previous code called
item.function.model_dump(), stripping the outer OpenAI wrapper before passing tools toapply_chat_template. This caused templates that rely on the standard{"type": "function", "function": {...}}structure to fall back to incorrect rendering.For example, with Kimi-K2.5, the old flat format was not recognized by the Jinja template, resulting in a raw JSON dump:
With the fix (
item.model_dump()), the template correctly identifies the OpenAI tool schema and renders the expected TypeScript namespace format that the model was trained on:Additionally, the old fallback logic (
[t if "function" in t else {"function": t} for t in tools]) was designed to wrap flat dicts outward. After switching toitem.model_dump(),"function"is already a key in every dict, making the fallback a no-op — it would retry with the exact same input. This PR fixes the fallback to unwrap instead (t["function"] if "function" in t else t), so templates that expect flat function-only dicts (e.g. some Mistral variants) still work on the second attempt.Test plan
test_jinja_uses_openai_tool_schema_first— verifiesapply_chat_templatereceives full OpenAI-shaped tools on the first calltest_jinja_tool_schema_fallback_to_flat_function— simulates a template rejection, verifies the retry passes flatfunction-only dictsChecklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci