omniai-google 1.6.3 → 1.8.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b2626eef567a04e92b73eaf64a39dc274c2651b479810cd62a94be1a920cf25f
-  data.tar.gz: 8c968e5ca2328ce77e604f91c36607c22ef6572bf2e7633505316fdd6bab93db
+  metadata.gz: 07e400994624881af7790d3a209eb28030237a4684d319edc11ebb8dc7052f2c
+  data.tar.gz: e9ea6d986267343627c9cc8d061aaea673f4cbaba35a6fffdd113e49fe674278
 SHA512:
-  metadata.gz: c104b92ca09ae4c29285b2ac35e7badd23d2c54e4cb7189b0f1fb2fa1d0ad0cac4770e6c8ceeaaf5468225db52913458e6d4b6b94a0054f5c58735b01a78911d
-  data.tar.gz: a1037ddc9ccc58f84dfba6a094d0c7e1481dbba103e4e4d9f16dfbfb21def59fee61e4c463e96e3807c0763db4920f84934133b66090c9a8d575ac20c4225913
+  metadata.gz: cc5d0209e6024d8860f8aa619e8bdc76469863f314d66810d5b525674c417d4a79d57f5ad0ac404401214cf29e7606ee8a2669026609780f563a173351048b5a
+  data.tar.gz: c0958efb3190c4df770d49862fdf3370584c4533dc8925350fb194490f3da5cefd5de58c7e47fda14bf76803575e438505fdaf13d00f6693e363c6cc2ed42ef7

data/README.md

@@ -34,36 +34,27 @@ Global configuration is supported for the following options:
 OmniAI::Google.configure do |config|
   config.api_key = 'sk-...' # default: ENV['GOOGLE_API_KEY']
   config.host = '...' # default: 'https://generativelanguage.googleapis.com'
-  config.version = 'v1beta' # default: 'v1'
+  config.version = OmniAI::Google::Config::Version::BETA # either 'v1' or 'v1beta'
 end
 ```
 
 ### Chat
 
-A chat completion is generated by passing in prompts using any a variety of formats:
+A chat completion is generated by passing in a simple text prompt:
 
 ```ruby
 completion = client.chat('Tell me a joke!')
-completion.choice.message.content # 'Why did the chicken cross the road? To get to the other side.'
+completion.text # 'Why did the chicken cross the road? To get to the other side.'
 ```
 
-```ruby
-completion = client.chat({
-  role: OmniAI::Chat::Role::USER,
-  content: 'Is it wise to jump off a bridge?'
-})
-completion.choice.message.content # 'No.'
-```
+A chat completion may also be generated by using the prompt builder:
 
 ```ruby
-completion = client.chat([
-  {
-    role: OmniAI::Chat::Role::USER,
-    content: 'You are a helpful assistant.'
-  },
-  'What is the capital of Canada?',
-])
-completion.choice.message.content # 'The capital of Canada is Ottawa.'
+completion = client.chat do |prompt|
+  prompt.system('Your are an expert in geography.')
+  prompt.user('What is the capital of Canada?')
+end
+completion.text # 'The capital of Canada is Ottawa.'
 ```
 
 #### Model
@@ -72,7 +63,7 @@ completion.choice.message.content # 'The capital of Canada is Ottawa.'
 
 ```ruby
 completion = client.chat('How fast is a cheetah?', model: OmniAI::Google::Chat::Model::GEMINI_FLASH)
-completion.choice.message.content # 'A cheetah can reach speeds over 100 km/h.'
+completion.text # 'A cheetah can reach speeds over 100 km/h.'
 ```
 
 [Google API Reference `model`](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/model-versioning#gemini-model-versions)
@@ -83,7 +74,7 @@ completion.choice.message.content # 'A cheetah can reach speeds over 100 km/h.'
 
 ```ruby
 completion = client.chat('Pick a number between 1 and 5', temperature: 2.0)
-completion.choice.message.content # '3'
+completion.text # '3'
 ```
 
 [Google API Reference `temperature`](https://ai.google.dev/api/rest/v1/GenerationConfig)
@@ -94,7 +85,7 @@ completion.choice.message.content # '3'
 
 ```ruby
 stream = proc do |chunk|
-  print(chunk.choice.delta.content) # 'Better', 'three', 'hours', ...
+  print(chunk.text) # 'Better', 'three', 'hours', ...
 end
 client.chat('Be poetic.', stream:)
 ```

data/lib/omniai/google/chat/choice_serializer.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides choice serialize / deserialize.
+      module ChoiceSerializer
+        # @param choice [OmniAI::Chat::Choice]
+        # @param context [Context]
+        # @return [Hash]
+        def self.serialize(choice, context:)
+          content = choice.message.serialize(context:)
+          { content: }
+        end
+
+        # @param data [Hash]
+        # @param context [Context]
+        # @return [OmniAI::Chat::Choice]
+        def self.deserialize(data, context:)
+          message = OmniAI::Chat::Message.deserialize(data['content'], context:)
+          OmniAI::Chat::Choice.new(message:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/content_serializer.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides content serialize / deserialize.
+      module ContentSerializer
+        # @param data [Hash]
+        # @param context [Context]
+        # @return [OmniAI::Chat::Text, OmniAI::Chat::ToolCall]
+        def self.deserialize(data, context:)
+          case
+          when data['text'] then data['text']
+          when data['functionCall'] then OmniAI::Chat::ToolCall.deserialize(data, context:)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/function_serializer.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides function serialize / deserialize.
+      module FunctionSerializer
+        # @param function [OmniAI::Chat::Function]
+        # @return [Hash]
+        def self.serialize(function, *)
+          {
+            name: function.name,
+            args: function.arguments,
+          }
+        end
+
+        # @param data [Hash]
+        # @return [OmniAI::Chat::Function]
+        def self.deserialize(data, *)
+          name = data['name']
+          arguments = data['args']
+          OmniAI::Chat::Function.new(name:, arguments:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/media_serializer.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides media serialize / deserialize.
+      module MediaSerializer
+        # @param media [OmniAI::Chat::Media]
+        # @return [Hash]
+        def self.serialize(media, *)
+          {
+            inlineData: {
+              mimeType: media.type,
+              data: media.data,
+            },
+          }
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/message_serializer.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides message serialize / deserialize.
+      module MessageSerializer
+        # @param message [OmniAI::Chat::Message]
+        # @param context [OmniAI::Context]
+        # @return [Hash]
+        def self.serialize(message, context:)
+          role = message.role
+          parts = (arrayify(message.content) + arrayify(message.tool_call_list)).map do |part|
+            case part
+            when String then { text: part }
+            else part.serialize(context:)
+            end
+          end
+
+          { role:, parts: }
+        end
+
+        # @param data [Hash]
+        # @param context [OmniAI::Context]
+        # @return [OmniAI::Chat::Message]
+        def self.deserialize(data, context:)
+          role = data['role']
+          parts = arrayify(data['parts']).map do |part|
+            case
+            when part['text'] then OmniAI::Chat::Text.deserialize(part, context:)
+            when part['functionCall'] then OmniAI::Chat::ToolCall.deserialize(part, context:)
+            when part['functionResponse'] then OmniAI::Chat::ToolCallResult.deserialize(part, context:)
+            end
+          end
+
+          tool_call_list = parts.select { |part| part.is_a?(OmniAI::Chat::ToolCall) }
+          content = parts.reject { |part| part.is_a?(OmniAI::Chat::ToolCall) }
+
+          OmniAI::Chat::Message.new(content:, role:, tool_call_list:)
+        end
+
+        # @param content [Object]
+        # @return [Array<Object>]
+        def self.arrayify(content)
+          return [] if content.nil?
+
+          content.is_a?(Array) ? content : [content]
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/payload_serializer.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides payload serialize / deserialize.
+      module PayloadSerializer
+        # @param payload [OmniAI::Chat::Payload]
+        # @param context [OmniAI::Context]
+        # @return [Hash]
+        def self.serialize(payload, context:)
+          candidates = payload.choices.map { |choice| choice.serialize(context:) }
+          usage_metadata = payload.usage&.serialize(context:)
+
+          {
+            candidates:,
+            usage_metadata:,
+          }
+        end
+
+        # @param data [Hash]
+        # @param context [OmniAI::Context]
+        # @return [OmniAI::Chat::Payload]
+        def self.deserialize(data, context:)
+          choices = data['candidates'].map { |candidate| OmniAI::Chat::Choice.deserialize(candidate, context:) }
+          usage = OmniAI::Chat::Usage.deserialize(data['usage_metadata'], context:) if data['usage_metadata']
+
+          OmniAI::Chat::Payload.new(choices:, usage:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/text_serializer.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides text serialize / deserialize.
+      module TextSerializer
+        # @param text [OmniAI::Chat::Text]
+        # @return [Hash]
+        def self.serialize(text, *)
+          { text: text.text }
+        end
+
+        # @param data [Hash]
+        # @return [OmniAI::Chat::Text]
+        def self.deserialize(data, *)
+          OmniAI::Chat::Text.new(data['text'])
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/tool_call_result_serializer.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides tool-call result serialize / deserialize.
+      module ToolCallResultSerializer
+        # @param tool_call_response [OmniAI::Chat::ToolCallResult]
+        # @return [Hash]
+        def self.serialize(tool_call_response, *)
+          {
+            functionResponse: {
+              name: tool_call_response.tool_call_id,
+              response: {
+                name: tool_call_response.tool_call_id,
+                content: tool_call_response.content,
+              },
+            },
+          }
+        end
+
+        # @param data [Hash]
+        # @return [ToolCallResult]
+        def self.deserialize(data, *)
+          tool_call_id = data['functionResponse']['name']
+          content = data['functionResponse']['response']['content']
+          OmniAI::Chat::ToolCallResult.new(content:, tool_call_id:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/tool_call_serializer.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides tool-call serialize / deserialize.
+      module ToolCallSerializer
+        # @param tool_call [OmniAI::Chat::ToolCall]
+        # @param context [OmniAI::Context]
+        # @return [Hash]
+        def self.serialize(tool_call, context:)
+          { functionCall: tool_call.function.serialize(context:) }
+        end
+
+        # @param data [Hash]
+        # @param context [OmniAI::Context]
+        # @return [OmniAI::Chat::ToolCall]
+        def self.deserialize(data, context:)
+          function = OmniAI::Chat::Function.deserialize(data['functionCall'], context:)
+          OmniAI::Chat::ToolCall.new(id: function.name, function:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/tool_serializer.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides tool serialize / deserialize.
+      module ToolSerializer
+        # @param tool [OmniAI::Tool]
+        def self.serialize(tool, *)
+          {
+            name: tool.name,
+            description: tool.description,
+            parameters: tool.parameters&.prepare,
+          }
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat/usage_serializer.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module OmniAI
+  module Google
+    class Chat
+      # Overrides usage serialize / deserialize.
+      module UsageSerializer
+        # @param usage [OmniAI::Chat::Usage]
+        # @return [Hash]
+        def self.serialize(usage, *)
+          {
+            prompt_token_count: usage.input_tokens,
+            candidates_token_count: usage.output_tokens,
+            total_token_count: usage.total_tokens,
+          }
+        end
+
+        # @param data [Hash]
+        # @return [OmniAI::Chat::Usage]
+        def self.deserialize(data, *)
+          input_tokens = data['prompt_token_count']
+          output_tokens = data['candidates_token_count']
+          total_tokens = data['total_token_count']
+          OmniAI::Chat::Usage.new(input_tokens:, output_tokens:, total_tokens:)
+        end
+      end
+    end
+  end
+end

data/lib/omniai/google/chat.rb

@@ -25,49 +25,47 @@ module OmniAI
 
       DEFAULT_MODEL = Model::GEMINI_PRO
 
-      TEXT_SERIALIZER = lambda do |content, *|
-        { text: content.text }
-      end
+      # @return [Context]
+      CONTEXT = Context.build do |context|
+        context.serializers[:text] = TextSerializer.method(:serialize)
+        context.deserializers[:text] = TextSerializer.method(:deserialize)
 
-      # @param [Message]
-      # @return [Hash]
-      # @example
-      #   message = Message.new(...)
-      #   MESSAGE_SERIALIZER.call(message)
-      MESSAGE_SERIALIZER = lambda do |message, context:|
-        parts = message.content.is_a?(String) ? [Text.new(message.content)] : message.content
-        role = message.system? ? Role::USER : message.role
-
-        {
-          role:,
-          parts: parts.map { |part| part.serialize(context:) },
-        }
-      end
+        context.serializers[:file] = MediaSerializer.method(:serialize)
+        context.serializers[:url] = MediaSerializer.method(:serialize)
 
-      # @param [Media]
-      # @return [Hash]
-      # @example
-      #   media = Media.new(...)
-      #   MEDIA_SERIALIZER.call(media)
-      MEDIA_SERIALIZER = lambda do |media, *|
-        {
-          inlineData: {
-            mimeType: media.type,
-            data: media.data,
-          },
-        }
-      end
+        context.serializers[:tool_call] = ToolCallSerializer.method(:serialize)
+        context.deserializers[:tool_call] = ToolCallSerializer.method(:deserialize)
 
-      # @return [Context]
-      CONTEXT = Context.build do |context|
-        context.serializers[:message] = MESSAGE_SERIALIZER
-        context.serializers[:text] = TEXT_SERIALIZER
-        context.serializers[:file] = MEDIA_SERIALIZER
-        context.serializers[:url] = MEDIA_SERIALIZER
+        context.serializers[:tool_call_result] = ToolCallResultSerializer.method(:serialize)
+        context.deserializers[:tool_call_result] = ToolCallResultSerializer.method(:deserialize)
+
+        context.serializers[:function] = FunctionSerializer.method(:serialize)
+        context.deserializers[:function] = FunctionSerializer.method(:deserialize)
+
+        context.serializers[:usage] = UsageSerializer.method(:serialize)
+        context.deserializers[:usage] = UsageSerializer.method(:deserialize)
+
+        context.serializers[:payload] = PayloadSerializer.method(:serialize)
+        context.deserializers[:payload] = PayloadSerializer.method(:deserialize)
+
+        context.serializers[:choice] = ChoiceSerializer.method(:serialize)
+        context.deserializers[:choice] = ChoiceSerializer.method(:deserialize)
+
+        context.serializers[:message] = MessageSerializer.method(:serialize)
+        context.deserializers[:message] = MessageSerializer.method(:deserialize)
+
+        context.deserializers[:content] = ContentSerializer.method(:deserialize)
+
+        context.serializers[:tool] = ToolSerializer.method(:serialize)
       end
 
       protected
 
+      # @return [Context]
+      def context
+        CONTEXT
+      end
+
       # @return [HTTP::Response]
       def request!
         @client
@@ -82,7 +80,8 @@ module OmniAI
       # @return [Hash]
       def payload
         OmniAI::Google.config.chat_options.merge({
-          contents:,
+          system_instruction: @prompt.messages.find(&:system?)&.serialize(context:),
+          contents: @prompt.messages.reject(&:system?).map { |message| message.serialize(context:) },
           tools:,
           generationConfig: generation_config,
         }).compact
@@ -93,7 +92,7 @@ module OmniAI
         return unless @tools
 
         [
-          function_declarations: @tools&.map(&:prepare),
+          function_declarations: @tools.map { |tool| tool.serialize(context:) },
         ]
       end
 
@@ -104,15 +103,6 @@ module OmniAI
         { temperature: @temperature }.compact
       end
 
-      # Example:
-      #
-      #   [{ role: 'user', parts: [{ text: '...' }] }]
-      #
-      # @return [Array<Hash>]
-      def contents
-        @prompt.serialize(context: CONTEXT)
-      end
-
       # @return [String]
       def path
         "/#{@client.version}/models/#{@model}:#{operation}"
@@ -122,6 +112,15 @@ module OmniAI
       def operation
         @stream ? 'streamGenerateContent' : 'generateContent'
       end
+
+      # @return [Array<Message>]
+      def build_tool_call_messages(tool_call_list)
+        content = tool_call_list.map do |tool_call|
+          ToolCallResult.new(tool_call_id: tool_call.id, content: execute_tool_call(tool_call))
+        end
+
+        [Message.new(role: 'function', content:)]
+      end
     end
   end
 end

data/lib/omniai/google/config.rb

@@ -4,8 +4,13 @@ module OmniAI
   module Google
     # Configuration for Google.
     class Config < OmniAI::Config
+      module Version
+        STABLE = 'v1'
+        BETA = 'v1beta'
+      end
+
       DEFAULT_HOST = 'https://generativelanguage.googleapis.com'
-      DEFAULT_VERSION = 'v1'
+      DEFAULT_VERSION = Version::STABLE
 
       # @!attribute [rw] version
       #   @return [String, nil]

data/lib/omniai/google/version.rb

@@ -2,6 +2,6 @@
 
 module OmniAI
   module Google
-    VERSION = '1.6.3'
+    VERSION = '1.8.0'
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: omniai-google
 version: !ruby/object:Gem::Version
-  version: 1.6.3
+  version: 1.8.0
 platform: ruby
 authors:
 - Kevin Sylvestre
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2024-07-18 00:00:00.000000000 Z
+date: 2024-08-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: event_stream_parser
@@ -63,9 +63,17 @@ files:
 - README.md
 - lib/omniai/google.rb
 - lib/omniai/google/chat.rb
-- lib/omniai/google/chat/response/chunk.rb
-- lib/omniai/google/chat/response/completion.rb
-- lib/omniai/google/chat/response/stream.rb
+- lib/omniai/google/chat/choice_serializer.rb
+- lib/omniai/google/chat/content_serializer.rb
+- lib/omniai/google/chat/function_serializer.rb
+- lib/omniai/google/chat/media_serializer.rb
+- lib/omniai/google/chat/message_serializer.rb
+- lib/omniai/google/chat/payload_serializer.rb
+- lib/omniai/google/chat/text_serializer.rb
+- lib/omniai/google/chat/tool_call_result_serializer.rb
+- lib/omniai/google/chat/tool_call_serializer.rb
+- lib/omniai/google/chat/tool_serializer.rb
+- lib/omniai/google/chat/usage_serializer.rb
 - lib/omniai/google/client.rb
 - lib/omniai/google/config.rb
 - lib/omniai/google/version.rb

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

@@ -1,29 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  module Google
-    class Chat
-      module Response
-        # A chunk given when streaming.
-        class Chunk < OmniAI::Chat::Response::Chunk
-          # @return [Array<OmniAI::Chat::Choice>]
-          def choices
-            @choices ||= [].tap do |choices|
-              @data['candidates'].each do |candidate|
-                candidate['content']['parts'].each do |part|
-                  choices << OmniAI::Chat::Response::DeltaChoice.new(data: {
-                    'index' => candidate['index'],
-                    'delta' => {
-                      'role' => candidate['content']['role'],
-                      'content' => part['text'],
-                    },
-                  })
-                end
-              end
-            end
-          end
-        end
-      end
-    end
-  end
-end

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

@@ -1,29 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  module Google
-    class Chat
-      module Response
-        # A completion returned by the API.
-        class Completion < OmniAI::Chat::Response::Completion
-          # @return [Array<OmniAI::Chat::Choice>]
-          def choices
-            @choices ||= [].tap do |entries|
-              @data['candidates'].each do |candidate|
-                candidate['content']['parts'].each do |part|
-                  entries << OmniAI::Chat::Response::MessageChoice.new(data: {
-                    'index' => candidate['index'],
-                    'message' => {
-                      'role' => candidate['content']['role'],
-                      'content' => part['text'],
-                    },
-                  })
-                end
-              end
-            end
-          end
-        end
-      end
-    end
-  end
-end

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

@@ -1,21 +0,0 @@
-# frozen_string_literal: true
-
-module OmniAI
-  module Google
-    class Chat
-      module Response
-        # A stream given when streaming.
-        class Stream < OmniAI::Chat::Response::Stream
-          # @yield [OmniAI::Chat::Chunk]
-          def stream!(&)
-            @response.body.each do |chunk|
-              @parser.feed(chunk) do |_, data|
-                yield(Chunk.new(data: JSON.parse(data)))
-              end
-            end
-          end
-        end
-      end
-    end
-  end
-end