Skip to content

Commit 2a5b7dc

Browse files
luispaterluispater
committed
fix(kimi): normalize tool and function parameter schemas
- Inline local `$ref` pointers and strip `$defs` and `definitions` from tool parameter schemas for Moonshot compatibility. - Ensure tool parameter root objects declare an explicit `type: "object"`. Closes: #5316
1 parent 2a5234d commit 2a5b7dc

2 files changed

Lines changed: 309 additions & 0 deletions

File tree

‎internal/runtime/executor/kimi_executor.go‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
144144
if err != nil {
145145
return resp, err
146146
}
147+
body = normalizeKimiTools(body)
147148
reporter.SetTranslatedReasoningEffort(body, e.Identifier())
148149

149150
url := kimiauth.KimiAPIBaseURL + "/v1/chat/completions"
@@ -268,6 +269,7 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
268269
if err != nil {
269270
return nil, err
270271
}
272+
body = normalizeKimiTools(body)
271273
reporter.SetTranslatedReasoningEffort(body, e.Identifier())
272274

273275
url := kimiauth.KimiAPIBaseURL + "/v1/chat/completions"
@@ -848,3 +850,84 @@ func normalizeKimiUpstreamModel(model string) string {
848850
}
849851
return normalized
850852
}
853+
854+
// normalizeKimiTools normalizes tool and legacy function parameter schemas for Moonshot.
855+
// It resolves and inlines local $ref pointers, strips $defs / definitions, and ensures
856+
// parameters root objects declare an explicit type: "object".
857+
func normalizeKimiTools(body []byte) []byte {
858+
if len(body) == 0 {
859+
return body
860+
}
861+
body = normalizeKimiToolList(body, "tools", true)
862+
body = normalizeKimiToolList(body, "functions", false)
863+
return body
864+
}
865+
866+
func normalizeKimiToolList(body []byte, arrayKey string, isTools bool) []byte {
867+
items := gjson.GetBytes(body, arrayKey)
868+
if !items.Exists() || !items.IsArray() {
869+
return body
870+
}
871+
arr := items.Array()
872+
if len(arr) == 0 {
873+
return body
874+
}
875+
876+
changed := false
877+
var updatedItems []string
878+
for _, item := range arr {
879+
itemRaw := item.Raw
880+
var paramPath string
881+
if isTools && item.Get("function.parameters").Exists() {
882+
paramPath = "function.parameters"
883+
} else if item.Get("parameters").Exists() {
884+
paramPath = "parameters"
885+
}
886+
887+
if paramPath != "" {
888+
rawParams := item.Get(paramPath)
889+
if rawParams.IsObject() {
890+
normalizedParams := normalizeKimiParametersSchema(rawParams.Raw)
891+
if normalizedParams != rawParams.Raw {
892+
if updated, errSet := sjson.SetRawBytes([]byte(itemRaw), paramPath, []byte(normalizedParams)); errSet == nil {
893+
itemRaw = string(updated)
894+
changed = true
895+
}
896+
}
897+
}
898+
}
899+
updatedItems = append(updatedItems, itemRaw)
900+
}
901+
902+
if !changed {
903+
return body
904+
}
905+
906+
out, errSetRaw := sjson.SetRawBytes(body, arrayKey, helps.JoinRawJSONStrings(updatedItems))
907+
if errSetRaw != nil {
908+
return body
909+
}
910+
return out
911+
}
912+
913+
func normalizeKimiParametersSchema(paramsRaw string) string {
914+
if strings.TrimSpace(paramsRaw) == "" {
915+
return paramsRaw
916+
}
917+
918+
inlined := util.InlineLocalRefs(paramsRaw)
919+
paramBytes := []byte(inlined)
920+
921+
if inlinedDefs := gjson.GetBytes(paramBytes, "$defs"); inlinedDefs.Exists() {
922+
paramBytes, _ = sjson.DeleteBytes(paramBytes, "$defs")
923+
}
924+
if inlinedDefinitions := gjson.GetBytes(paramBytes, "definitions"); inlinedDefinitions.Exists() {
925+
paramBytes, _ = sjson.DeleteBytes(paramBytes, "definitions")
926+
}
927+
928+
if rootType := gjson.GetBytes(paramBytes, "type"); !rootType.Exists() {
929+
paramBytes, _ = sjson.SetBytes(paramBytes, "type", "object")
930+
}
931+
932+
return string(paramBytes)
933+
}

‎internal/runtime/executor/kimi_executor_test.go‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -703,3 +703,229 @@ func TestNormalizeKimiUpstreamModel(t *testing.T) {
703703
}
704704
}
705705
}
706+
707+
func TestKimiExecutorNormalizesToolSchemasForMoonshot(t *testing.T) {
708+
var upstreamBody []byte
709+
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) {
710+
var errRead error
711+
upstreamBody, errRead = io.ReadAll(req.Body)
712+
if errRead != nil {
713+
return nil, errRead
714+
}
715+
return &http.Response{
716+
StatusCode: http.StatusOK,
717+
Header: http.Header{"Content-Type": []string{"application/json"}},
718+
Body: io.NopCloser(strings.NewReader(
719+
`{"id":"chatcmpl_test","object":"chat.completion","created":1,"model":"k3","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`,
720+
)),
721+
}, nil
722+
}))
723+
724+
executor := NewKimiExecutor(&config.Config{})
725+
auth := &cliproxyauth.Auth{
726+
Attributes: map[string]string{},
727+
Metadata: map[string]any{"access_token": "test-token"},
728+
}
729+
730+
// Tool with $defs and $ref sibling type, matching the shape used by Codex Desktop app tools
731+
payload := []byte(`{
732+
"model":"kimi-k3",
733+
"messages":[{"role":"user","content":"hello"}],
734+
"tools":[
735+
{
736+
"type":"function",
737+
"function":{
738+
"name":"codex_app__automation_update",
739+
"description":"Update automation",
740+
"parameters":{
741+
"$defs":{
742+
"value":{
743+
"type":"string"
744+
}
745+
},
746+
"type":"object",
747+
"properties":{
748+
"value":{
749+
"$ref":"#/$defs/value",
750+
"type":"string",
751+
"description":"A value"
752+
}
753+
}
754+
}
755+
}
756+
}
757+
]
758+
}`)
759+
760+
_, err := executor.Execute(ctx, auth, cliproxyexecutor.Request{
761+
Model: "kimi-k3",
762+
Payload: payload,
763+
}, cliproxyexecutor.Options{
764+
SourceFormat: sdktranslator.FormatOpenAI,
765+
OriginalRequest: payload,
766+
})
767+
if err != nil {
768+
t.Fatalf("Execute() error = %v", err)
769+
}
770+
771+
tool := gjson.GetBytes(upstreamBody, "tools.0.function")
772+
if !tool.Exists() {
773+
t.Fatalf("upstream tool function not found: %s", upstreamBody)
774+
}
775+
776+
// Moonshot rejects $ref with sibling type; verify $ref is inlined and removed
777+
if ref := tool.Get("parameters.properties.value.$ref"); ref.Exists() {
778+
t.Fatalf("upstream tool parameter still contains $ref: %s", tool.Get("parameters").Raw)
779+
}
780+
if got := tool.Get("parameters.properties.value.type").String(); got != "string" {
781+
t.Fatalf("upstream tool parameter value.type = %q, want %q", got, "string")
782+
}
783+
if got := tool.Get("parameters.properties.value.description").String(); got != "A value" {
784+
t.Fatalf("upstream tool parameter value.description = %q, want %q", got, "A value")
785+
}
786+
// Verify $defs container was pruned
787+
if defs := tool.Get("parameters.$defs"); defs.Exists() {
788+
t.Fatalf("upstream tool parameter still contains $defs: %s", tool.Get("parameters").Raw)
789+
}
790+
// Verify explicit object type
791+
if got := tool.Get("parameters.type").String(); got != "object" {
792+
t.Fatalf("upstream tool parameters.type = %q, want %q", got, "object")
793+
}
794+
}
795+
796+
func TestKimiExecutorStreamNormalizesToolSchemasFromResponses(t *testing.T) {
797+
var upstreamBody []byte
798+
ctx := context.WithValue(context.Background(), "cliproxy.roundtripper", kimiRoundTripperFunc(func(req *http.Request) (*http.Response, error) {
799+
var errRead error
800+
upstreamBody, errRead = io.ReadAll(req.Body)
801+
if errRead != nil {
802+
return nil, errRead
803+
}
804+
return &http.Response{
805+
StatusCode: http.StatusOK,
806+
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
807+
Body: io.NopCloser(strings.NewReader(
808+
"data: {\"id\":\"chatcmpl_1\",\"object\":\"chat.completion.chunk\",\"created\":1,\"model\":\"k3\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n" +
809+
"data: [DONE]\n\n",
810+
)),
811+
}, nil
812+
}))
813+
814+
executor := NewKimiExecutor(&config.Config{})
815+
auth := &cliproxyauth.Auth{
816+
Attributes: map[string]string{},
817+
Metadata: map[string]any{"access_token": "test-token"},
818+
}
819+
820+
// Codex Responses format payload with tools containing $defs and $ref
821+
payload := []byte(`{
822+
"model":"kimi-k3",
823+
"input":[{"type":"message","role":"user","content":"hello"}],
824+
"tools":[
825+
{
826+
"type":"function",
827+
"name":"codex_app__automation_update",
828+
"description":"Update automation",
829+
"parameters":{
830+
"$defs":{
831+
"sub":{
832+
"type":"string"
833+
}
834+
},
835+
"properties":{
836+
"field":{
837+
"$ref":"#/$defs/sub",
838+
"type":"string"
839+
}
840+
}
841+
}
842+
}
843+
]
844+
}`)
845+
846+
streamResult, err := executor.ExecuteStream(ctx, auth, cliproxyexecutor.Request{
847+
Model: "kimi-k3",
848+
Payload: payload,
849+
}, cliproxyexecutor.Options{
850+
SourceFormat: sdktranslator.FormatOpenAIResponse,
851+
OriginalRequest: payload,
852+
Stream: true,
853+
})
854+
if err != nil {
855+
t.Fatalf("ExecuteStream() error = %v", err)
856+
}
857+
if streamResult != nil && streamResult.Chunks != nil {
858+
for range streamResult.Chunks {
859+
}
860+
}
861+
862+
tool := gjson.GetBytes(upstreamBody, "tools.0.function")
863+
if !tool.Exists() {
864+
t.Fatalf("upstream tool function not found in stream body: %s", upstreamBody)
865+
}
866+
if ref := tool.Get("parameters.properties.field.$ref"); ref.Exists() {
867+
t.Fatalf("upstream stream tool parameter still contains $ref: %s", tool.Get("parameters").Raw)
868+
}
869+
if got := tool.Get("parameters.properties.field.type").String(); got != "string" {
870+
t.Fatalf("upstream stream tool field.type = %q, want %q", got, "string")
871+
}
872+
if defs := tool.Get("parameters.$defs"); defs.Exists() {
873+
t.Fatalf("upstream stream tool parameter still contains $defs: %s", tool.Get("parameters").Raw)
874+
}
875+
if got := tool.Get("parameters.type").String(); got != "object" {
876+
t.Fatalf("upstream stream tool parameters.type = %q, want %q", got, "object")
877+
}
878+
}
879+
880+
func TestNormalizeKimiToolsDirect(t *testing.T) {
881+
input := []byte(`{
882+
"tools":[
883+
{
884+
"type":"function",
885+
"function":{
886+
"name":"test_fn",
887+
"parameters":{
888+
"definitions":{
889+
"prop":{"type":"number"}
890+
},
891+
"properties":{
892+
"count":{"$ref":"#/definitions/prop","description":"item count"}
893+
}
894+
}
895+
}
896+
}
897+
],
898+
"functions":[
899+
{
900+
"name":"legacy_fn",
901+
"parameters":{
902+
"properties":{"name":{"type":"string"}}
903+
}
904+
}
905+
]
906+
}`)
907+
908+
normalized := normalizeKimiTools(input)
909+
910+
toolParams := gjson.GetBytes(normalized, "tools.0.function.parameters")
911+
if toolParams.Get("definitions").Exists() {
912+
t.Errorf("definitions was not stripped: %s", toolParams.Raw)
913+
}
914+
if toolParams.Get("properties.count.$ref").Exists() {
915+
t.Errorf("$ref was not inlined: %s", toolParams.Raw)
916+
}
917+
if got := toolParams.Get("properties.count.type").String(); got != "number" {
918+
t.Errorf("properties.count.type = %q, want number", got)
919+
}
920+
if got := toolParams.Get("properties.count.description").String(); got != "item count" {
921+
t.Errorf("properties.count.description = %q, want 'item count'", got)
922+
}
923+
if got := toolParams.Get("type").String(); got != "object" {
924+
t.Errorf("tools.0.function.parameters.type = %q, want object", got)
925+
}
926+
927+
fnParams := gjson.GetBytes(normalized, "functions.0.parameters")
928+
if got := fnParams.Get("type").String(); got != "object" {
929+
t.Errorf("functions.0.parameters.type = %q, want object", got)
930+
}
931+
}

0 commit comments

Comments
 (0)