llms 0.1.0

data/lib/llms/parsers/anthropic_chat_response_stream_parser.rb

@@ -0,0 +1,184 @@
+require_relative '../stream/events'
+require_relative './partial_json_parser'
+
+module LLMs
+  module Parsers
+    class AnthropicChatResponseStreamParser
+      include PartialJsonParser
+
+      def initialize(emitter)
+        @emitter = emitter
+        @received_jsons = []
+
+        # Message metadata
+        @id = nil
+        @type = nil
+        @role = nil
+        @model = nil
+        @content = []
+        @stop_reason = nil
+        @stop_sequence = nil
+        @usage = nil
+      end
+
+      def full_response
+        {
+          'id' => @id,
+          'type' => @type,
+          'role' => @role,
+          'model' => @model,
+          'content' => @content,
+          'stop_reason' => @stop_reason,
+          'stop_sequence' => @stop_sequence,
+          'usage' => @usage
+        }
+      end
+
+      def handle_json(json)
+        @received_jsons << json
+
+        case json['type']
+        when 'message_start'
+          handle_message_start(json['message'])
+        when 'content_block_start'
+          handle_content_block_start(json)
+        when 'content_block_delta'
+          handle_content_block_delta(json)
+        when 'content_block_stop'
+          handle_content_block_stop(json)
+        when 'message_delta'
+          handle_message_delta(json)
+        when 'message_stop'
+          handle_message_stop
+        end
+      end
+
+      private
+
+      def handle_message_start(message)
+        @id = message['id']
+        @type = message['type']
+        @role = message['role']
+        @model = message['model']
+        @content = message['content'].dup
+        @usage = message['usage'].dup
+        @emitter.emit(:message_started, Stream::Events::MessageStarted.new(@id))
+        @emitter.emit(:usage_updated, Stream::Events::UsageUpdated.new(@id, @usage))
+      end
+
+      def handle_content_block_start(json)
+        index = json['index']
+        block = json['content_block'].dup
+        @content[index] = block
+
+        if block['type'] == 'text' && block['text'] && !block['text'].empty?
+          ## May never happen, but just in case
+          @emitter.emit(:text_delta, Stream::Events::TextDelta.new(block['text']))
+
+        elsif block['type'] == 'thinking' && block['thinking'] && !block['thinking'].empty?
+          ## Maybe never happens, but just in case
+          @emitter.emit(:thinking_delta, Stream::Events::ThinkingDelta.new(@id, block['thinking']))
+
+        elsif block['type'] == 'tool_use'
+          @emitter.emit(:tool_call_started, Stream::Events::ToolCallStarted.new(
+            @id,
+            block['id'],
+            index,
+            block['name'],
+            block['input'].dup
+          ))
+        end
+      end
+
+      def handle_content_block_delta(json)
+        index = json['index']
+        current_block = @content[index]
+
+        case json['delta']['type']
+        when 'text_delta'
+          text = json['delta']['text']
+          current_block['text'] ||= ''
+          current_block['text'] << text
+          @emitter.emit(:text_delta, Stream::Events::TextDelta.new(@id, text))
+
+        when 'thinking_delta'
+          thinking = json['delta']['thinking']
+          current_block['thinking'] ||= ''
+          current_block['thinking'] << thinking
+          @emitter.emit(:thinking_delta, Stream::Events::ThinkingDelta.new(@id, thinking))
+
+        when 'input_json_delta'
+          if current_block['type'] == 'tool_use'
+            handle_tool_use_delta(index, json['delta']['partial_json'])
+          end
+        end
+      end
+
+      def handle_tool_use_delta(index, partial_json)
+        current_block = @content[index]
+
+        if current_block['input'] == {}
+          current_block['input'] = ''
+        end
+        current_block['input'] << partial_json # This is an empty string first time
+
+        @emitter.emit(:tool_call_arguments_json_delta, Stream::Events::ToolCallArgumentsJsonDelta.new(
+          @id,
+          current_block['id'],
+          index,
+          partial_json
+        ))
+
+        parsed, _ = attempt_parse_json(current_block['input'])
+        if parsed
+          @emitter.emit(:tool_call_arguments_updated, Stream::Events::ToolCallArgumentsUpdated.new(
+            @id,
+            current_block['id'],
+            index,
+            parsed
+          ))
+        end
+      end
+
+      def handle_content_block_stop(json)
+        index = json['index']
+        current_block = @content[index]
+
+        if current_block['type'] == 'tool_use'
+          parse_tool_use_input(index)
+
+          @emitter.emit(:tool_call_completed, Stream::Events::ToolCallCompleted.new(
+            @id,
+            current_block['id'],
+            index,
+            current_block['name'],
+            current_block['input']
+          ))
+        end
+      end
+
+      def handle_message_stop
+        @emitter.emit(:message_completed, Stream::Events::MessageCompleted.new(@id, full_response))
+      end
+
+      def parse_tool_use_input(index)
+        input = @content[index]['input'].to_s.strip
+        @content[index]['input'] = input.empty? ? {} : JSON.parse(input)
+      end
+
+      def handle_message_delta(json)
+        @stop_reason = json['delta']['stop_reason']
+        @stop_sequence = json['delta']['stop_sequence']
+        update_usage(json['usage']) if json['usage']
+        @emitter.emit(:usage_updated, Stream::Events::UsageUpdated.new(@id, @usage))
+      end
+
+      def update_usage(usage)
+        usage.each do |key, value|
+          @usage[key] = value
+        end
+      end
+
+    end
+  end
+end

data/lib/llms/parsers/google_gemini_chat_response_stream_parser.rb

@@ -0,0 +1,128 @@
+require_relative './sse_chat_response_stream_parser'
+require_relative '../stream/events'
+
+module LLMs
+  module Parsers
+    class GoogleGeminiChatResponseStreamParser < SSEChatResponseStreamParser
+
+      attr_reader :current_message_id
+
+      def full_response
+        fr = {
+          'candidates' => @candidates,
+          'modelVersion' => @model_version,
+          'usageMetadata' => @usage_metadata
+        }
+
+        pp fr
+
+        fr
+      end
+
+      protected
+
+      def initialize_state
+        @candidates = []
+        @model_version = nil
+        @usage_metadata = nil
+        @current_message_id = nil
+        @tool_call_count = 0
+      end
+
+      def handle_json(json)
+        update_candidates(json['candidates']) if json['candidates']
+        update_model_version(json['modelVersion']) if json['modelVersion']
+        update_usage_metadata(json['usageMetadata']) if json['usageMetadata']
+      end
+
+      private
+
+      def update_candidates(candidates)
+        candidates.each_with_index do |candidate, index|
+          @candidates[index] ||= {}
+          current_candidate = @candidates[index]
+
+          if @current_message_id.nil?
+            @current_message_id = "gemini-#{Time.now.to_i}"
+            @emitter.emit(:message_started, Stream::Events::MessageStarted.new(@current_message_id))
+          end
+
+          if content = candidate['content']
+            update_candidate_content(current_candidate, content)
+          end
+
+          if finish_reason = candidate['finishReason']
+            current_candidate['finishReason'] = finish_reason
+            @emitter.emit(:message_completed, Stream::Events::MessageCompleted.new(@current_message_id, full_response))
+          end
+        end
+      end
+
+      def update_candidate_content(current_candidate, content)
+        current_candidate['content'] ||= {}
+
+        if parts = content['parts']
+          current_candidate['content']['parts'] ||= []
+          current_candidate['content']['parts'] += parts
+
+          parts.each do |part|
+            if part['text']
+              @emitter.emit(:text_delta, Stream::Events::TextDelta.new(@current_message_id, part['text']))
+            end
+
+            if part['functionCall']
+              tool_call_id = "tool_call#{@tool_call_count}"
+
+              @emitter.emit(:tool_call_started, Stream::Events::ToolCallStarted.new(
+                              @current_message_id,
+                              tool_call_id,
+                              @tool_call_count,
+                              part['functionCall']['name'],
+                              {}
+                            ))
+
+              args = part['functionCall']['args']
+
+              @emitter.emit(:tool_call_arguments_json_delta, Stream::Events::ToolCallArgumentsJsonDelta.new(
+                              @current_message_id,
+                              tool_call_id,
+                              @tool_call_count,
+                              JSON.dump(args)
+                            ))
+
+              @emitter.emit(:tool_call_arguments_updated, Stream::Events::ToolCallArgumentsUpdated.new(
+                              @current_message_id,
+                              tool_call_id,
+                              @tool_call_count,
+                              args
+                            ))
+
+              @emitter.emit(:tool_call_completed, Stream::Events::ToolCallCompleted.new(
+                              @current_message_id,
+                              tool_call_id,
+                              @tool_call_count,
+                              part['functionCall']['name'],
+                              args
+                            ))
+
+              @tool_call_count += 1
+            end
+          end
+        end
+
+        if role = content['role']
+          current_candidate['role'] = role
+        end
+      end
+
+      def update_model_version(version)
+        @model_version = version
+      end
+
+      def update_usage_metadata(metadata)
+        @usage_metadata = metadata
+        @emitter.emit(:usage_updated, Stream::Events::UsageUpdated.new(@current_message_id, @usage_metadata))
+      end
+    end
+  end
+end

data/lib/llms/parsers/open_ai_compatible_chat_response_stream_parser.rb

@@ -0,0 +1,170 @@
+require_relative './sse_chat_response_stream_parser'
+require_relative '../stream/events'
+
+module LLMs
+  module Parsers
+    class OpenAICompatibleChatResponseStreamParser < SSEChatResponseStreamParser
+
+      def full_response
+        # to match the format for non-streamed responses, all tool call arguments
+        # must be serialized back to a JSON string
+        converted_choices = @choices.map do |c|
+          dup_c = c.dup
+          dup_c['tool_calls']&.each do |tc|
+            tc['function']['arguments'] = JSON.dump(tc['function']['arguments'])
+          end
+          dup_c
+        end
+        {
+          'id' => @id,
+          'model' => @model,
+          # TODO I think this should be converted_choices, but providers are inconsistent in their response formats - switch to this after more testing
+          'choices' => @choices, #converted_choices, 
+          'usage' => @usage,
+          'created' => @created
+        }
+      end
+
+      protected
+
+      def initialize_state
+        @id = nil
+        @model = nil
+        @choices = []
+        @usage = nil
+        @created = nil
+      end
+
+      def handle_json(json)
+        update_id(json['id']) if json['id']
+        update_choices(json['choices']) if json['choices']
+        update_model(json['model']) if json['model']
+        update_usage(json['usage']) if json['usage']
+        update_created(json['created']) if json['created']
+      end
+
+      private
+
+      def update_id(id)
+        if @id.nil?
+          @id = id
+          @emitter.emit(:message_started, Stream::Events::MessageStarted.new(id))
+        elsif @id != id
+          puts "WARNING: id mismatch: #{@id} != #{id}"
+        end
+      end
+
+      def update_choices(choices)
+        choices.each_with_index do |choice, index|
+          @choices[index] ||= { 'message' => {} }
+          current_choice = @choices[index]['message']
+
+          if delta = choice['delta']
+            update_choice_delta(current_choice, delta)
+          end
+
+          if finish_reason = choice['finish_reason']
+            current_choice['finish_reason'] = finish_reason
+            @emitter.emit(:message_completed, Stream::Events::MessageCompleted.new(@id, full_response))
+          end
+        end
+      end
+
+      def update_choice_delta(current_choice, delta)
+        if role = delta['role']
+          current_choice['role'] = role
+        end
+
+        if content = delta['content']
+          current_choice['content'] ||= ''
+          current_choice['content'] += content
+          @emitter.emit(:text_delta, Stream::Events::TextDelta.new(@id, content))
+        end
+
+        if tool_calls = delta['tool_calls']
+          tool_calls = [tool_calls] unless tool_calls.is_a?(Array)
+          update_tool_calls(current_choice, tool_calls)
+        end
+      end
+
+      def update_tool_calls(current_choice, tool_calls)
+        current_choice['tool_calls'] ||= []
+
+        tool_calls.each do |tool_call|
+          tool_index = tool_call['index']
+
+          if new_call = current_choice['tool_calls'][tool_index].nil?
+            current_choice['tool_calls'][tool_index] = {
+              'id' => tool_call['id'],
+              'type' => tool_call['type'],
+              'function' => {
+                'name' => tool_call['function']['name'],
+                'arguments' => '' # Not this: ( tool_call['function']['arguments'].dup ) - since some providers append anyway
+              }
+            }
+          end
+
+          if new_call
+            @emitter.emit(:tool_call_started, Stream::Events::ToolCallStarted.new(
+              @id,
+              tool_call['id'],
+              tool_index,
+              tool_call['function']['name'],
+              {}
+            ))
+          end
+
+          if arguments = tool_call['function']['arguments']
+            current_tool_call = current_choice['tool_calls'][tool_index]
+            current_tool_call['function']['arguments'] += arguments
+
+            @emitter.emit(:tool_call_arguments_json_delta, Stream::Events::ToolCallArgumentsJsonDelta.new(
+              @id,
+              current_tool_call['id'],
+              tool_index,
+              arguments
+            ))
+
+            ## finish_reason"=>"tool_calls" <--- this is the finish reason when all tool calls completed
+            ## TODO use that instead?
+
+            parsed, corrected = attempt_parse_json(current_tool_call['function']['arguments'])
+
+            if parsed
+              if corrected
+                @emitter.emit(:tool_call_arguments_updated, Stream::Events::ToolCallArgumentsUpdated.new(
+                  @id,
+                  current_tool_call['id'],
+                  tool_index,
+                  parsed
+                ))
+              else
+                @emitter.emit(:tool_call_completed, Stream::Events::ToolCallCompleted.new(
+                  @id,
+                  current_tool_call['id'],
+                  tool_index,
+                  tool_call['function']['name'],
+                  parsed
+                ))
+              end
+            end
+
+          end
+        end
+      end
+
+      def update_model(model)
+        @model = model
+      end
+
+      def update_usage(usage)
+        @usage = usage
+        @emitter.emit(:usage_updated, Stream::Events::UsageUpdated.new(@id, @usage))
+      end
+
+      def update_created(created)
+        @created = created
+      end
+    end
+  end
+end

data/lib/llms/parsers/partial_json_parser.rb

@@ -0,0 +1,77 @@
+require 'json'
+
+module LLMs
+  module Parsers
+    module PartialJsonParser
+      def attempt_parse_json(json)
+        parsed = nil
+        corrected = false
+
+        begin
+          parsed = JSON.parse(json)
+        rescue JSON::ParserError
+          # Track unclosed delimiters
+          unclosed = []
+          in_string = false
+          escape_next = false
+
+          json.each_char.with_index do |char, i|
+            if escape_next
+              escape_next = false
+              next
+            end
+
+            case char
+            when '\\'
+              escape_next = true
+            when '"'
+              unless escape_next
+                if in_string
+                  if unclosed.last == :quote
+                    unclosed.pop
+                  end
+                  in_string = false
+                else
+                  unclosed.push(:quote)
+                  in_string = true
+                end
+              end
+            when '{'
+              unclosed.push(:brace) unless in_string
+            when '['
+              unclosed.push(:bracket) unless in_string
+            when '}'
+              if !in_string && unclosed.last == :brace
+                unclosed.pop
+              end
+            when ']'
+              if !in_string && unclosed.last == :bracket
+                unclosed.pop
+              end
+            end
+          end
+
+          # Build correction by closing delimiters in reverse order
+          correction = unclosed.reverse.map do |type|
+            case type
+            when :quote then '"'
+            when :brace then '}'
+            when :bracket then ']'
+            end
+          end.join
+
+          # Try parsing with correction
+          begin
+            corrected = true
+            corrected_json = json + correction
+            parsed = JSON.parse(corrected_json)
+          rescue JSON::ParserError
+            parsed = nil
+          end
+        end
+
+        [parsed, corrected]
+      end
+    end
+  end
+end

data/lib/llms/parsers/sse_chat_response_stream_parser.rb

@@ -0,0 +1,72 @@
+require 'json'
+require_relative './partial_json_parser'
+
+module LLMs
+  module Parsers
+    class SSEChatResponseStreamParser
+      include PartialJsonParser
+
+      def initialize(emitter)
+        @emitter = emitter
+        @buffer = ''
+        initialize_state
+      end
+
+      def add_data(data)
+        @buffer += data
+        process_buffer
+      end
+
+      def full_response
+        raise NotImplementedError, "Subclasses must implement full_response"
+      end
+
+      protected
+
+      def initialize_state
+        # Override in subclasses to initialize parser state
+      end
+
+      def process_buffer
+        while line = get_next_line
+          process_line(line)
+        end
+      end
+
+      def process_line(line)
+        if line.start_with?('data: ')
+          data = line[6..-1]
+          if data == '[DONE]'
+            handle_done
+          else
+            json = parse_line_data(data)
+            handle_json(json)
+          end
+        end
+      end
+
+      # Override in subclasses to rescue JSON parse errors if needed for the provider (shouldn't actually be needed for any?)
+      def parse_line_data(data)
+        JSON.parse(data)
+      end
+
+      def handle_json(json)
+        # Override in subclasses to handle JSON data
+      end
+
+      def handle_done
+        # Override in subclasses if needed
+      end
+
+      private
+
+      def get_next_line
+        if i = @buffer.index("\n")
+          line = @buffer[0...i].strip
+          @buffer = @buffer[(i + 1)..-1]
+          line
+        end
+      end
+    end
+  end
+end