omniai 1.7.0 → 1.8.1

data/lib/omniai/chat/response.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # Used when processing everything at once.
+    class Response
+      # @return [Hash]
+      attr_accessor :data
+
+      # @param data [Hash]
+      # @param context [Context, nil]
+      def initialize(data:, context: nil)
+        @data = data
+        @context = context
+      end
+
+      # @return [Payload]
+      def completion
+        @completion ||= Payload.deserialize(@data, context: @context)
+      end
+
+      # @return [Usage, nil]
+      def usage
+        completion.usage
+      end
+
+      # @return [Array<Choice>]
+      def choices
+        completion.choices
+      end
+
+      # @return [Array<Message>]
+      def messages
+        completion.messages
+      end
+
+      # @param index [Integer]
+      # @return [Choice]
+      def choice(index: 0)
+        completion.choice(index:)
+      end
+
+      # @param index [Integer]
+      # @return [Message]
+      def message(index: 0)
+        completion.message(index:)
+      end
+
+      # @return [String]
+      def text
+        message.text
+      end
+
+      # @return [Boolean]
+      def text?
+        message.text?
+      end
+
+      # @return [Array<ToolCall>]
+      def tool_call_list
+        choice.tool_call_list
+      end
+
+      # @return [Boolean]
+      def tool_call_list?
+        tool_call_list&.any?
+      end
+    end
+  end
+end

data/lib/omniai/chat/stream.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # Used when streaming to process chunks of data.
+    class Stream
+      # @param logger [OmniAI::Client]
+      # @param body [HTTP::Response::Body]
+      # @param context [Context, nil]
+      def initialize(body:, logger: nil, context: nil)
+        @body = body
+        @logger = logger
+        @context = context
+      end
+
+      # @yield [payload]
+      # @yieldparam payload [OmniAI::Chat::Payload]
+      def stream!(&)
+        @body.each do |chunk|
+          parser.feed(chunk) do |type, data, id|
+            process!(type, data, id, &)
+          end
+        end
+      end
+
+      protected
+
+      # @param type [String]
+      # @param data [String]
+      # @param id [String]
+      def log(type, data, id)
+        arguments = [
+          ("type=#{type.inspect}" if type),
+          ("data=#{data.inspect}" if data),
+          ("id=#{id.inspect}" if id),
+        ].compact
+
+        @logger&.debug("Stream#process! #{arguments.join(' ')}")
+      end
+
+      # @param type [String]
+      # @param data [String]
+      # @param id [String]
+      #
+      # @yield [payload]
+      # @yieldparam payload [OmniAI::Chat::Payload]
+      def process!(type, data, id, &block)
+        log(type, data, id)
+
+        return if data.eql?('[DONE]')
+
+        block.call(Payload.deserialize(JSON.parse(data), context: @context))
+      end
+
+      # @return [EventStreamParser::Parser]
+      def parser
+        @parser ||= EventStreamParser::Parser.new
+      end
+    end
+  end
+end

data/lib/omniai/chat/text.rb

@@ -25,7 +25,7 @@ module OmniAI
 
       # @param data [Hash]
       def self.deserialize(data, context: nil)
-        deserialize = context&.deserializers&.[](:text)
+        deserialize = context&.deserializer(:text)
         return deserialize.call(data, context:) if deserialize
 
         new(data['text'])
@@ -35,7 +35,7 @@ module OmniAI
       #
       # @return [Hash]
       def serialize(context: nil)
-        serializer = context&.serializers&.[](:text)
+        serializer = context&.serializer(:text)
         return serializer.call(self, context:) if serializer
 
         { type: 'text', text: @text }

data/lib/omniai/chat/tool_call.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # A tool-call that includes an ID / function.
+    class ToolCall
+      # @return [String]
+      attr_accessor :id
+
+      # @return [Function]
+      attr_accessor :function
+
+      # @param id [String]
+      # @param function [Function]
+      def initialize(id:, function:)
+        @id = id
+        @function = function
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} id=#{id.inspect} function=#{function.inspect}>"
+      end
+
+      # @param data [Hash]
+      # @param context [Context] optional
+      #
+      # @return [Function]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:tool_call)
+        return deserialize.call(data, context:) if deserialize
+
+        id = data['id']
+        function = Function.deserialize(data['function'], context:)
+
+        new(id:, function:)
+      end
+
+      # @param context [Context] optional
+      #
+      # @return [Hash]
+      def serialize(context: nil)
+        serializer = context&.serializer(:tool_call)
+        return serializer.call(self, context:) if serializer
+
+        {
+          id: @id,
+          type: 'function',
+          function: @function.serialize(context:),
+        }
+      end
+    end
+  end
+end

data/lib/omniai/chat/tool_call_message.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # A specific message that contains the result of a tool call.
+    class ToolCallMessage < Message
+      # @return [String]
+      attr_accessor :tool_call_id
+
+      # @param content [String]
+      # @param tool_call_id [String]
+      def initialize(content:, tool_call_id:, role: OmniAI::Chat::Role::TOOL)
+        super(content:, role:)
+        @tool_call_id = tool_call_id
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} content=#{content.inspect} tool_call_id=#{tool_call_id.inspect}>"
+      end
+
+      # Usage:
+      #
+      #   ToolCall.deserialize({ 'role' => 'tool', content: '{ 'temperature': 0 }' }) # => #<ToolCall ...>
+      #
+      # @param data [Hash]
+      # @param context [Context] optional
+      #
+      # @return [ToolMessage]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:tool_message)
+        return deserialize.call(data, context:) if deserialize
+
+        role = data['role']
+        content = JSON.parse(data['content'])
+        tool_call_id = data['tool_call_id']
+
+        new(role:, content:, tool_call_id:)
+      end
+
+      # Usage:
+      #
+      #   message.serialize # => { role: :user, content: 'Hello!' }
+      #   message.serialize # => { role: :user, content: [{ type: 'text', text: 'Hello!' }] }
+      #
+      # @param context [Context] optional
+      #
+      # @return [Hash]
+      def serialize(context: nil)
+        serializer = context&.serializer(:tool_message)
+        return serializer.call(self, context:) if serializer
+
+        role = @role
+        content = JSON.generate(@content)
+        tool_call_id = @tool_call_id
+
+        { role:, content:, tool_call_id: }
+      end
+    end
+  end
+end

data/lib/omniai/chat/tool_call_result.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # The result of a tool call.
+    class ToolCallResult
+      # @return [Object]
+      attr_accessor :content
+
+      # @return [ToolCall]
+      attr_accessor :tool_call_id
+
+      # @param content [Object]
+      # @param tool_call [ToolCall]
+      def initialize(content:, tool_call_id:)
+        @content = content
+        @tool_call_id = tool_call_id
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} content=#{content.inspect} tool_call_id=#{tool_call_id.inspect}>"
+      end
+
+      # @param context [Context] optional
+      # @return [Hash]
+      def serialize(context: nil)
+        serializer = context&.serializer(:tool_call_result)
+        return serializer.call(self, context:) if serializer
+
+        content = JSON.generate(@content)
+        tool_call_id = @tool_call_id
+
+        { content:, tool_call_id: }
+      end
+
+      # @param data [Hash]
+      # @param context [Context] optional
+      # @return [ToolCallResult]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:tool_call_result)
+        return deserialize.call(data, context:) if deserialize
+
+        content = JSON.parse(data['content'])
+        tool_call_id = data['tool_call_id']
+
+        new(content:, tool_call_id:)
+      end
+    end
+  end
+end

data/lib/omniai/chat/url.rb

@@ -28,7 +28,7 @@ module OmniAI
 
       # @param data [Hash]
       def self.deserialize(data, context: nil)
-        deserialize = context&.deserializers&.[](:url)
+        deserialize = context&.deserializer(:url)
         return deserialize.call(data, context:) if deserialize
 
         type = /(?<type>\w+)_url/.match(data['type'])[:type]
@@ -45,7 +45,7 @@ module OmniAI
           content = fetch!
           Text.new("<file>#{filename}: #{content}</file>").serialize(context:)
         else
-          serializer = context&.serializers&.[](:url)
+          serializer = context&.serializer(:url)
           return serializer.call(self, context:) if serializer
 
           {

data/lib/omniai/chat/usage.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Chat
+    # The usage of a chat in terms of tokens (input / output / total).
+    class Usage
+      # @return [Integer]
+      attr_accessor :input_tokens
+
+      # @return [Integer]
+      attr_accessor :output_tokens
+
+      # @return [Integer]
+      attr_accessor :total_tokens
+
+      # @param input_tokens [Integer]
+      # @param output_tokens [Integer]
+      # @param total_tokens [Integer]
+      def initialize(input_tokens:, output_tokens:, total_tokens:)
+        @input_tokens = input_tokens
+        @output_tokens = output_tokens
+        @total_tokens = total_tokens
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} input_tokens=#{input_tokens} output_tokens=#{output_tokens} total_tokens=#{total_tokens}>"
+      end
+
+      # @param data [Hash]
+      # @param context [OmniAI::Context] optional
+      #
+      # @return [OmniAI::Chat::Usage]
+      def self.deserialize(data, context: nil)
+        deserialize = context&.deserializer(:usage)
+        return deserialize.call(data, context:) if deserialize
+
+        input_tokens = data['input_tokens']
+        output_tokens = data['output_tokens']
+        total_tokens = data['total_tokens']
+
+        new(input_tokens:, output_tokens:, total_tokens:)
+      end
+
+      # @param context [OmniAI::Context] optional
+      #
+      # @return [Hash]
+      def serialize(context: nil)
+        serialize = context&.serializer(:usage)
+        return serialize.call(self, context:) if serialize
+
+        {
+          input_tokens:,
+          output_tokens:,
+          total_tokens:,
+        }
+      end
+    end
+  end
+end

data/lib/omniai/chat.rb

@@ -27,14 +27,23 @@ module OmniAI
   class Chat
     JSON_PROMPT = 'Respond with valid JSON. Do not include any non-JSON in the response.'
 
-    # An error raised when a chat makes a tool-call for a tool that cannot be found.
-    class ToolCallLookupError < Error
-      def initialize(tool_call)
-        super("missing tool for tool_call=#{tool_call.inspect}")
+    # An error raised for tool-call issues.
+    class ToolCallError < Error
+      # @param tool_call [OmniAI::Chat::ToolCall]
+      # @param message [String]
+      def initialize(tool_call:, message:)
+        super(message)
         @tool_call = tool_call
       end
     end
 
+    # An error raised when a tool-call is missing.
+    class ToolCallMissingError < ToolCallError
+      def initialize(tool_call:)
+        super(tool_call:, message: "missing tool for tool_call=#{tool_call.inspect}")
+      end
+    end
+
     module Role
       ASSISTANT = 'assistant'
       USER = 'user'
@@ -87,10 +96,22 @@ module OmniAI
 
     protected
 
+    # Override  to provide an context for serializers / deserializes for a provider.
+    #
+    # @return [Context, nil]
+    def context
+      nil
+    end
+
+    # @return [Logger, nil]
+    def logger
+      @client.logger
+    end
+
     # Used to spawn another chat with the same configuration using different messages.
     #
     # @param prompt [OmniAI::Chat::Prompt]
-    # @return [OmniAI::Chat::Prompt]
+    # @return [OmniAI::Chat]
     def spawn!(prompt)
       self.class.new(
         prompt,
@@ -100,7 +121,7 @@ module OmniAI
         stream: @stream,
         tools: @tools,
         format: @format
-      ).process!
+      )
     end
 
     # @return [Hash]
@@ -114,7 +135,7 @@ module OmniAI
     end
 
     # @param response [HTTP::Response]
-    # @return [OmniAI::Chat::Response::Completion]
+    # @return [OmniAI::Chat::Response]
     def parse!(response:)
       if @stream
         stream!(response:)
@@ -124,31 +145,32 @@ module OmniAI
     end
 
     # @param response [HTTP::Response]
-    # @return [OmniAI::Chat::Response::Completion]
+    # @return [OmniAI::Chat::Response]
     def complete!(response:)
-      completion = self.class::Response::Completion.new(data: response.parse)
-
-      if @tools && completion.tool_call_list.any?
-        spawn!([
-          *@prompt.serialize,
-          *completion.choices.map(&:message).map(&:data),
-          *(completion.tool_call_list.map { |tool_call| execute_tool_call(tool_call) }),
-        ])
+      completion = self.class::Response.new(data: response.parse, context:)
+
+      if @tools && completion.tool_call_list?
+        spawn!(
+          @prompt.dup.tap do |prompt|
+            prompt.messages += completion.messages
+            prompt.messages += build_tool_call_messages(completion.tool_call_list)
+          end
+        ).process!
       else
         completion
       end
     end
 
     # @param response [HTTP::Response]
-    # @return [OmniAI::Chat::Response::Stream]
+    # @return [OmniAI::Chat::Stream]
     def stream!(response:)
       raise Error, "#{self.class.name}#stream! unstreamable" unless @stream
 
-      self.class::Response::Stream.new(response:).stream! do |chunk|
+      self.class::Stream.new(body: response.body, logger:, context:).stream! do |chunk|
         case @stream
         when IO, StringIO
-          if chunk.content?
-            @stream << chunk.content
+          if chunk.text
+            @stream << chunk.text
             @stream.flush
           end
         else @stream.call(chunk)
@@ -160,31 +182,36 @@ module OmniAI
 
     # @return [HTTP::Response]
     def request!
+      logger&.debug("Chat#request! payload=#{payload.inspect}")
+
       @client
         .connection
         .accept(:json)
         .post(path, json: payload)
     end
 
+    # @param tool_call_list [Array<OmniAI::Chat::ToolCall>]
+    # @return [Array<Message>]
+    def build_tool_call_messages(tool_call_list)
+      tool_call_list.map do |tool_call|
+        content = execute_tool_call(tool_call)
+        ToolCallMessage.new(content:, tool_call_id: tool_call.id)
+      end
+    end
+
+    # @raise [ToolCallError]
     # @param tool_call [OmniAI::Chat::ToolCall]
+    # @return [ToolCallResult]
     def execute_tool_call(tool_call)
-      function = tool_call.function
+      logger&.debug("Chat#execute_tool_call tool_call=#{tool_call.inspect}")
 
-      tool = @tools.find { |entry| function.name == entry.name } || raise(ToolCallLookupError, tool_call)
-      result = tool.call(function.arguments)
+      function = tool_call.function
+      tool = @tools.find { |entry| function.name == entry.name } || raise(ToolCallMissingError, tool_call)
+      content = tool.call(function.arguments)
 
-      prepare_tool_call_message(tool_call:, content: result)
-    end
+      logger&.debug("Chat#execute_tool_call content=#{content.inspect}")
 
-    # @param tool_call [OmniAI::Chat::ToolCall]
-    # @param content [String]
-    def prepare_tool_call_message(tool_call:, content:)
-      {
-        role: Role::TOOL,
-        name: tool_call.function.name,
-        tool_call_id: tool_call.id,
-        content:,
-      }
+      content
     end
   end
 end

data/lib/omniai/context.rb

@@ -39,5 +39,17 @@ module OmniAI
       @serializers = {}
       @deserializers = {}
     end
+
+    # @param name [Symbol]
+    # @return [Proc, nil]
+    def serializer(name)
+      @serializers[name]
+    end
+
+    # @param name [Symbol]
+    # @return [Proc, nil]
+    def deserializer(name)
+      @deserializers[name]
+    end
   end
 end

data/lib/omniai/embed/response.rb

@@ -22,7 +22,7 @@ module OmniAI
       # @return [Usage]
       def usage
         @usage ||= begin
-          deserializer = @context&.deserializers&.[](:usage)
+          deserializer = @context&.deserializer(:usage)
 
           if deserializer
             deserializer.call(@data, context: @context)
@@ -45,7 +45,7 @@ module OmniAI
       # @return [Array<Array<Float>>]
       def embeddings
         @embeddings ||= begin
-          deserializer = @context&.deserializers&.[](:embeddings)
+          deserializer = @context&.deserializer(:embeddings)
 
           if deserializer
             deserializer.call(@data, context: @context)

data/lib/omniai/tool/parameters.rb

@@ -24,10 +24,10 @@ module OmniAI
       end
 
       # @return [Hash]
-      def prepare
+      def serialize
         {
           type: @type,
-          properties: @properties.transform_values(&:prepare),
+          properties: @properties.transform_values(&:serialize),
           required: @required,
         }.compact
       end

data/lib/omniai/tool/property.rb

@@ -60,7 +60,7 @@ module OmniAI
       end
 
       # @example
-      #   property.prepare
+      #   property.serialize
       #   # {
       #   #   type: 'string',
       #   #   description: 'The unit (e.g. "fahrenheit" or "celsius").'
@@ -68,7 +68,7 @@ module OmniAI
       #   # }
       #
       # @return [Hash]
-      def prepare
+      def serialize
         {
           type: @type,
           description: @description,

data/lib/omniai/tool.rb

@@ -44,7 +44,7 @@ module OmniAI
     end
 
     # @example
-    #   tool.prepare
+    #   tool.serialize
     #   # {
     #   #   type: 'function',
     #   #   function: {
@@ -53,21 +53,28 @@ module OmniAI
     #   #     parameters: {
     #   #       type: 'object',
     #   #       properties: {
-    #   #         n: OmniAI::Tool::Property.integer(description: 'The nth Fibonacci number to calculate')
+    #   #         n: {
+    #   #           description: 'The nth Fibonacci number to calculate')
+    #   #           type: 'integer'
+    #   #         }
     #   #       },
     #   #       required: ['n']
     #   #     }
     #   #   }
     #   # }
     #
+    # @param context [Context] optional
     # @return [Hash]
-    def prepare
+    def serialize(context: nil)
+      serialize = context&.serializer(:tool)
+      return serialize.call(self, context:) if serialize
+
       {
         type: 'function',
         function: {
           name: @name,
           description: @description,
-          parameters: @parameters&.prepare,
+          parameters: @parameters.is_a?(Tool::Parameters) ? @parameters.serialize : @parameters,
         }.compact,
       }
     end
@@ -79,7 +86,7 @@ module OmniAI
     # @param args [Hash]
     # @return [String]
     def call(args = {})
-      @function.call(**(@parameters ? @parameters.parse(args) : args))
+      @function.call(**(@parameters.is_a?(Tool::Parameters) ? @parameters.parse(args) : args))
     end
   end
 end

data/lib/omniai/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module OmniAI
-  VERSION = '1.7.0'
+  VERSION = '1.8.1'
 end