Skip to content

Commit dc04d8b

Browse files
luispaterluispater
committed
feat(translator): enhance response aggregation and annotation handling
- Implemented logic to aggregate text blocks until `message_stop` for improved consistency. - Introduced support for message annotations like `citations_delta` in content responses. - Added methods (`finalizeAssistantMessage`, `appendMessageAnnotation`) to handle message grouping and annotation appending cleanly. - Updated unit tests to verify message aggregation, annotation handling, and suppression of unwanted native events. Closes: router-for-me#3801
1 parent 58bf645 commit dc04d8b

2 files changed

Lines changed: 198 additions & 41 deletions

File tree

‎internal/translator/claude/openai/responses/claude_openai-responses_response.go‎

Lines changed: 110 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,23 @@ import (
1414
)
1515

1616
type claudeToResponsesState struct {
17-
Seq int
18-
ResponseID string
19-
CreatedAt int64
20-
CurrentMsgID string
21-
CurrentFCID string
22-
InTextBlock bool
23-
InFuncBlock bool
24-
FuncArgsBuf map[int]*strings.Builder // index -> args
17+
Seq int
18+
ResponseID string
19+
CreatedAt int64
20+
CurrentMsgID string
21+
CurrentFCID string
22+
InTextBlock bool
23+
InFuncBlock bool
24+
MessageOpen bool
25+
ContentPartOpen bool
26+
FuncArgsBuf map[int]*strings.Builder // index -> args
2527
// function call bookkeeping for output aggregation
2628
FuncNames map[int]string // index -> function name
2729
FuncCallIDs map[int]string // index -> call id
2830
// message text aggregation
29-
TextBuf strings.Builder
30-
CurrentTextBuf strings.Builder
31+
TextBuf strings.Builder
32+
CurrentTextBuf strings.Builder
33+
MessageAnnotations []any
3134
// reasoning state
3235
ReasoningActive bool
3336
ReasoningItemID string
@@ -57,6 +60,57 @@ func emitEvent(event string, payload []byte) []byte {
5760
return translatorcommon.SSEEventData(event, payload)
5861
}
5962

63+
func noSSEOutput(out [][]byte) [][]byte {
64+
if out == nil {
65+
return [][]byte{}
66+
}
67+
return out
68+
}
69+
70+
func (st *claudeToResponsesState) appendMessageAnnotation(annotation any) {
71+
if annotation == nil {
72+
return
73+
}
74+
st.MessageAnnotations = append(st.MessageAnnotations, annotation)
75+
}
76+
77+
func (st *claudeToResponsesState) finalizeAssistantMessage(nextSeq func() int) [][]byte {
78+
if !st.MessageOpen {
79+
return nil
80+
}
81+
fullText := st.TextBuf.String()
82+
var out [][]byte
83+
done := []byte(`{"type":"response.output_text.done","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"text":"","logprobs":[]}`)
84+
done, _ = sjson.SetBytes(done, "sequence_number", nextSeq())
85+
done, _ = sjson.SetBytes(done, "item_id", st.CurrentMsgID)
86+
done, _ = sjson.SetBytes(done, "text", fullText)
87+
out = append(out, emitEvent("response.output_text.done", done))
88+
89+
partDone := []byte(`{"type":"response.content_part.done","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""}}`)
90+
partDone, _ = sjson.SetBytes(partDone, "sequence_number", nextSeq())
91+
partDone, _ = sjson.SetBytes(partDone, "item_id", st.CurrentMsgID)
92+
partDone, _ = sjson.SetBytes(partDone, "part.text", fullText)
93+
if len(st.MessageAnnotations) > 0 {
94+
partDone, _ = sjson.SetBytes(partDone, "part.annotations", st.MessageAnnotations)
95+
}
96+
out = append(out, emitEvent("response.content_part.done", partDone))
97+
98+
final := []byte(`{"type":"response.output_item.done","sequence_number":0,"output_index":0,"item":{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}}`)
99+
final, _ = sjson.SetBytes(final, "sequence_number", nextSeq())
100+
final, _ = sjson.SetBytes(final, "item.id", st.CurrentMsgID)
101+
final, _ = sjson.SetBytes(final, "item.content.0.text", fullText)
102+
if len(st.MessageAnnotations) > 0 {
103+
final, _ = sjson.SetBytes(final, "item.content.0.annotations", st.MessageAnnotations)
104+
}
105+
out = append(out, emitEvent("response.output_item.done", final))
106+
107+
st.InTextBlock = false
108+
st.MessageOpen = false
109+
st.ContentPartOpen = false
110+
st.CurrentTextBuf.Reset()
111+
return out
112+
}
113+
60114
// ConvertClaudeResponseToOpenAIResponses converts Claude SSE to OpenAI Responses SSE events.
61115
func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName string, originalRequestRawJSON, requestRawJSON, rawJSON []byte, param *any) [][]byte {
62116
if *param == nil {
@@ -83,10 +137,13 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
83137
// Reset per-message aggregation state
84138
st.TextBuf.Reset()
85139
st.CurrentTextBuf.Reset()
140+
st.MessageAnnotations = nil
86141
st.ReasoningBuf.Reset()
87142
st.ReasoningActive = false
88143
st.InTextBlock = false
89144
st.InFuncBlock = false
145+
st.MessageOpen = false
146+
st.ContentPartOpen = false
90147
st.CurrentMsgID = ""
91148
st.CurrentFCID = ""
92149
st.ReasoningItemID = ""
@@ -125,24 +182,29 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
125182
case "content_block_start":
126183
cb := root.Get("content_block")
127184
if !cb.Exists() {
128-
return out
185+
return noSSEOutput(out)
129186
}
130187
idx := int(root.Get("index").Int())
131188
typ := cb.Get("type").String()
132189
if typ == "text" {
133-
// open message item + content part
134190
st.InTextBlock = true
135-
st.CurrentTextBuf.Reset()
136-
st.CurrentMsgID = fmt.Sprintf("msg_%s_0", st.ResponseID)
137-
item := []byte(`{"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"id":"","type":"message","status":"in_progress","content":[],"role":"assistant"}}`)
138-
item, _ = sjson.SetBytes(item, "sequence_number", nextSeq())
139-
item, _ = sjson.SetBytes(item, "item.id", st.CurrentMsgID)
140-
out = append(out, emitEvent("response.output_item.added", item))
141-
142-
part := []byte(`{"type":"response.content_part.added","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""}}`)
143-
part, _ = sjson.SetBytes(part, "sequence_number", nextSeq())
144-
part, _ = sjson.SetBytes(part, "item_id", st.CurrentMsgID)
145-
out = append(out, emitEvent("response.content_part.added", part))
191+
if st.CurrentMsgID == "" {
192+
st.CurrentMsgID = fmt.Sprintf("msg_%s_0", st.ResponseID)
193+
}
194+
if !st.MessageOpen {
195+
item := []byte(`{"type":"response.output_item.added","sequence_number":0,"output_index":0,"item":{"id":"","type":"message","status":"in_progress","content":[],"role":"assistant"}}`)
196+
item, _ = sjson.SetBytes(item, "sequence_number", nextSeq())
197+
item, _ = sjson.SetBytes(item, "item.id", st.CurrentMsgID)
198+
out = append(out, emitEvent("response.output_item.added", item))
199+
st.MessageOpen = true
200+
}
201+
if !st.ContentPartOpen {
202+
part := []byte(`{"type":"response.content_part.added","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""}}`)
203+
part, _ = sjson.SetBytes(part, "sequence_number", nextSeq())
204+
part, _ = sjson.SetBytes(part, "item_id", st.CurrentMsgID)
205+
out = append(out, emitEvent("response.content_part.added", part))
206+
st.ContentPartOpen = true
207+
}
146208
} else if typ == "tool_use" {
147209
st.InFuncBlock = true
148210
st.CurrentFCID = cb.Get("id").String()
@@ -187,7 +249,7 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
187249
case "content_block_delta":
188250
d := root.Get("delta")
189251
if !d.Exists() {
190-
return out
252+
return noSSEOutput(out)
191253
}
192254
dt := d.Get("type").String()
193255
if dt == "text_delta" {
@@ -202,6 +264,9 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
202264
st.CurrentTextBuf.WriteString(t.String())
203265
}
204266
} else if dt == "input_json_delta" {
267+
if !st.InFuncBlock || st.CurrentFCID == "" {
268+
return [][]byte{}
269+
}
205270
idx := int(root.Get("index").Int())
206271
if pj := d.Get("partial_json"); pj.Exists() {
207272
if st.FuncArgsBuf[idx] == nil {
@@ -233,26 +298,16 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
233298
st.ReasoningSignature = signature.String()
234299
}
235300
}
301+
return [][]byte{}
302+
} else if dt == "citations_delta" {
303+
if citation := d.Get("citation"); citation.Exists() {
304+
st.appendMessageAnnotation(citation.Value())
305+
}
306+
return [][]byte{}
236307
}
237308
case "content_block_stop":
238309
idx := int(root.Get("index").Int())
239310
if st.InTextBlock {
240-
fullText := st.CurrentTextBuf.String()
241-
done := []byte(`{"type":"response.output_text.done","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"text":"","logprobs":[]}`)
242-
done, _ = sjson.SetBytes(done, "sequence_number", nextSeq())
243-
done, _ = sjson.SetBytes(done, "item_id", st.CurrentMsgID)
244-
done, _ = sjson.SetBytes(done, "text", fullText)
245-
out = append(out, emitEvent("response.output_text.done", done))
246-
partDone := []byte(`{"type":"response.content_part.done","sequence_number":0,"item_id":"","output_index":0,"content_index":0,"part":{"type":"output_text","annotations":[],"logprobs":[],"text":""}}`)
247-
partDone, _ = sjson.SetBytes(partDone, "sequence_number", nextSeq())
248-
partDone, _ = sjson.SetBytes(partDone, "item_id", st.CurrentMsgID)
249-
partDone, _ = sjson.SetBytes(partDone, "part.text", fullText)
250-
out = append(out, emitEvent("response.content_part.done", partDone))
251-
final := []byte(`{"type":"response.output_item.done","sequence_number":0,"output_index":0,"item":{"id":"","type":"message","status":"completed","content":[{"type":"output_text","text":""}],"role":"assistant"}}`)
252-
final, _ = sjson.SetBytes(final, "sequence_number", nextSeq())
253-
final, _ = sjson.SetBytes(final, "item.id", st.CurrentMsgID)
254-
final, _ = sjson.SetBytes(final, "item.content.0.text", fullText)
255-
out = append(out, emitEvent("response.output_item.done", final))
256311
st.InTextBlock = false
257312
} else if st.InFuncBlock {
258313
args := "{}"
@@ -304,6 +359,7 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
304359
st.ReasoningActive = false
305360
st.ReasoningPartAdded = false
306361
}
362+
return noSSEOutput(out)
307363
case "message_delta":
308364
if usage := root.Get("usage"); usage.Exists() {
309365
if v := usage.Get("output_tokens"); v.Exists() {
@@ -315,7 +371,9 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
315371
st.UsageSeen = true
316372
}
317373
}
374+
return [][]byte{}
318375
case "message_stop":
376+
out = append(out, st.finalizeAssistantMessage(nextSeq)...)
319377

320378
completed := []byte(`{"type":"response.completed","sequence_number":0,"response":{"id":"","object":"response","created_at":0,"status":"completed","background":false,"error":null}}`)
321379
completed, _ = sjson.SetBytes(completed, "sequence_number", nextSeq())
@@ -407,6 +465,9 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
407465
item := []byte(`{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`)
408466
item, _ = sjson.SetBytes(item, "id", st.CurrentMsgID)
409467
item, _ = sjson.SetBytes(item, "content.0.text", st.TextBuf.String())
468+
if len(st.MessageAnnotations) > 0 {
469+
item, _ = sjson.SetBytes(item, "content.0.annotations", st.MessageAnnotations)
470+
}
410471
outputsWrapper, _ = sjson.SetRawBytes(outputsWrapper, "arr.-1", item)
411472
}
412473
// function_call items (in ascending index order for determinism)
@@ -466,7 +527,7 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
466527
out = append(out, emitEvent("response.completed", completed))
467528
}
468529

469-
return out
530+
return noSSEOutput(out)
470531
}
471532

472533
// ConvertClaudeResponseToOpenAIResponsesNonStream aggregates Claude SSE into a single OpenAI Responses JSON.
@@ -506,6 +567,7 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(_ context.Context, _ string
506567
reasoningActive bool
507568
reasoningItemID string
508569
reasoningSig string
570+
annotations []any
509571
inputTokens int64
510572
outputTokens int64
511573
)
@@ -592,6 +654,10 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(_ context.Context, _ string
592654
reasoningSig = signature.String()
593655
}
594656
}
657+
case "citations_delta":
658+
if citation := d.Get("citation"); citation.Exists() {
659+
annotations = append(annotations, citation.Value())
660+
}
595661
}
596662

597663
case "content_block_stop":
@@ -692,6 +758,9 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(_ context.Context, _ string
692758
item := []byte(`{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`)
693759
item, _ = sjson.SetBytes(item, "id", currentMsgID)
694760
item, _ = sjson.SetBytes(item, "content.0.text", textBuf.String())
761+
if len(annotations) > 0 {
762+
item, _ = sjson.SetBytes(item, "content.0.annotations", annotations)
763+
}
695764
outputsWrapper, _ = sjson.SetRawBytes(outputsWrapper, "arr.-1", item)
696765
}
697766
if len(toolCalls) > 0 {

‎internal/translator/claude/openai/responses/claude_openai-responses_response_test.go‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"strings"
66
"testing"
77

8+
sdktranslator "github.com/router-for-me/CLIProxyAPI/v7/sdk/translator"
89
"github.com/tidwall/gjson"
910
)
1011

@@ -29,6 +30,15 @@ func parseClaudeResponsesSSEEvent(t *testing.T, chunk []byte) (string, gjson.Res
2930
return event, gjson.Parse(data)
3031
}
3132

33+
func translateClaudeResponsesStreamThroughRegistry(chunks [][]byte) [][]byte {
34+
var param any
35+
var outputs [][]byte
36+
for _, chunk := range chunks {
37+
outputs = append(outputs, sdktranslator.TranslateStream(context.Background(), sdktranslator.FormatClaude, sdktranslator.FormatOpenAIResponse, "claude-test", nil, nil, chunk, &param)...)
38+
}
39+
return outputs
40+
}
41+
3242
func TestConvertClaudeResponseToOpenAIResponses_ThinkingIncludesSignature(t *testing.T) {
3343
signature := "claude_sig_123"
3444
chunks := [][]byte{
@@ -78,6 +88,84 @@ func TestConvertClaudeResponseToOpenAIResponses_ThinkingIncludesSignature(t *tes
7888
}
7989
}
8090

91+
func TestConvertClaudeResponseToOpenAIResponses_SuppressesSignatureDeltaPassthrough(t *testing.T) {
92+
chunk := []byte(`data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"claude_sig_123"}}`)
93+
94+
outputs := translateClaudeResponsesStreamThroughRegistry([][]byte{chunk})
95+
if len(outputs) != 0 {
96+
t.Fatalf("expected signature_delta to be suppressed, got %d chunks", len(outputs))
97+
}
98+
}
99+
100+
func TestConvertClaudeResponseToOpenAIResponses_AggregatesTextBlocksUntilMessageStop(t *testing.T) {
101+
chunks := [][]byte{
102+
[]byte(`data: {"type":"message_start","message":{"id":"msg_123","usage":{"input_tokens":1,"output_tokens":0}}}`),
103+
[]byte(`data: {"type":"content_block_start","index":4,"content_block":{"type":"text","text":""}}`),
104+
[]byte(`data: {"type":"content_block_delta","index":4,"delta":{"type":"text_delta","text":"**对比竞品**\n- "}}`),
105+
[]byte(`data: {"type":"content_block_stop","index":4}`),
106+
[]byte(`data: {"type":"content_block_start","index":5,"content_block":{"type":"server_tool_use","id":"srv_123","name":"web_search","input":{}}}`),
107+
[]byte(`data: {"type":"content_block_delta","index":5,"delta":{"type":"input_json_delta","partial_json":"{\"query\":\"Qwen3\"}"}}`),
108+
[]byte(`data: {"type":"content_block_stop","index":5}`),
109+
[]byte(`data: {"type":"content_block_start","index":6,"content_block":{"type":"web_search_tool_result","tool_use_id":"srv_123","content":[{"type":"web_search_result","title":"Example","url":"https://example.com"}]}}`),
110+
[]byte(`data: {"type":"content_block_stop","index":6}`),
111+
[]byte(`data: {"type":"content_block_delta","index":5,"delta":{"type":"citations_delta","citation":{"type":"web_search_result_location","cited_text":"Qwen 3.7 Max","url":"https://example.com","title":"Example"}}}`),
112+
[]byte(`data: {"type":"content_block_start","index":7,"content_block":{"type":"text","text":""}}`),
113+
[]byte(`data: {"type":"content_block_delta","index":7,"delta":{"type":"text_delta","text":"Qwen 3.7 Max leads."}}`),
114+
[]byte(`data: {"type":"content_block_stop","index":7}`),
115+
[]byte(`data: {"type":"message_delta","usage":{"output_tokens":12}}`),
116+
[]byte(`data: {"type":"message_stop"}`),
117+
}
118+
119+
outputs := translateClaudeResponsesStreamThroughRegistry(chunks)
120+
121+
counts := map[string]int{}
122+
var outputTextDone gjson.Result
123+
var completed gjson.Result
124+
for _, output := range outputs {
125+
event, data := parseClaudeResponsesSSEEvent(t, output)
126+
counts[event]++
127+
if event == "response.output_text.done" {
128+
outputTextDone = data
129+
}
130+
if event == "response.completed" {
131+
completed = data
132+
}
133+
if strings.HasPrefix(event, "content_block_") || event == "message_delta" {
134+
t.Fatalf("unexpected anthropic-native event leaked: %s", event)
135+
}
136+
}
137+
138+
if counts["response.output_item.added"] != 1 {
139+
t.Fatalf("response.output_item.added count = %d, want 1", counts["response.output_item.added"])
140+
}
141+
if counts["response.content_part.added"] != 1 {
142+
t.Fatalf("response.content_part.added count = %d, want 1", counts["response.content_part.added"])
143+
}
144+
if counts["response.output_text.done"] != 1 {
145+
t.Fatalf("response.output_text.done count = %d, want 1", counts["response.output_text.done"])
146+
}
147+
if counts["response.content_part.done"] != 1 {
148+
t.Fatalf("response.content_part.done count = %d, want 1", counts["response.content_part.done"])
149+
}
150+
if counts["response.output_item.done"] != 1 {
151+
t.Fatalf("response.output_item.done count = %d, want 1", counts["response.output_item.done"])
152+
}
153+
if counts["response.function_call_arguments.delta"] != 0 {
154+
t.Fatalf("response.function_call_arguments.delta count = %d, want 0", counts["response.function_call_arguments.delta"])
155+
}
156+
157+
wantText := "**对比竞品**\n- Qwen 3.7 Max leads."
158+
if got := outputTextDone.Get("text").String(); got != wantText {
159+
t.Fatalf("output_text.done text = %q, want %q", got, wantText)
160+
}
161+
if got := completed.Get("response.output.0.content.0.text").String(); got != wantText {
162+
t.Fatalf("completed message text = %q, want %q", got, wantText)
163+
}
164+
if got := completed.Get("response.output.0.content.0.annotations.0.type").String(); got != "web_search_result_location" {
165+
t.Fatalf("completed annotation type = %q", got)
166+
}
167+
}
168+
81169
func TestConvertClaudeResponseToOpenAIResponsesNonStream_ThinkingIncludesSignature(t *testing.T) {
82170
signature := "claude_sig_nonstream"
83171
raw := []byte(strings.Join([]string{

0 commit comments

Comments
 (0)