feat: remove tool_choice for image_generation when disabled
- Added logic to remove `tool_choice` entries of type `image_generation` from payloads when `disable-image-generation` is enabled. - Updated `ApplyPayloadConfigWithRoot` to handle new removal logic. - Added unit tests to verify `tool_choice` removal behavior.
This commit is contained in:
@@ -151,6 +151,7 @@ func ApplyPayloadConfigWithRoot(cfg *config.Config, model, protocol, root string
|
||||
|
||||
if cfg.DisableImageGeneration {
|
||||
out = removeToolTypeFromPayloadWithRoot(out, root, "image_generation")
|
||||
out = removeToolChoiceFromPayloadWithRoot(out, root, "image_generation")
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -242,6 +243,55 @@ func removeToolTypeFromPayloadWithRoot(payload []byte, root string, toolType str
|
||||
return removeToolTypeFromToolsArray(payload, toolsPath, toolType)
|
||||
}
|
||||
|
||||
func removeToolChoiceFromPayloadWithRoot(payload []byte, root string, toolType string) []byte {
|
||||
if len(payload) == 0 {
|
||||
return payload
|
||||
}
|
||||
toolType = strings.TrimSpace(toolType)
|
||||
if toolType == "" {
|
||||
return payload
|
||||
}
|
||||
toolChoicePath := buildPayloadPath(root, "tool_choice")
|
||||
return removeToolChoiceFromPayload(payload, toolChoicePath, toolType)
|
||||
}
|
||||
|
||||
func removeToolChoiceFromPayload(payload []byte, toolChoicePath string, toolType string) []byte {
|
||||
choice := gjson.GetBytes(payload, toolChoicePath)
|
||||
if !choice.Exists() {
|
||||
return payload
|
||||
}
|
||||
if choice.Type == gjson.String {
|
||||
if strings.EqualFold(strings.TrimSpace(choice.String()), toolType) {
|
||||
updated, errDel := sjson.DeleteBytes(payload, toolChoicePath)
|
||||
if errDel == nil {
|
||||
return updated
|
||||
}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
if choice.Type != gjson.JSON {
|
||||
return payload
|
||||
}
|
||||
choiceType := strings.TrimSpace(choice.Get("type").String())
|
||||
if strings.EqualFold(choiceType, toolType) {
|
||||
updated, errDel := sjson.DeleteBytes(payload, toolChoicePath)
|
||||
if errDel == nil {
|
||||
return updated
|
||||
}
|
||||
return payload
|
||||
}
|
||||
if strings.EqualFold(choiceType, "tool") {
|
||||
name := strings.TrimSpace(choice.Get("name").String())
|
||||
if strings.EqualFold(name, toolType) {
|
||||
updated, errDel := sjson.DeleteBytes(payload, toolChoicePath)
|
||||
if errDel == nil {
|
||||
return updated
|
||||
}
|
||||
}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
func removeToolTypeFromToolsArray(payload []byte, toolsPath string, toolType string) []byte {
|
||||
tools := gjson.GetBytes(payload, toolsPath)
|
||||
if !tools.Exists() || !tools.IsArray() {
|
||||
|
||||
Reference in New Issue
Block a user