diff --git a/lib/openai/helpers/streaming/chat_completion_stream.rb b/lib/openai/helpers/streaming/chat_completion_stream.rb index 874ef7153..11a795aea 100644 --- a/lib/openai/helpers/streaming/chat_completion_stream.rb +++ b/lib/openai/helpers/streaming/chat_completion_stream.rb @@ -78,7 +78,10 @@ class ChatCompletionStreamState def initialize(response_format: nil, input_tools: nil) @current_completion_snapshot = nil - @choice_event_states = [] + @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 @rich_response_format = response_format.is_a?(Class) ? response_format : nil @@ -105,6 +108,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 @@ -118,7 +122,18 @@ 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) + @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) + + index_choice_snapshot!(choice_snapshot) + end + + return completion_snapshot end completion_snapshot = @current_completion_snapshot @@ -134,21 +149,23 @@ def accumulate_chunk(chunk) end def accumulate_choice!(choice, completion_snapshot) - choice_snapshot = completion_snapshot.choices[choice.index] + choice_snapshot = @choice_snapshots_by_index[choice.index] if choice_snapshot.nil? choice_snapshot = create_new_choice_snapshot(choice) - completion_snapshot.choices[choice.index] = choice_snapshot + insert_choice_snapshot!(completion_snapshot.choices, choice_snapshot) 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) end @@ -180,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 = completion_snapshot.choices[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 @@ -225,16 +244,13 @@ 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 = tool_calls[tool_call_delta.index] - next unless tool_call.type == :function && tool_call_delta.function + 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) tool_call.function.parsed @@ -290,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 = 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) @@ -402,6 +418,51 @@ def find_input_tool(name) @input_tools.find { |tool| tool.dig(:function, :name) == name } 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 + + @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") + + 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] @@ -509,6 +570,22 @@ def accumulate_delta(acc, delta) next 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"] + 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 + delta_value.each do |delta_entry| unless delta_entry.is_a?(Hash) raise( @@ -529,12 +606,22 @@ 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_entries_by_index[index] + + 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 && !entries_ordered + acc_value.sort_by! { |entry| entry[:index] || entry["index"] } + end else acc[key] = acc_value end @@ -556,14 +643,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 @@ -573,7 +660,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 @@ -640,12 +727,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&.[](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 new file mode 100644 index 000000000..a8ff82b42 --- /dev/null +++ b/test/openai/helpers/chat_completion_stream_test.rb @@ -0,0 +1,270 @@ +# 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 + + 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_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| + {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 + + 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| + if event.defined_class == Integer && event.method_id == :== + comparisons += 1 + end + end + + trace.enable { yield } + comparisons + ensure + 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( + 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) + 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, + { + id: "chatcmpl-index-test", + object: :"chat.completion.chunk", + created: 1, + model: "gpt-4o-mini", + choices: choices + } + ) + end +end