From b16d6d504ee70d78a152117fec928b74ef2cb882 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 20 Aug 2026 17:48:52 +0000 Subject: [PATCH 1/4] fix(streaming): keep indexed snapshots compact --- .../streaming/chat_completion_stream.rb | 55 +++++++-- .../helpers/chat_completion_stream_test.rb | 116 ++++++++++++++++++ 2 files changed, 159 insertions(+), 12 deletions(-) create mode 100644 test/openai/helpers/chat_completion_stream_test.rb diff --git a/lib/openai/helpers/streaming/chat_completion_stream.rb b/lib/openai/helpers/streaming/chat_completion_stream.rb index 874ef7153..40e969e52 100644 --- a/lib/openai/helpers/streaming/chat_completion_stream.rb +++ b/lib/openai/helpers/streaming/chat_completion_stream.rb @@ -78,7 +78,7 @@ class ChatCompletionStreamState def initialize(response_format: nil, input_tools: nil) @current_completion_snapshot = nil - @choice_event_states = [] + @choice_event_states = {} @input_tools = Array(input_tools) @response_format = response_format @rich_response_format = response_format.is_a?(Class) ? response_format : nil @@ -105,6 +105,7 @@ def handle_chunk(chunk) return [] unless chunk.is_a?(OpenAI::Chat::ChatCompletionChunk) + validate_chunk_indices!(chunk) @current_completion_snapshot = accumulate_chunk(chunk) build_events(chunk: chunk, completion_snapshot: @current_completion_snapshot) end @@ -134,11 +135,12 @@ def accumulate_chunk(chunk) end def accumulate_choice!(choice, completion_snapshot) - choice_snapshot = completion_snapshot.choices[choice.index] + choice_snapshot = find_choice_snapshot(completion_snapshot, choice.index) if choice_snapshot.nil? choice_snapshot = create_new_choice_snapshot(choice) - completion_snapshot.choices[choice.index] = choice_snapshot + completion_snapshot.choices << choice_snapshot + completion_snapshot.choices.sort_by!(&:index) else update_existing_choice_snapshot(choice, choice_snapshot) end @@ -188,7 +190,7 @@ def build_events(chunk:, completion_snapshot:) def build_choice_events(choice, completion_snapshot) choice_state = get_choice_state(choice) - choice_snapshot = completion_snapshot.choices[choice.index] + choice_snapshot = find_choice_snapshot(completion_snapshot, choice.index) content_delta_events(choice, choice_snapshot) + tool_call_delta_events(choice, choice_snapshot) + @@ -233,8 +235,8 @@ def tool_call_delta_events(choice, choice_snapshot) return events unless tool_calls choice.delta.tool_calls.each do |tool_call_delta| - tool_call = tool_calls[tool_call_delta.index] - next unless tool_call.type == :function && tool_call_delta.function + tool_call = find_tool_call_snapshot(tool_calls, tool_call_delta.index) + next unless tool_call&.type == :function && tool_call_delta.function parsed_args = if tool_call.function.respond_to?(:parsed) tool_call.function.parsed @@ -294,7 +296,7 @@ def parse_tool_calls!(delta_tool_calls, snapshot_tool_calls) return unless delta_tool_calls && snapshot_tool_calls delta_tool_calls.each do |tool_call_chunk| - tool_call_snapshot = snapshot_tool_calls[tool_call_chunk.index] + tool_call_snapshot = find_tool_call_snapshot(snapshot_tool_calls, tool_call_chunk.index) next unless tool_call_snapshot&.type == :function input_tool = find_input_tool(tool_call_snapshot.function.name) @@ -402,6 +404,30 @@ def find_input_tool(name) @input_tools.find { |tool| tool.dig(:function, :name) == name } end + def find_choice_snapshot(completion_snapshot, index) + completion_snapshot.choices.find { |choice| choice.index == index } + end + + def find_tool_call_snapshot(tool_calls, index) + tool_calls.find { |tool_call| tool_call[:index] == index } + end + + def validate_chunk_indices!(chunk) + chunk.choices.each do |choice| + validate_stream_index!(choice.index, "choice") + + Array(choice.delta.tool_calls).each do |tool_call| + validate_stream_index!(tool_call.index, "tool call") + end + end + end + + def validate_stream_index!(index, kind) + return unless index.negative? + + raise StreamError.new("Invalid streamed #{kind} index #{index}: expected a non-negative integer") + end + def parse_function_tool_arguments(function) return nil unless function[:arguments] @@ -529,10 +555,15 @@ def accumulate_delta(acc, delta) ) end - if acc_value[index].nil? - acc_value[index] = delta_entry - elsif acc_value[index].is_a?(Hash) - acc_value[index] = accumulate_delta(acc_value[index], delta_entry) + acc_entry = acc_value.find do |entry| + entry.is_a?(Hash) && (entry[:index] || entry["index"]) == index + end + + if acc_entry.nil? + acc_value << delta_entry + acc_value.sort_by! { |entry| entry[:index] || entry["index"] } + else + accumulate_delta(acc_entry, delta_entry) end end else @@ -645,7 +676,7 @@ def tool_done_event(choice_snapshot, tool_index) @done_tool_calls.add(tool_index) - tool_call = choice_snapshot.message.tool_calls&.[](tool_index) + tool_call = choice_snapshot.message.tool_calls&.find { |entry| entry[:index] == tool_index } return nil unless tool_call&.type == :function parsed_args = parse_function_tool_arguments(tool_call.function) diff --git a/test/openai/helpers/chat_completion_stream_test.rb b/test/openai/helpers/chat_completion_stream_test.rb new file mode 100644 index 000000000..ba7c1f9b8 --- /dev/null +++ b/test/openai/helpers/chat_completion_stream_test.rb @@ -0,0 +1,116 @@ +# frozen_string_literal: true + +require_relative "../test_helper" + +class OpenAI::Test::ChatCompletionStreamTest < Minitest::Test + def test_stream_indices_do_not_create_sparse_snapshots + [127, 128, 1_000_000_000].each do |index| + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + events = stream_tool_call(state, choice_index: index, tool_index: index) + completion = state.get_final_completion + + assert_equal([index], completion.choices.map(&:index)) + + tool_calls = completion.choices.first.message.tool_calls + assert_equal([index], tool_calls.map { |tool_call| tool_call[:index] }) + assert_equal("{\"value\":true}", tool_calls.first.function.arguments) + + deltas = events.select { |event| event.type == :"tool_calls.function.arguments.delta" } + assert_equal([index, index], deltas.map(&:index)) + + done = events.find { |event| event.type == :"tool_calls.function.arguments.done" } + assert_equal(index, done.index) + end + end + + def test_stream_rejects_negative_indices_before_accumulation + invalid_chunks = [ + ["choice", build_chunk(choice_index: -1, delta: {role: :assistant})], + [ + "tool call", + build_chunk( + choice_index: 0, + delta: { + role: :assistant, + tool_calls: [tool_call_delta(index: -1, arguments: "")] + } + ) + ] + ] + + invalid_chunks.each do |kind, chunk| + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + error = assert_raises(OpenAI::Helpers::Streaming::StreamError) do + state.handle_chunk(chunk) + end + + assert_match("Invalid streamed #{kind} index -1", error.message) + assert_nil(state.current_completion_snapshot) + end + end + + def test_interleaved_indices_remain_ordered_and_independent + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk(build_chunk(choice_index: 1, delta: {role: :assistant, content: "one"})) + stream_tool_call(state, choice_index: 0, tool_index: 1, finish: false) + stream_tool_call(state, choice_index: 0, tool_index: 0, finish: false) + + completion = state.get_final_completion + assert_equal([0, 1], completion.choices.map(&:index)) + + tool_calls = completion.choices.first.message.tool_calls + assert_equal([0, 1], tool_calls.map { |tool_call| tool_call[:index] }) + assert_equal(["{\"value\":true}", "{\"value\":true}"], tool_calls.map { |tool_call| tool_call.function.arguments }) + end + + private + + def stream_tool_call(state, choice_index:, tool_index:, finish: true) + events = state.handle_chunk( + build_chunk( + choice_index: choice_index, + delta: { + role: :assistant, + tool_calls: [tool_call_delta(index: tool_index, arguments: "")] + } + ) + ) + events.concat( + state.handle_chunk( + build_chunk( + choice_index: choice_index, + delta: {tool_calls: [{index: tool_index, function: {arguments: "{\"value\":true}"}}]} + ) + ) + ) + if finish + events.concat( + state.handle_chunk(build_chunk(choice_index: choice_index, delta: {}, finish_reason: :tool_calls)) + ) + end + + events + end + + def tool_call_delta(index:, arguments:) + { + index: index, + id: "call_#{index}", + type: :function, + function: {name: "test", arguments: arguments} + } + end + + def build_chunk(choice_index:, delta:, finish_reason: nil) + OpenAI::Internal::Type::Converter.coerce( + OpenAI::Chat::ChatCompletionChunk, + { + id: "chatcmpl-index-test", + object: :"chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: [{index: choice_index, delta: delta, finish_reason: finish_reason}] + } + ) + end +end From b8c68801635016472f58f689ecb56401917ec8ee Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 20 Aug 2026 19:49:20 +0000 Subject: [PATCH 2/4] fix(streaming): index compact snapshots --- .../streaming/chat_completion_stream.rb | 96 ++++++++++++------- .../helpers/chat_completion_stream_test.rb | 70 +++++++++++++- 2 files changed, 132 insertions(+), 34 deletions(-) diff --git a/lib/openai/helpers/streaming/chat_completion_stream.rb b/lib/openai/helpers/streaming/chat_completion_stream.rb index 40e969e52..c265e8d27 100644 --- a/lib/openai/helpers/streaming/chat_completion_stream.rb +++ b/lib/openai/helpers/streaming/chat_completion_stream.rb @@ -78,6 +78,8 @@ class ChatCompletionStreamState def initialize(response_format: nil, input_tools: nil) @current_completion_snapshot = nil + @choice_snapshots_by_index = {} + @tool_call_snapshots_by_choice_index = {} @choice_event_states = {} @input_tools = Array(input_tools) @response_format = response_format @@ -119,15 +121,25 @@ def get_choice_state(choice) def accumulate_chunk(chunk) if @current_completion_snapshot.nil? - return convert_initial_chunk_into_snapshot(chunk) + completion_snapshot = convert_initial_chunk_into_snapshot(chunk) + completion_snapshot.choices.each do |choice_snapshot| + next if @choice_snapshots_by_index.key?(choice_snapshot.index) + + index_choice_snapshot!(choice_snapshot) + end + + return completion_snapshot end completion_snapshot = @current_completion_snapshot + appended_choice = false chunk.choices.each do |choice| - accumulate_choice!(choice, completion_snapshot) + appended_choice = true if accumulate_choice!(choice, completion_snapshot) end + completion_snapshot.choices.sort_by!(&:index) if appended_choice + completion_snapshot.usage = chunk.usage if chunk.usage completion_snapshot.system_fingerprint = chunk.system_fingerprint if chunk.system_fingerprint @@ -135,24 +147,27 @@ def accumulate_chunk(chunk) end def accumulate_choice!(choice, completion_snapshot) - choice_snapshot = find_choice_snapshot(completion_snapshot, choice.index) + choice_snapshot = @choice_snapshots_by_index[choice.index] + appended_choice = choice_snapshot.nil? - if choice_snapshot.nil? + if appended_choice choice_snapshot = create_new_choice_snapshot(choice) completion_snapshot.choices << choice_snapshot - completion_snapshot.choices.sort_by!(&:index) else update_existing_choice_snapshot(choice, choice_snapshot) end + tool_calls_by_index = index_choice_snapshot!(choice_snapshot) + if choice.finish_reason choice_snapshot.finish_reason = choice.finish_reason handle_finish_reason(choice.finish_reason, completion_snapshot) end - parse_tool_calls!(choice.delta.tool_calls, choice_snapshot.message.tool_calls) + parse_tool_calls!(choice.delta.tool_calls, tool_calls_by_index) accumulate_logprobs!(choice.logprobs, choice_snapshot) + appended_choice end def create_new_choice_snapshot(choice) @@ -182,23 +197,25 @@ def build_events(chunk:, completion_snapshot:) ) choice_events = chunk.choices.flat_map do |choice| - build_choice_events(choice, completion_snapshot) + build_choice_events(choice) end [chunk_event] + choice_events end - def build_choice_events(choice, completion_snapshot) + def build_choice_events(choice) choice_state = get_choice_state(choice) - choice_snapshot = find_choice_snapshot(completion_snapshot, choice.index) + choice_snapshot = @choice_snapshots_by_index.fetch(choice.index) + tool_calls_by_index = @tool_call_snapshots_by_choice_index.fetch(choice.index) content_delta_events(choice, choice_snapshot) + - tool_call_delta_events(choice, choice_snapshot) + + tool_call_delta_events(choice, tool_calls_by_index) + logprobs_delta_events(choice, choice_snapshot) + choice_state.get_done_events( choice_chunk: choice, choice_snapshot: choice_snapshot, - response_format: @response_format + response_format: @response_format, + tool_calls_by_index: tool_calls_by_index ) end @@ -227,15 +244,12 @@ def content_delta_events(choice, choice_snapshot) events end - def tool_call_delta_events(choice, choice_snapshot) + def tool_call_delta_events(choice, tool_calls_by_index) events = [] return events unless choice.delta.tool_calls - tool_calls = choice_snapshot.message.tool_calls - return events unless tool_calls - choice.delta.tool_calls.each do |tool_call_delta| - tool_call = find_tool_call_snapshot(tool_calls, tool_call_delta.index) + tool_call = tool_calls_by_index[tool_call_delta.index] next unless tool_call&.type == :function && tool_call_delta.function parsed_args = if tool_call.function.respond_to?(:parsed) @@ -292,11 +306,11 @@ def handle_finish_reason(finish_reason, completion_snapshot) end end - def parse_tool_calls!(delta_tool_calls, snapshot_tool_calls) - return unless delta_tool_calls && snapshot_tool_calls + def parse_tool_calls!(delta_tool_calls, tool_calls_by_index) + return unless delta_tool_calls delta_tool_calls.each do |tool_call_chunk| - tool_call_snapshot = find_tool_call_snapshot(snapshot_tool_calls, tool_call_chunk.index) + tool_call_snapshot = tool_calls_by_index[tool_call_chunk.index] next unless tool_call_snapshot&.type == :function input_tool = find_input_tool(tool_call_snapshot.function.name) @@ -404,12 +418,15 @@ def find_input_tool(name) @input_tools.find { |tool| tool.dig(:function, :name) == name } end - def find_choice_snapshot(completion_snapshot, index) - completion_snapshot.choices.find { |choice| choice.index == index } - end + def index_choice_snapshot!(choice_snapshot) + @choice_snapshots_by_index[choice_snapshot.index] = choice_snapshot + + tool_calls_by_index = {} + Array(choice_snapshot.message.tool_calls).each do |tool_call| + tool_calls_by_index[tool_call[:index]] ||= tool_call + end - def find_tool_call_snapshot(tool_calls, index) - tool_calls.find { |tool_call| tool_call[:index] == index } + @tool_call_snapshots_by_choice_index[choice_snapshot.index] = tool_calls_by_index end def validate_chunk_indices!(chunk) @@ -535,6 +552,16 @@ def accumulate_delta(acc, delta) next end + acc_entries_by_index = {} + acc_value.each do |entry| + next unless entry.is_a?(Hash) + + entry_index = entry[:index] || entry["index"] + acc_entries_by_index[entry_index] ||= entry if entry_index + end + + appended_entry = false + delta_value.each do |delta_entry| unless delta_entry.is_a?(Hash) raise( @@ -555,17 +582,20 @@ def accumulate_delta(acc, delta) ) end - acc_entry = acc_value.find do |entry| - entry.is_a?(Hash) && (entry[:index] || entry["index"]) == index - end + acc_entry = acc_entries_by_index[index] if acc_entry.nil? acc_value << delta_entry - acc_value.sort_by! { |entry| entry[:index] || entry["index"] } + acc_entries_by_index[index] = delta_entry + appended_entry = true else accumulate_delta(acc_entry, delta_entry) end end + + if appended_entry + acc_value.sort_by! { |entry| entry[:index] || entry["index"] } + end else acc[key] = acc_value end @@ -587,14 +617,14 @@ def initialize(input_tools:) @current_tool_call_index = nil end - def get_done_events(choice_chunk:, choice_snapshot:, response_format:) + def get_done_events(choice_chunk:, choice_snapshot:, response_format:, tool_calls_by_index:) events = [] if choice_snapshot.finish_reason events.concat(content_done_events(choice_snapshot, response_format)) if @current_tool_call_index && !@done_tool_calls.include?(@current_tool_call_index) - event = tool_done_event(choice_snapshot, @current_tool_call_index) + event = tool_done_event(tool_calls_by_index, @current_tool_call_index) events << event if event end end @@ -604,7 +634,7 @@ def get_done_events(choice_chunk:, choice_snapshot:, response_format:) events.concat(content_done_events(choice_snapshot, response_format)) if @current_tool_call_index - event = tool_done_event(choice_snapshot, @current_tool_call_index) + event = tool_done_event(tool_calls_by_index, @current_tool_call_index) events << event if event end end @@ -671,12 +701,12 @@ def logprobs_done_events(choice_snapshot) events end - def tool_done_event(choice_snapshot, tool_index) + def tool_done_event(tool_calls_by_index, tool_index) return nil if @done_tool_calls.include?(tool_index) @done_tool_calls.add(tool_index) - tool_call = choice_snapshot.message.tool_calls&.find { |entry| entry[:index] == tool_index } + tool_call = tool_calls_by_index[tool_index] return nil unless tool_call&.type == :function parsed_args = parse_function_tool_arguments(tool_call.function) diff --git a/test/openai/helpers/chat_completion_stream_test.rb b/test/openai/helpers/chat_completion_stream_test.rb index ba7c1f9b8..58ab0aa27 100644 --- a/test/openai/helpers/chat_completion_stream_test.rb +++ b/test/openai/helpers/chat_completion_stream_test.rb @@ -63,8 +63,70 @@ def test_interleaved_indices_remain_ordered_and_independent assert_equal(["{\"value\":true}", "{\"value\":true}"], tool_calls.map { |tool_call| tool_call.function.arguments }) end + def test_tool_call_snapshot_lookups_are_linear + tool_call_count = 128 + tool_calls = Array.new(tool_call_count) do |index| + tool_call_delta(index: index, arguments: "") + end + + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk( + build_chunk(choice_index: 0, delta: {role: :assistant, tool_calls: tool_calls}) + ) + + updates = Array.new(tool_call_count) do |index| + {index: index, function: {arguments: "{}"}} + end + + comparisons = count_integer_equality_calls do + state.handle_chunk( + build_chunk( + choice_index: 0, + delta: {tool_calls: updates}, + finish_reason: :tool_calls + ) + ) + end + + final_tool_calls = state.get_final_completion.choices.first.message.tool_calls + assert_equal(tool_call_count, final_tool_calls.length) + assert_equal(["{}"] * tool_call_count, final_tool_calls.map { |tool_call| tool_call.function.arguments }) + assert_operator(comparisons, :<=, tool_call_count * 32) + end + + def test_choice_snapshot_lookups_are_linear + choice_count = 128 + choices = Array.new(choice_count) do |index| + {index: index, delta: {role: :assistant, content: "choice #{index}"}, finish_reason: nil} + end + + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk(build_chunk_with_choices([])) + + comparisons = count_integer_equality_calls do + state.handle_chunk(build_chunk_with_choices(choices)) + end + + assert_equal((0...choice_count).to_a, state.get_final_completion.choices.map(&:index)) + assert_operator(comparisons, :<=, choice_count * 32) + end + private + def count_integer_equality_calls + comparisons = 0 + trace = TracePoint.new(:c_call) do |event| + if event.defined_class == Integer && event.method_id == :== + comparisons += 1 + end + end + + trace.enable { yield } + comparisons + ensure + trace&.disable + end + def stream_tool_call(state, choice_index:, tool_index:, finish: true) events = state.handle_chunk( build_chunk( @@ -102,6 +164,12 @@ def tool_call_delta(index:, arguments:) end def build_chunk(choice_index:, delta:, finish_reason: nil) + build_chunk_with_choices( + [{index: choice_index, delta: delta, finish_reason: finish_reason}] + ) + end + + def build_chunk_with_choices(choices) OpenAI::Internal::Type::Converter.coerce( OpenAI::Chat::ChatCompletionChunk, { @@ -109,7 +177,7 @@ def build_chunk(choice_index:, delta:, finish_reason: nil) object: :"chat.completion.chunk", created: 1, model: "gpt-4o-mini", - choices: [{index: choice_index, delta: delta, finish_reason: finish_reason}] + choices: choices } ) end From b168bf178eadb002a4aeafb6946882f6e39fd2d5 Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Thu, 20 Aug 2026 23:09:01 +0000 Subject: [PATCH 3/4] fix(streaming): avoid repeated choice sorting --- .../streaming/chat_completion_stream.rb | 34 ++++++++++---- .../helpers/chat_completion_stream_test.rb | 46 +++++++++++++++++++ 2 files changed, 72 insertions(+), 8 deletions(-) diff --git a/lib/openai/helpers/streaming/chat_completion_stream.rb b/lib/openai/helpers/streaming/chat_completion_stream.rb index c265e8d27..82507f88b 100644 --- a/lib/openai/helpers/streaming/chat_completion_stream.rb +++ b/lib/openai/helpers/streaming/chat_completion_stream.rb @@ -80,6 +80,7 @@ def initialize(response_format: nil, input_tools: nil) @current_completion_snapshot = nil @choice_snapshots_by_index = {} @tool_call_snapshots_by_choice_index = {} + @choice_snapshots_sorted = true @choice_event_states = {} @input_tools = Array(input_tools) @response_format = response_format @@ -122,6 +123,10 @@ def get_choice_state(choice) def accumulate_chunk(chunk) if @current_completion_snapshot.nil? completion_snapshot = convert_initial_chunk_into_snapshot(chunk) + @choice_snapshots_sorted = completion_snapshot.choices.each_cons(2).all? do |previous, current| + previous.index <= current.index + end + completion_snapshot.choices.each do |choice_snapshot| next if @choice_snapshots_by_index.key?(choice_snapshot.index) @@ -132,14 +137,11 @@ def accumulate_chunk(chunk) end completion_snapshot = @current_completion_snapshot - appended_choice = false chunk.choices.each do |choice| - appended_choice = true if accumulate_choice!(choice, completion_snapshot) + accumulate_choice!(choice, completion_snapshot) end - completion_snapshot.choices.sort_by!(&:index) if appended_choice - completion_snapshot.usage = chunk.usage if chunk.usage completion_snapshot.system_fingerprint = chunk.system_fingerprint if chunk.system_fingerprint @@ -148,11 +150,10 @@ def accumulate_chunk(chunk) def accumulate_choice!(choice, completion_snapshot) choice_snapshot = @choice_snapshots_by_index[choice.index] - appended_choice = choice_snapshot.nil? - if appended_choice + if choice_snapshot.nil? choice_snapshot = create_new_choice_snapshot(choice) - completion_snapshot.choices << choice_snapshot + insert_choice_snapshot!(completion_snapshot.choices, choice_snapshot) else update_existing_choice_snapshot(choice, choice_snapshot) end @@ -167,7 +168,6 @@ def accumulate_choice!(choice, completion_snapshot) parse_tool_calls!(choice.delta.tool_calls, tool_calls_by_index) accumulate_logprobs!(choice.logprobs, choice_snapshot) - appended_choice end def create_new_choice_snapshot(choice) @@ -429,6 +429,24 @@ def index_choice_snapshot!(choice_snapshot) @tool_call_snapshots_by_choice_index[choice_snapshot.index] = tool_calls_by_index end + def insert_choice_snapshot!(choice_snapshots, choice_snapshot) + unless @choice_snapshots_sorted + choice_snapshots.sort_by!(&:index) + @choice_snapshots_sorted = true + end + + if choice_snapshots.empty? || choice_snapshots.last.index < choice_snapshot.index + choice_snapshots << choice_snapshot + return + end + + insertion_index = choice_snapshots.bsearch_index do |existing_choice| + existing_choice.index > choice_snapshot.index + end + + choice_snapshots.insert(insertion_index || choice_snapshots.length, choice_snapshot) + end + def validate_chunk_indices!(chunk) chunk.choices.each do |choice| validate_stream_index!(choice.index, "choice") diff --git a/test/openai/helpers/chat_completion_stream_test.rb b/test/openai/helpers/chat_completion_stream_test.rb index 58ab0aa27..3db7fa510 100644 --- a/test/openai/helpers/chat_completion_stream_test.rb +++ b/test/openai/helpers/chat_completion_stream_test.rb @@ -111,8 +111,54 @@ def test_choice_snapshot_lookups_are_linear assert_operator(comparisons, :<=, choice_count * 32) end + def test_in_order_choice_appends_do_not_resort_existing_snapshots + choice_count = 128 + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk(build_chunk(choice_index: 0, delta: {role: :assistant})) + + index_reads = count_choice_snapshot_index_reads do + 1.upto(choice_count - 1) do |index| + state.handle_chunk(build_chunk(choice_index: index, delta: {role: :assistant})) + end + end + + assert_equal((0...choice_count).to_a, state.get_final_completion.choices.map(&:index)) + assert_operator(index_reads, :<=, choice_count * 16) + end + + def test_initial_choice_order_is_preserved_until_a_choice_is_appended + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk( + build_chunk_with_choices( + [ + {index: 2, delta: {role: :assistant}, finish_reason: nil}, + {index: 0, delta: {role: :assistant}, finish_reason: nil} + ] + ) + ) + + assert_equal([2, 0], state.current_completion_snapshot.choices.map(&:index)) + + state.handle_chunk(build_chunk(choice_index: 1, delta: {role: :assistant})) + assert_equal([0, 1, 2], state.current_completion_snapshot.choices.map(&:index)) + end + private + def count_choice_snapshot_index_reads + reads = 0 + trace = TracePoint.new(:call) do |event| + if event.method_id == :index && event.self.is_a?(OpenAI::Models::Chat::ParsedChoice) + reads += 1 + end + end + + trace.enable { yield } + reads + ensure + trace&.disable + end + def count_integer_equality_calls comparisons = 0 trace = TracePoint.new(:c_call) do |event| From 32bf0169bb08e89687a41f241dcf01859f8e30ad Mon Sep 17 00:00:00 2001 From: Justin Beckwith Date: Fri, 21 Aug 2026 01:21:49 +0000 Subject: [PATCH 4/4] fix(streaming): avoid repeated tool call sorting --- .../streaming/chat_completion_stream.rb | 12 +++++- .../helpers/chat_completion_stream_test.rb | 40 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/lib/openai/helpers/streaming/chat_completion_stream.rb b/lib/openai/helpers/streaming/chat_completion_stream.rb index 82507f88b..11a795aea 100644 --- a/lib/openai/helpers/streaming/chat_completion_stream.rb +++ b/lib/openai/helpers/streaming/chat_completion_stream.rb @@ -571,11 +571,17 @@ def accumulate_delta(acc, delta) end acc_entries_by_index = {} + entries_ordered = true + last_entry_index = nil acc_value.each do |entry| next unless entry.is_a?(Hash) entry_index = entry[:index] || entry["index"] - acc_entries_by_index[entry_index] ||= entry if entry_index + next unless entry_index + + entries_ordered &&= last_entry_index.nil? || last_entry_index <= entry_index + last_entry_index = entry_index + acc_entries_by_index[entry_index] ||= entry end appended_entry = false @@ -605,13 +611,15 @@ def accumulate_delta(acc, delta) if acc_entry.nil? acc_value << delta_entry acc_entries_by_index[index] = delta_entry + entries_ordered &&= last_entry_index.nil? || last_entry_index < index + last_entry_index = index appended_entry = true else accumulate_delta(acc_entry, delta_entry) end end - if appended_entry + if appended_entry && !entries_ordered acc_value.sort_by! { |entry| entry[:index] || entry["index"] } end else diff --git a/test/openai/helpers/chat_completion_stream_test.rb b/test/openai/helpers/chat_completion_stream_test.rb index 3db7fa510..a8ff82b42 100644 --- a/test/openai/helpers/chat_completion_stream_test.rb +++ b/test/openai/helpers/chat_completion_stream_test.rb @@ -94,6 +94,32 @@ def test_tool_call_snapshot_lookups_are_linear assert_operator(comparisons, :<=, tool_call_count * 32) end + def test_in_order_tool_call_appends_do_not_resort_existing_snapshots + tool_call_count = 128 + state = OpenAI::Helpers::Streaming::ChatCompletionStreamState.new + state.handle_chunk( + build_chunk( + choice_index: 0, + delta: {role: :assistant, tool_calls: [tool_call_delta(index: 0, arguments: "")]} + ) + ) + + sorts = count_array_sort_by_calls do + 1.upto(tool_call_count - 1) do |index| + state.handle_chunk( + build_chunk( + choice_index: 0, + delta: {tool_calls: [tool_call_delta(index: index, arguments: "")]} + ) + ) + end + end + + final_tool_calls = state.get_final_completion.choices.first.message.tool_calls + assert_equal((0...tool_call_count).to_a, final_tool_calls.map { |tool_call| tool_call[:index] }) + assert_equal(0, sorts) + end + def test_choice_snapshot_lookups_are_linear choice_count = 128 choices = Array.new(choice_count) do |index| @@ -173,6 +199,20 @@ def count_integer_equality_calls trace&.disable end + def count_array_sort_by_calls + sorts = 0 + trace = TracePoint.new(:c_call) do |event| + if event.defined_class == Array && event.method_id == :sort_by! + sorts += 1 + end + end + + trace.enable { yield } + sorts + ensure + trace&.disable + end + def stream_tool_call(state, choice_index:, tool_index:, finish: true) events = state.handle_chunk( build_chunk(