Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 55 additions & 16 deletions anthropic/stream.go
Original file line number Diff line number Diff line change
Expand Up @@ -154,9 +154,11 @@ func handleContent(delta OpenAIMessage, state *StreamState) []AnthropicStreamEve
content = *delta.Content.Text
}

// Close thinking block if open
signature := reasoningOpaque(delta)

// Close thinking block if open, attaching signature to the same block.
if state.ThinkingBlockOpen {
events = append(events, closeThinkingBlock(state)...)
events = append(events, closeThinkingBlock(state, signature)...)
}

// Close tool block if open
Expand Down Expand Up @@ -198,9 +200,13 @@ func handleContent(delta OpenAIMessage, state *StreamState) []AnthropicStreamEve
func handleToolCalls(delta OpenAIMessage, state *StreamState) []AnthropicStreamEvent {
var events []AnthropicStreamEvent

// Close thinking block if open
signature := reasoningOpaque(delta)
consumedThinkingSignature := false

// Close thinking block if open, attaching signature to the same block.
if state.ThinkingBlockOpen {
events = append(events, closeThinkingBlock(state)...)
events = append(events, closeThinkingBlock(state, signature)...)
consumedThinkingSignature = signature != ""
}

// Handle reasoning opaque in tool calls
Expand All @@ -213,9 +219,11 @@ func handleToolCalls(delta OpenAIMessage, state *StreamState) []AnthropicStreamE
state.ContentBlockOpen = false
}

// Handle reasoning opaque
if delta.ReasoningOpaque != nil && *delta.ReasoningOpaque != "" {
events = append(events, handleReasoningOpaque(delta, state)...)
// Handle reasoning opaque only when there was no open thinking block to attach it to.
if !consumedThinkingSignature && !state.ThinkingBlockOpen {
if signature != "" {
events = append(events, handleReasoningOpaque(signature, state)...)
}
}

for _, toolCall := range delta.ToolCalls {
Expand Down Expand Up @@ -286,9 +294,13 @@ func handleToolCalls(delta OpenAIMessage, state *StreamState) []AnthropicStreamE
func handleFinish(choice OpenAIChunkChoice, chunk OpenAIChatCompletionChunk, state *StreamState) []AnthropicStreamEvent {
var events []AnthropicStreamEvent

// Close thinking block if open
signature := reasoningOpaque(choice.Delta)
consumedThinkingSignature := false

// Close thinking block if open, attaching signature to the same block.
if state.ThinkingBlockOpen {
events = append(events, closeThinkingBlock(state)...)
events = append(events, closeThinkingBlock(state, signature)...)
consumedThinkingSignature = signature != ""
}

// Close any open content block
Expand All @@ -304,9 +316,11 @@ func handleFinish(choice OpenAIChunkChoice, chunk OpenAIChatCompletionChunk, sta
state.ContentBlockOpen = false
state.ContentBlockIndex++

// Handle reasoning opaque for non-tool blocks
if !toolBlockOpen {
events = append(events, handleReasoningOpaque(choice.Delta, state)...)
// Handle reasoning opaque for non-tool blocks when there was no open thinking block.
if !toolBlockOpen && !consumedThinkingSignature {
if signature != "" {
events = append(events, handleReasoningOpaque(signature, state)...)
}
}
}

Expand Down Expand Up @@ -344,10 +358,10 @@ func handleFinish(choice OpenAIChunkChoice, chunk OpenAIChatCompletionChunk, sta
return events
}

func handleReasoningOpaque(delta OpenAIMessage, state *StreamState) []AnthropicStreamEvent {
func handleReasoningOpaque(signature string, state *StreamState) []AnthropicStreamEvent {
var events []AnthropicStreamEvent

if delta.ReasoningOpaque != nil && *delta.ReasoningOpaque != "" {
if signature != "" {
events = append(events,
AnthropicStreamEvent{
Type: "content_block_start",
Expand All @@ -357,12 +371,20 @@ func handleReasoningOpaque(delta OpenAIMessage, state *StreamState) []AnthropicS
Thinking: "",
},
},
AnthropicStreamEvent{
Type: "content_block_delta",
Index: intPtr(state.ContentBlockIndex),
Delta: &AnthropicContentDelta{
Type: "thinking_delta",
Thinking: "",
},
},
AnthropicStreamEvent{
Type: "content_block_delta",
Index: intPtr(state.ContentBlockIndex),
Delta: &AnthropicContentDelta{
Type: "signature_delta",
Signature: *delta.ReasoningOpaque,
Signature: signature,
},
},
AnthropicStreamEvent{
Expand All @@ -376,10 +398,20 @@ func handleReasoningOpaque(delta OpenAIMessage, state *StreamState) []AnthropicS
return events
}

func closeThinkingBlock(state *StreamState) []AnthropicStreamEvent {
func closeThinkingBlock(state *StreamState, signature string) []AnthropicStreamEvent {
var events []AnthropicStreamEvent

if state.ThinkingBlockOpen {
if signature != "" {
events = append(events, AnthropicStreamEvent{
Type: "content_block_delta",
Index: intPtr(state.ContentBlockIndex),
Delta: &AnthropicContentDelta{
Type: "signature_delta",
Signature: signature,
},
})
}
events = append(events,
AnthropicStreamEvent{
Type: "content_block_stop",
Expand All @@ -393,6 +425,13 @@ func closeThinkingBlock(state *StreamState) []AnthropicStreamEvent {
return events
}

func reasoningOpaque(delta OpenAIMessage) string {
if delta.ReasoningOpaque != nil {
return *delta.ReasoningOpaque
}
return ""
}

func isToolBlockOpen(state *StreamState) bool {
if !state.ContentBlockOpen {
return false
Expand Down
124 changes: 124 additions & 0 deletions anthropic/stream_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -361,6 +361,130 @@ func TestConvertOpenAIChunkToAnthropicEvents_Finish(t *testing.T) {
}
}

func TestConvertOpenAIChunkToAnthropicEvents_ThinkingFinishWithSignature(t *testing.T) {
state := NewStreamState()
state.MessageStartSent = true

thinkingChunk := OpenAIChatCompletionChunk{
ID: "msg_123",
Model: "claude-3-sonnet-20240229",
Choices: []OpenAIChunkChoice{{
Index: 0,
Delta: OpenAIMessage{
ReasoningText: stringPtr("Let me think about this..."),
},
}},
}
if _, err := ConvertOpenAIChunkToAnthropicEvents(thinkingChunk, state); err != nil {
t.Fatalf("Thinking conversion failed: %v", err)
}

finishChunk := OpenAIChatCompletionChunk{
ID: "msg_123",
Model: "claude-3-sonnet-20240229",
Choices: []OpenAIChunkChoice{{
Index: 0,
Delta: OpenAIMessage{
ReasoningOpaque: stringPtr("sig_123"),
},
FinishReason: "stop",
}},
Usage: &OpenAIUsage{PromptTokens: 10, CompletionTokens: 25},
}

events, err := ConvertOpenAIChunkToAnthropicEvents(finishChunk, state)
if err != nil {
t.Fatalf("Finish conversion failed: %v", err)
}

if len(events) != 4 {
t.Fatalf("Expected 4 events (signature + stop + delta + stop), got %d", len(events))
}

cd := contentDelta(t, events[0])
if cd.Type != "signature_delta" {
t.Fatalf("Expected first delta to be signature_delta, got %q", cd.Type)
}
if cd.Signature != "sig_123" {
t.Fatalf("Expected signature sig_123, got %q", cd.Signature)
}

if events[1].Type != "content_block_stop" {
t.Fatalf("Expected second event content_block_stop, got %q", events[1].Type)
}
if events[2].Type != "message_delta" {
t.Fatalf("Expected third event message_delta, got %q", events[2].Type)
}
if events[3].Type != "message_stop" {
t.Fatalf("Expected fourth event message_stop, got %q", events[3].Type)
}
}

func TestConvertOpenAIChunkToAnthropicEvents_ThinkingToToolCallWithSignature(t *testing.T) {
state := NewStreamState()
state.MessageStartSent = true

thinkingChunk := OpenAIChatCompletionChunk{
ID: "msg_123",
Model: "claude-3-sonnet-20240229",
Choices: []OpenAIChunkChoice{{
Index: 0,
Delta: OpenAIMessage{
ReasoningText: stringPtr("Analyzing..."),
},
}},
}
if _, err := ConvertOpenAIChunkToAnthropicEvents(thinkingChunk, state); err != nil {
t.Fatalf("Thinking conversion failed: %v", err)
}

toolChunk := OpenAIChatCompletionChunk{
ID: "msg_123",
Model: "claude-3-sonnet-20240229",
Choices: []OpenAIChunkChoice{{
Index: 0,
Delta: OpenAIMessage{
ReasoningOpaque: stringPtr("sig_456"),
ToolCalls: []OpenAIToolCall{{
Index: intPtr(0),
ID: "call_1",
Function: OpenAIToolCallFunction{
Name: "search",
},
}},
},
}},
}

events, err := ConvertOpenAIChunkToAnthropicEvents(toolChunk, state)
if err != nil {
t.Fatalf("Tool conversion failed: %v", err)
}

if len(events) != 3 {
t.Fatalf("Expected 3 events (signature + stop + tool_start), got %d", len(events))
}

cd := contentDelta(t, events[0])
if cd.Type != "signature_delta" {
t.Fatalf("Expected first delta to be signature_delta, got %q", cd.Type)
}
if cd.Signature != "sig_456" {
t.Fatalf("Expected signature sig_456, got %q", cd.Signature)
}

if events[1].Type != "content_block_stop" {
t.Fatalf("Expected second event content_block_stop, got %q", events[1].Type)
}
if events[2].Type != "content_block_start" || events[2].ContentBlock == nil || events[2].ContentBlock.Type != "tool_use" {
t.Fatalf("Expected third event tool_use content_block_start, got %#v", events[2])
}

if events[2].Index == nil || *events[2].Index != 1 {
t.Fatalf("Expected tool block index 1, got %v", events[2].Index)
}
}

func TestConvertOpenAIChunkToAnthropicEvents_ComplexFlow(t *testing.T) {
state := NewStreamState()

Expand Down