test commited on
Commit
bb5fa10
·
1 Parent(s): 594db87

Fix Kimi tool-call reasoning_content normalization

Browse files
internal/runtime/executor/kimi_executor.go CHANGED
@@ -20,6 +20,7 @@ import (
20
  cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
21
  sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
22
  log "github.com/sirupsen/logrus"
 
23
  "github.com/tidwall/sjson"
24
  )
25
 
@@ -94,6 +95,10 @@ func (e *KimiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
94
 
95
  requestedModel := payloadRequestedModel(opts, req.Model)
96
  body = applyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
 
 
 
 
97
 
98
  url := kimiauth.KimiAPIBaseURL + "/chat/completions"
99
  httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
@@ -189,6 +194,10 @@ func (e *KimiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
189
  }
190
  requestedModel := payloadRequestedModel(opts, req.Model)
191
  body = applyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
 
 
 
 
192
 
193
  url := kimiauth.KimiAPIBaseURL + "/chat/completions"
194
  httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
@@ -291,6 +300,150 @@ func (e *KimiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth,
291
  return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
292
  }
293
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
294
  // Refresh refreshes the Kimi token using the refresh token.
295
  func (e *KimiExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
296
  log.Debugf("kimi executor: refresh called")
 
20
  cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
21
  sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
22
  log "github.com/sirupsen/logrus"
23
+ "github.com/tidwall/gjson"
24
  "github.com/tidwall/sjson"
25
  )
26
 
 
95
 
96
  requestedModel := payloadRequestedModel(opts, req.Model)
97
  body = applyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
98
+ body, err = normalizeKimiToolMessageLinks(body)
99
+ if err != nil {
100
+ return resp, err
101
+ }
102
 
103
  url := kimiauth.KimiAPIBaseURL + "/chat/completions"
104
  httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
 
194
  }
195
  requestedModel := payloadRequestedModel(opts, req.Model)
196
  body = applyPayloadConfigWithRoot(e.cfg, baseModel, to.String(), "", body, originalTranslated, requestedModel)
197
+ body, err = normalizeKimiToolMessageLinks(body)
198
+ if err != nil {
199
+ return nil, err
200
+ }
201
 
202
  url := kimiauth.KimiAPIBaseURL + "/chat/completions"
203
  httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
 
300
  return cliproxyexecutor.Response{Payload: []byte(translated)}, nil
301
  }
302
 
303
+ func normalizeKimiToolMessageLinks(body []byte) ([]byte, error) {
304
+ if len(body) == 0 || !gjson.ValidBytes(body) {
305
+ return body, nil
306
+ }
307
+
308
+ messages := gjson.GetBytes(body, "messages")
309
+ if !messages.Exists() || !messages.IsArray() {
310
+ return body, nil
311
+ }
312
+
313
+ out := body
314
+ pending := make([]string, 0)
315
+ patched := 0
316
+ patchedReasoning := 0
317
+ ambiguous := 0
318
+ latestReasoning := ""
319
+ hasLatestReasoning := false
320
+
321
+ removePending := func(id string) {
322
+ for idx := range pending {
323
+ if pending[idx] != id {
324
+ continue
325
+ }
326
+ pending = append(pending[:idx], pending[idx+1:]...)
327
+ return
328
+ }
329
+ }
330
+
331
+ msgs := messages.Array()
332
+ for msgIdx := range msgs {
333
+ msg := msgs[msgIdx]
334
+ role := strings.TrimSpace(msg.Get("role").String())
335
+ switch role {
336
+ case "assistant":
337
+ reasoning := msg.Get("reasoning_content")
338
+ if reasoning.Exists() {
339
+ reasoningText := reasoning.String()
340
+ if strings.TrimSpace(reasoningText) != "" {
341
+ latestReasoning = reasoningText
342
+ hasLatestReasoning = true
343
+ }
344
+ }
345
+
346
+ toolCalls := msg.Get("tool_calls")
347
+ if !toolCalls.Exists() || !toolCalls.IsArray() || len(toolCalls.Array()) == 0 {
348
+ continue
349
+ }
350
+
351
+ if !reasoning.Exists() || strings.TrimSpace(reasoning.String()) == "" {
352
+ reasoningText := fallbackAssistantReasoning(msg, hasLatestReasoning, latestReasoning)
353
+ path := fmt.Sprintf("messages.%d.reasoning_content", msgIdx)
354
+ next, err := sjson.SetBytes(out, path, reasoningText)
355
+ if err != nil {
356
+ return body, fmt.Errorf("kimi executor: failed to set assistant reasoning_content: %w", err)
357
+ }
358
+ out = next
359
+ patchedReasoning++
360
+ }
361
+
362
+ for _, tc := range toolCalls.Array() {
363
+ id := strings.TrimSpace(tc.Get("id").String())
364
+ if id == "" {
365
+ continue
366
+ }
367
+ pending = append(pending, id)
368
+ }
369
+ case "tool":
370
+ toolCallID := strings.TrimSpace(msg.Get("tool_call_id").String())
371
+ if toolCallID == "" {
372
+ toolCallID = strings.TrimSpace(msg.Get("call_id").String())
373
+ if toolCallID != "" {
374
+ path := fmt.Sprintf("messages.%d.tool_call_id", msgIdx)
375
+ next, err := sjson.SetBytes(out, path, toolCallID)
376
+ if err != nil {
377
+ return body, fmt.Errorf("kimi executor: failed to set tool_call_id from call_id: %w", err)
378
+ }
379
+ out = next
380
+ patched++
381
+ }
382
+ }
383
+ if toolCallID == "" {
384
+ if len(pending) == 1 {
385
+ toolCallID = pending[0]
386
+ path := fmt.Sprintf("messages.%d.tool_call_id", msgIdx)
387
+ next, err := sjson.SetBytes(out, path, toolCallID)
388
+ if err != nil {
389
+ return body, fmt.Errorf("kimi executor: failed to infer tool_call_id: %w", err)
390
+ }
391
+ out = next
392
+ patched++
393
+ } else if len(pending) > 1 {
394
+ ambiguous++
395
+ }
396
+ }
397
+ if toolCallID != "" {
398
+ removePending(toolCallID)
399
+ }
400
+ }
401
+ }
402
+
403
+ if patched > 0 || patchedReasoning > 0 {
404
+ log.WithFields(log.Fields{
405
+ "patched_tool_messages": patched,
406
+ "patched_reasoning_messages": patchedReasoning,
407
+ }).Debug("kimi executor: normalized tool message fields")
408
+ }
409
+ if ambiguous > 0 {
410
+ log.WithFields(log.Fields{
411
+ "ambiguous_tool_messages": ambiguous,
412
+ "pending_tool_calls": len(pending),
413
+ }).Warn("kimi executor: tool messages missing tool_call_id with ambiguous candidates")
414
+ }
415
+
416
+ return out, nil
417
+ }
418
+
419
+ func fallbackAssistantReasoning(msg gjson.Result, hasLatest bool, latest string) string {
420
+ if hasLatest && strings.TrimSpace(latest) != "" {
421
+ return latest
422
+ }
423
+
424
+ content := msg.Get("content")
425
+ if content.Type == gjson.String {
426
+ if text := strings.TrimSpace(content.String()); text != "" {
427
+ return text
428
+ }
429
+ }
430
+ if content.IsArray() {
431
+ parts := make([]string, 0, len(content.Array()))
432
+ for _, item := range content.Array() {
433
+ text := strings.TrimSpace(item.Get("text").String())
434
+ if text == "" {
435
+ continue
436
+ }
437
+ parts = append(parts, text)
438
+ }
439
+ if len(parts) > 0 {
440
+ return strings.Join(parts, "\n")
441
+ }
442
+ }
443
+
444
+ return "[reasoning unavailable]"
445
+ }
446
+
447
  // Refresh refreshes the Kimi token using the refresh token.
448
  func (e *KimiExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
449
  log.Debugf("kimi executor: refresh called")
internal/runtime/executor/kimi_executor_test.go ADDED
@@ -0,0 +1,205 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ package executor
2
+
3
+ import (
4
+ "testing"
5
+
6
+ "github.com/tidwall/gjson"
7
+ )
8
+
9
+ func TestNormalizeKimiToolMessageLinks_UsesCallIDFallback(t *testing.T) {
10
+ body := []byte(`{
11
+ "messages":[
12
+ {"role":"assistant","tool_calls":[{"id":"list_directory:1","type":"function","function":{"name":"list_directory","arguments":"{}"}}]},
13
+ {"role":"tool","call_id":"list_directory:1","content":"[]"}
14
+ ]
15
+ }`)
16
+
17
+ out, err := normalizeKimiToolMessageLinks(body)
18
+ if err != nil {
19
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
20
+ }
21
+
22
+ got := gjson.GetBytes(out, "messages.1.tool_call_id").String()
23
+ if got != "list_directory:1" {
24
+ t.Fatalf("messages.1.tool_call_id = %q, want %q", got, "list_directory:1")
25
+ }
26
+ }
27
+
28
+ func TestNormalizeKimiToolMessageLinks_InferSinglePendingID(t *testing.T) {
29
+ body := []byte(`{
30
+ "messages":[
31
+ {"role":"assistant","tool_calls":[{"id":"call_123","type":"function","function":{"name":"read_file","arguments":"{}"}}]},
32
+ {"role":"tool","content":"file-content"}
33
+ ]
34
+ }`)
35
+
36
+ out, err := normalizeKimiToolMessageLinks(body)
37
+ if err != nil {
38
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
39
+ }
40
+
41
+ got := gjson.GetBytes(out, "messages.1.tool_call_id").String()
42
+ if got != "call_123" {
43
+ t.Fatalf("messages.1.tool_call_id = %q, want %q", got, "call_123")
44
+ }
45
+ }
46
+
47
+ func TestNormalizeKimiToolMessageLinks_AmbiguousMissingIDIsNotInferred(t *testing.T) {
48
+ body := []byte(`{
49
+ "messages":[
50
+ {"role":"assistant","tool_calls":[
51
+ {"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}},
52
+ {"id":"call_2","type":"function","function":{"name":"read_file","arguments":"{}"}}
53
+ ]},
54
+ {"role":"tool","content":"result-without-id"}
55
+ ]
56
+ }`)
57
+
58
+ out, err := normalizeKimiToolMessageLinks(body)
59
+ if err != nil {
60
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
61
+ }
62
+
63
+ if gjson.GetBytes(out, "messages.1.tool_call_id").Exists() {
64
+ t.Fatalf("messages.1.tool_call_id should be absent for ambiguous case, got %q", gjson.GetBytes(out, "messages.1.tool_call_id").String())
65
+ }
66
+ }
67
+
68
+ func TestNormalizeKimiToolMessageLinks_PreservesExistingToolCallID(t *testing.T) {
69
+ body := []byte(`{
70
+ "messages":[
71
+ {"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}]},
72
+ {"role":"tool","tool_call_id":"call_1","call_id":"different-id","content":"result"}
73
+ ]
74
+ }`)
75
+
76
+ out, err := normalizeKimiToolMessageLinks(body)
77
+ if err != nil {
78
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
79
+ }
80
+
81
+ got := gjson.GetBytes(out, "messages.1.tool_call_id").String()
82
+ if got != "call_1" {
83
+ t.Fatalf("messages.1.tool_call_id = %q, want %q", got, "call_1")
84
+ }
85
+ }
86
+
87
+ func TestNormalizeKimiToolMessageLinks_InheritsPreviousReasoningForAssistantToolCalls(t *testing.T) {
88
+ body := []byte(`{
89
+ "messages":[
90
+ {"role":"assistant","content":"plan","reasoning_content":"previous reasoning"},
91
+ {"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}]}
92
+ ]
93
+ }`)
94
+
95
+ out, err := normalizeKimiToolMessageLinks(body)
96
+ if err != nil {
97
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
98
+ }
99
+
100
+ got := gjson.GetBytes(out, "messages.1.reasoning_content").String()
101
+ if got != "previous reasoning" {
102
+ t.Fatalf("messages.1.reasoning_content = %q, want %q", got, "previous reasoning")
103
+ }
104
+ }
105
+
106
+ func TestNormalizeKimiToolMessageLinks_InsertsFallbackReasoningWhenMissing(t *testing.T) {
107
+ body := []byte(`{
108
+ "messages":[
109
+ {"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}]}
110
+ ]
111
+ }`)
112
+
113
+ out, err := normalizeKimiToolMessageLinks(body)
114
+ if err != nil {
115
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
116
+ }
117
+
118
+ reasoning := gjson.GetBytes(out, "messages.0.reasoning_content")
119
+ if !reasoning.Exists() {
120
+ t.Fatalf("messages.0.reasoning_content should exist")
121
+ }
122
+ if reasoning.String() != "[reasoning unavailable]" {
123
+ t.Fatalf("messages.0.reasoning_content = %q, want %q", reasoning.String(), "[reasoning unavailable]")
124
+ }
125
+ }
126
+
127
+ func TestNormalizeKimiToolMessageLinks_UsesContentAsReasoningFallback(t *testing.T) {
128
+ body := []byte(`{
129
+ "messages":[
130
+ {"role":"assistant","content":[{"type":"text","text":"first line"},{"type":"text","text":"second line"}],"tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}]}
131
+ ]
132
+ }`)
133
+
134
+ out, err := normalizeKimiToolMessageLinks(body)
135
+ if err != nil {
136
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
137
+ }
138
+
139
+ got := gjson.GetBytes(out, "messages.0.reasoning_content").String()
140
+ if got != "first line\nsecond line" {
141
+ t.Fatalf("messages.0.reasoning_content = %q, want %q", got, "first line\nsecond line")
142
+ }
143
+ }
144
+
145
+ func TestNormalizeKimiToolMessageLinks_ReplacesEmptyReasoningContent(t *testing.T) {
146
+ body := []byte(`{
147
+ "messages":[
148
+ {"role":"assistant","content":"assistant summary","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}],"reasoning_content":""}
149
+ ]
150
+ }`)
151
+
152
+ out, err := normalizeKimiToolMessageLinks(body)
153
+ if err != nil {
154
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
155
+ }
156
+
157
+ got := gjson.GetBytes(out, "messages.0.reasoning_content").String()
158
+ if got != "assistant summary" {
159
+ t.Fatalf("messages.0.reasoning_content = %q, want %q", got, "assistant summary")
160
+ }
161
+ }
162
+
163
+ func TestNormalizeKimiToolMessageLinks_PreservesExistingAssistantReasoning(t *testing.T) {
164
+ body := []byte(`{
165
+ "messages":[
166
+ {"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}],"reasoning_content":"keep me"}
167
+ ]
168
+ }`)
169
+
170
+ out, err := normalizeKimiToolMessageLinks(body)
171
+ if err != nil {
172
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
173
+ }
174
+
175
+ got := gjson.GetBytes(out, "messages.0.reasoning_content").String()
176
+ if got != "keep me" {
177
+ t.Fatalf("messages.0.reasoning_content = %q, want %q", got, "keep me")
178
+ }
179
+ }
180
+
181
+ func TestNormalizeKimiToolMessageLinks_RepairsIDsAndReasoningTogether(t *testing.T) {
182
+ body := []byte(`{
183
+ "messages":[
184
+ {"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"list_directory","arguments":"{}"}}],"reasoning_content":"r1"},
185
+ {"role":"tool","call_id":"call_1","content":"[]"},
186
+ {"role":"assistant","tool_calls":[{"id":"call_2","type":"function","function":{"name":"read_file","arguments":"{}"}}]},
187
+ {"role":"tool","call_id":"call_2","content":"file"}
188
+ ]
189
+ }`)
190
+
191
+ out, err := normalizeKimiToolMessageLinks(body)
192
+ if err != nil {
193
+ t.Fatalf("normalizeKimiToolMessageLinks() error = %v", err)
194
+ }
195
+
196
+ if got := gjson.GetBytes(out, "messages.1.tool_call_id").String(); got != "call_1" {
197
+ t.Fatalf("messages.1.tool_call_id = %q, want %q", got, "call_1")
198
+ }
199
+ if got := gjson.GetBytes(out, "messages.3.tool_call_id").String(); got != "call_2" {
200
+ t.Fatalf("messages.3.tool_call_id = %q, want %q", got, "call_2")
201
+ }
202
+ if got := gjson.GetBytes(out, "messages.2.reasoning_content").String(); got != "r1" {
203
+ t.Fatalf("messages.2.reasoning_content = %q, want %q", got, "r1")
204
+ }
205
+ }