omniai 1.6.6 → 1.8.0

data/lib/omniai/context.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module OmniAI
+  # Used to handle the setup of serializer / deserializer required per provide (e.g. Anthropic / Google / etc).
+  #
+  # Usage:
+  #
+  #   OmniAI::Context.build do |context|
+  #     context.serializers[:prompt] = (prompt, context:) -> { ... }
+  #     context.serializers[:message] = (prompt, context:) -> { ... }
+  #     context.serializers[:file] = (prompt, context:) -> { ... }
+  #     context.serializers[:text] = (prompt, context:) -> { ... }
+  #     context.serializers[:url] = (prompt, context:) -> { ... }
+  #     context.deserializers[:prompt] = (data, context:) -> { Prompt.new(...) }
+  #     context.deserializers[:message] = (data, context:) -> { Message.new(...) }
+  #     context.deserializers[:file] = (data, context:) -> { File.new(...) }
+  #     context.deserializers[:text] = (data, context:) -> { Text.new(...) }
+  #     context.deserializers[:url] = (data, context:) -> { URL.new(...) }
+  #   end
+  class Context
+    # @return [Hash]
+    attr_accessor :serializers
+
+    # @return [Hash]
+    attr_reader :deserializers
+
+    # @yield [context]
+    # @yieldparam context [Context]
+    #
+    # @return [Context]
+    def self.build(&block)
+      new.tap do |context|
+        block&.call(context)
+      end
+    end
+
+    # @return [Context]
+    def initialize
+      @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

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Embed
+    # The response returned by the API.
+    class Response
+      # @return [Hash]
+      attr_accessor :data
+
+      # @param data [Hash]
+      # @param context [OmniAI::Context] optional
+      def initialize(data:, context: nil)
+        @data = data
+        @context = context
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name}>"
+      end
+
+      # @return [Usage]
+      def usage
+        @usage ||= begin
+          deserializer = @context&.deserializer(:usage)
+
+          if deserializer
+            deserializer.call(@data, context: @context)
+          else
+            prompt_tokens = @data.dig('usage', 'prompt_tokens')
+            total_tokens = @data.dig('usage', 'total_tokens')
+
+            Usage.new(prompt_tokens:, total_tokens:)
+          end
+        end
+      end
+
+      # @param index [Integer] optional
+      #
+      # @return [Array<Float>]
+      def embedding(index: 0)
+        embeddings[index]
+      end
+
+      # @return [Array<Array<Float>>]
+      def embeddings
+        @embeddings ||= begin
+          deserializer = @context&.deserializer(:embeddings)
+
+          if deserializer
+            deserializer.call(@data, context: @context)
+          else
+            @data['data'].map { |embedding| embedding['embedding'] }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/omniai/embed/usage.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module OmniAI
+  class Embed
+    # Token usage returned by the API.
+    class Usage
+      # @return [Integer]
+      attr_accessor :prompt_tokens
+
+      # @return [Integer]
+      attr_accessor :total_tokens
+
+      # @param prompt_tokens Integer
+      # @param total_tokens Integer
+      def initialize(prompt_tokens:, total_tokens:)
+        @prompt_tokens = prompt_tokens
+        @total_tokens = total_tokens
+      end
+
+      # @return [String]
+      def inspect
+        "#<#{self.class.name} prompt_tokens=#{@prompt_tokens} total_tokens=#{@total_tokens}>"
+      end
+    end
+  end
+end

data/lib/omniai/embed.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module OmniAI
+  # An abstract class that provides a consistent interface for processing embedding requests.
+  #
+  # Usage:
+  #
+  #   class OmniAI::OpenAI::Embed < OmniAI::Embed
+  #     module Model
+  #       SMALL = "text-embedding-3-small"
+  #       LARGE = "text-embedding-3-large"
+  #       ADA = "text-embedding-3-002"
+  #     end
+  #
+  #     protected
+  #
+  #     # @return [Hash]
+  #     def payload
+  #       { ... }
+  #     end
+  #
+  #     # @return [String]
+  #     def path
+  #       "..."
+  #     end
+  #   end
+  #
+  #   client.embed(input, model: "...")
+  class Embed
+    def self.process!(...)
+      new(...).process!
+    end
+
+    # @param input [String] required
+    # @param client [Client] the client
+    # @param model [String] required
+    #
+    # @return [Response]
+    def initialize(input, client:, model:)
+      @input = input
+      @client = client
+      @model = model
+    end
+
+    # @raise [Error]
+    # @return [Response]
+    def process!
+      response = request!
+      raise HTTPError, response.flush unless response.status.ok?
+
+      parse!(response:)
+    end
+
+    protected
+
+    # @param response [HTTP::Response]
+    # @return [Response]
+    def parse!(response:)
+      Response.new(data: response.parse)
+    end
+
+    # @return [HTTP::Response]
+    def request!
+      @client
+        .connection
+        .accept(:json)
+        .post(path, json: payload)
+    end
+
+    # @return [Hash]
+    def payload
+      raise NotImplementedError, "#{self.class.name}#payload undefined"
+    end
+
+    # @return [String]
+    def path
+      raise NotImplementedError, "#{self.class.name}#path undefined"
+    end
+  end
+end

data/lib/omniai/tool.rb

@@ -44,7 +44,7 @@ module OmniAI
     end
 
     # @example
-    #   tool.prepare
+    #   tool.serialize
     #   # {
     #   #   type: 'function',
     #   #   function: {
@@ -60,8 +60,12 @@ module OmniAI
     #   #   }
     #   # }
     #
+    # @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: {

data/lib/omniai/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: omniai
 version: !ruby/object:Gem::Version
-  version: 1.6.6
+  version: 1.8.0
 platform: ruby
 authors:
 - Kevin Sylvestre
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-07-19 00:00:00.000000000 Z
+date: 2024-08-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: event_stream_parser
@@ -68,33 +68,33 @@ files:
 - exe/omniai
 - lib/omniai.rb
 - lib/omniai/chat.rb
+- lib/omniai/chat/choice.rb
 - lib/omniai/chat/content.rb
-- lib/omniai/chat/context.rb
 - lib/omniai/chat/file.rb
+- lib/omniai/chat/function.rb
 - lib/omniai/chat/media.rb
 - lib/omniai/chat/message.rb
+- lib/omniai/chat/message/builder.rb
+- lib/omniai/chat/payload.rb
 - lib/omniai/chat/prompt.rb
-- lib/omniai/chat/response/choice.rb
-- lib/omniai/chat/response/chunk.rb
-- lib/omniai/chat/response/completion.rb
-- lib/omniai/chat/response/delta.rb
-- lib/omniai/chat/response/delta_choice.rb
-- lib/omniai/chat/response/function.rb
-- lib/omniai/chat/response/message.rb
-- lib/omniai/chat/response/message_choice.rb
-- lib/omniai/chat/response/part.rb
-- lib/omniai/chat/response/payload.rb
-- lib/omniai/chat/response/resource.rb
-- lib/omniai/chat/response/stream.rb
-- lib/omniai/chat/response/tool_call.rb
-- lib/omniai/chat/response/usage.rb
+- lib/omniai/chat/response.rb
+- lib/omniai/chat/stream.rb
 - lib/omniai/chat/text.rb
+- lib/omniai/chat/tool_call.rb
+- lib/omniai/chat/tool_call_message.rb
+- lib/omniai/chat/tool_call_result.rb
 - lib/omniai/chat/url.rb
+- lib/omniai/chat/usage.rb
 - lib/omniai/cli.rb
 - lib/omniai/cli/base_handler.rb
 - lib/omniai/cli/chat_handler.rb
+- lib/omniai/cli/embed_handler.rb
 - lib/omniai/client.rb
 - lib/omniai/config.rb
+- lib/omniai/context.rb
+- lib/omniai/embed.rb
+- lib/omniai/embed/response.rb
+- lib/omniai/embed/usage.rb
 - lib/omniai/instrumentation.rb
 - lib/omniai/speak.rb
 - lib/omniai/tool.rb

data/lib/omniai/chat/context.rb

@@ -1,42 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    # Used to handle the setup of serializer / deserializer methods for each type.
-    #
-    # Usage:
-    #
-    #   OmniAI::Chat::Context.build do |context|
-    #     context.serializers[:prompt] = (prompt, context:) -> { ... }
-    #     context.serializers[:message] = (prompt, context:) -> { ... }
-    #     context.serializers[:file] = (prompt, context:) -> { ... }
-    #     context.serializers[:text] = (prompt, context:) -> { ... }
-    #     context.serializers[:url] = (prompt, context:) -> { ... }
-    #     context.deserializers[:prompt] = (data, context:) -> { Prompt.new(...) }
-    #     context.deserializers[:message] = (data, context:) -> { Message.new(...) }
-    #     context.deserializers[:file] = (data, context:) -> { File.new(...) }
-    #     context.deserializers[:text] = (data, context:) -> { Text.new(...) }
-    #     context.deserializers[:url] = (data, context:) -> { URL.new(...) }
-    #   end
-    class Context
-      # @return [Hash]
-      attr_accessor :serializers
-
-      # @return [Hash]
-      attr_reader :deserializers
-
-      # @return [Context]
-      def self.build(&block)
-        new.tap do |context|
-          block&.call(context)
-        end
-      end
-
-      # @return [Context]
-      def initialize
-        @serializers = {}
-        @deserializers = {}
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/choice.rb

@@ -1,35 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # For use with MessageChoice or DeltaChoice.
-      class Choice < Resource
-        # @return [Integer]
-        def index
-          @data['index']
-        end
-
-        # @return [Part]
-        def part
-          raise NotImplementedError, "#{self.class.name}#part undefined"
-        end
-
-        # @return [ToolCallList]
-        def tool_call_list
-          part.tool_call_list
-        end
-
-        # @return [String, nil]
-        def content
-          part.content
-        end
-
-        # @return [Boolean]
-        def content?
-          !content.nil?
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/chunk.rb

@@ -1,15 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A chunk returned by the API.
-      class Chunk < Payload
-        # @return [Array<DeltaChoice>]
-        def choices
-          @choices ||= @data['choices'].map { |data| DeltaChoice.new(data:) }
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/completion.rb

@@ -1,15 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A completion returned by the API.
-      class Completion < Payload
-        # @return [Array<MessageChoice>]
-        def choices
-          @choices ||= @data['choices'].map { |data| MessageChoice.new(data:) }
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/delta.rb

@@ -1,11 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A delta returned by the API.
-      class Delta < Part
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/delta_choice.rb

@@ -1,25 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A delta choice returned by the API.
-      class DeltaChoice < Choice
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} index=#{index} delta=#{delta.inspect}>"
-        end
-
-        # @return [Delta]
-        def delta
-          @delta ||= Delta.new(data: @data['delta'])
-        end
-
-        # @return [Delta]
-        def part
-          delta
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/function.rb

@@ -1,25 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A function returned by the API.
-      class Function < Resource
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} name=#{name.inspect} arguments=#{arguments.inspect}>"
-        end
-
-        # @return [String]
-        def name
-          @data['name']
-        end
-
-        # @return [Hash, nil]
-        def arguments
-          JSON.parse(@data['arguments']) if @data['arguments']
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/message.rb

@@ -1,11 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A message returned by the API.
-      class Message < Part
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/message_choice.rb

@@ -1,25 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A choice returned by the API.
-      class MessageChoice < Choice
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} index=#{index} message=#{message.inspect}>"
-        end
-
-        # @return [Message]
-        def message
-          @message ||= Message.new(data: @data['message'])
-        end
-
-        # @return [Message]
-        def part
-          message
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/part.rb

@@ -1,38 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # Either a delta or message.
-      class Part < Resource
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} role=#{role.inspect} content=#{content.inspect}>"
-        end
-
-        # @return [String]
-        def role
-          @data['role'] || Role::USER
-        end
-
-        # @return [String, nil]
-        def content
-          @data['content']
-        end
-
-        # @return [Array<ToolCall>]
-        def tool_call_list
-          return [] unless @data['tool_calls']
-
-          @tool_call_list ||= @data['tool_calls'].map { |tool_call_data| ToolCall.new(data: tool_call_data) }
-        end
-
-        # @param index [Integer]
-        # @return [ToolCall, nil]
-        def tool_call(index: 0)
-          tool_call_list[index]
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/payload.rb

@@ -1,72 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A chunk or completion.
-      class Payload < Resource
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} id=#{id.inspect} choices=#{choices.inspect}>"
-        end
-
-        # @return [String]
-        def id
-          @data['id']
-        end
-
-        # @return [Time]
-        def created
-          Time.at(@data['created']) if @data['created']
-        end
-
-        # @return [Time]
-        def updated
-          Time.at(@data['updated']) if @data['updated']
-        end
-
-        # @return [String]
-        def model
-          @data['model']
-        end
-
-        # @return [Array<Choice>]
-        def choices
-          raise NotImplementedError, "#{self.class.name}#choices undefined"
-        end
-
-        # @param index [Integer]
-        # @return [DeltaChoice]
-        def choice(index: 0)
-          choices[index]
-        end
-
-        # @param index [Integer]
-        # @return [Part]
-        def part(index: 0)
-          choice(index:).part
-        end
-
-        # @return [Usage]
-        def usage
-          @usage ||= Usage.new(data: @data['usage']) if @data['usage']
-        end
-
-        # @return [String, nil]
-        def content
-          choice.content
-        end
-
-        # @return [Boolean]
-        def content?
-          choice.content?
-        end
-
-        # @return [Array<ToolCall>]
-        def tool_call_list
-          choice.tool_call_list
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/resource.rb

@@ -1,22 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A generic data to handle response.
-      class Resource
-        attr_accessor :data
-
-        # @param data [Hash]
-        def initialize(data:)
-          @data = data
-        end
-
-        # @return [String]
-        def inspect
-          "#<#{self.class.name} data=#{@data.inspect}>"
-        end
-      end
-    end
-  end
-end

data/lib/omniai/chat/response/stream.rb

@@ -1,27 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  class Chat
-    module Response
-      # A stream given when streaming.
-      class Stream
-        # @param response [HTTP::Response]
-        def initialize(response:)
-          @response = response
-          @parser = EventStreamParser::Parser.new
-        end
-
-        # @yield [OmniAI::Chat::Chunk]
-        def stream!
-          @response.body.each do |chunk|
-            @parser.feed(chunk) do |_, data|
-              next if data.eql?('[DONE]')
-
-              yield(Chunk.new(data: JSON.parse(data)))
-            end
-          end
-        end
-      end
-    end
-  end
-end