omniai-google 3.6.0 → 3.7.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bf6bab2505b42bcb8b61ca4e428b77bbffb9f961e183a759a5ad4c4caa6faded
-  data.tar.gz: 46a4e0576c2b311ef25bb97384e268d9ae831d9c5eab192bd6192c42578fdffd
+  metadata.gz: 74848884db95fef34b6c7306b2814ecf8d50797963e690a93b77e69c93cc7d3a
+  data.tar.gz: 9fbb92e5fe4fe36d9086182391ae383180607b5e97b36a17e458dc200bef73e1
 SHA512:
-  metadata.gz: b076336eef0202b7a489003f8ac1acf6a75149e5f4c279c94452e3735075f15ef13e1849506fcd4e4bcd47471b1713affdbe794ddd69884c7144e656058fffb5
-  data.tar.gz: 2fa90885dddc31135e74b6985ac089b2e016150c83c40b1470b29b4f754c90c23387643a902cbaae9951c502e1598284fff84198d022e98d1d8731c619c87279
+  metadata.gz: c92c79303ec17e46832115d075aa29760ba8d1f52111286d51d370c49c5266bb2e0ec9956e058a133de69bd0bb28c8ea4fa52efab34edea64f9cdc896397069e
+  data.tar.gz: 5fc61a87b13a8bd3ba73bc6b8c150c7da1d77aef938d7c826f9eac541b78ea288fb0698bbaf9a7609d7a19a85607b3d04d97e82cd013a46f64c319c86329ea60

data/lib/omniai/google/client.rb

@@ -85,8 +85,9 @@ module OmniAI
       #
       # @param input [String, Array<String>, Array<Integer>] required
       # @param model [String] optional
-      def embed(input, model: Embed::DEFAULT_MODEL)
-        Embed.process!(input, model:, client: self)
+      # @param options [Hash] provider-specific options (e.g. task_type: "RETRIEVAL_DOCUMENT")
+      def embed(input, model: Embed::DEFAULT_MODEL, **)
+        Embed.process!(input, model:, client: self, **)
       end
 
       # @raise [OmniAI::Error]

data/lib/omniai/google/embed.rb

@@ -2,7 +2,7 @@
 
 module OmniAI
   module Google
-    # An Google embed implementation.
+    # A Google embed implementation.
     #
     # Usage:
     #
@@ -14,67 +14,125 @@ module OmniAI
         TEXT_EMBEDDING_004 = "text-embedding-004"
         TEXT_EMBEDDING_005 = "text-embedding-005"
         TEXT_MULTILINGUAL_EMBEDDING_002 = "text-multilingual-embedding-002"
+        GEMINI_EMBEDDING_001 = "gemini-embedding-001"
+        GEMINI_EMBEDDING_2_PREVIEW = "gemini-embedding-2-preview"
         EMBEDDING = TEXT_EMBEDDING_004
         MULTILINGUAL_EMBEDDING = TEXT_MULTILINGUAL_EMBEDDING_002
       end
 
       DEFAULT_MODEL = Model::EMBEDDING
 
-      DEFAULT_EMBEDDINGS_DESERIALIZER = proc do |data, *|
+      BATCH_EMBED_CONTENTS_DESERIALIZER = proc do |data, *|
         data["embeddings"].map { |embedding| embedding["values"] }
       end
 
-      VERTEX_EMBEDDINGS_DESERIALIZER = proc do |data, *|
+      PREDICT_EMBEDDINGS_DESERIALIZER = proc do |data, *|
         data["predictions"].map { |prediction| prediction["embeddings"]["values"] }
       end
 
-      VERTEX_USAGE_DESERIALIZER = proc do |data, *|
-        tokens = data["predictions"].map { |prediction| prediction["embeddings"]["statistics"]["token_count"] }.sum
+      PREDICT_USAGE_DESERIALIZER = proc do |data, *|
+        tokens = data["predictions"].sum { |prediction| prediction["embeddings"]["statistics"]["token_count"] }
 
         Usage.new(prompt_tokens: tokens, total_tokens: tokens)
       end
 
+      EMBED_CONTENT_DESERIALIZER = proc do |data, *|
+        [data["embedding"]["values"]]
+      end
+
+      USAGE_METADATA_DESERIALIZER = proc do |data, *|
+        prompt_tokens = data.dig("usageMetadata", "promptTokenCount")
+        total_tokens = data.dig("usageMetadata", "totalTokenCount")
+
+        Usage.new(prompt_tokens:, total_tokens:)
+      end
+
       # @return [Context]
-      DEFAULT_CONTEXT = Context.build do |context|
-        context.deserializers[:embeddings] = DEFAULT_EMBEDDINGS_DESERIALIZER
+      BATCH_EMBED_CONTENTS_CONTEXT = Context.build do |context|
+        context.deserializers[:embeddings] = BATCH_EMBED_CONTENTS_DESERIALIZER
+        context.deserializers[:usage] = USAGE_METADATA_DESERIALIZER
       end
 
       # @return [Context]
-      VERTEX_CONTEXT = Context.build do |context|
-        context.deserializers[:embeddings] = VERTEX_EMBEDDINGS_DESERIALIZER
-        context.deserializers[:usage] = VERTEX_USAGE_DESERIALIZER
+      PREDICT_CONTEXT = Context.build do |context|
+        context.deserializers[:embeddings] = PREDICT_EMBEDDINGS_DESERIALIZER
+        context.deserializers[:usage] = PREDICT_USAGE_DESERIALIZER
+      end
+
+      # @return [Context]
+      EMBED_CONTENT_CONTEXT = Context.build do |context|
+        context.deserializers[:embeddings] = EMBED_CONTENT_DESERIALIZER
+        context.deserializers[:usage] = USAGE_METADATA_DESERIALIZER
       end
 
     protected
 
-      # @return [Boolean]
-      def vertex?
-        @client.vertex?
+      # Determines which endpoint to use based on client and model configuration.
+      # Routes gemini-embedding-2-* models to embedContent on Vertex, as Google's
+      # Vertex AI requires this endpoint for newer multimodal embedding models.
+      #
+      # @return [Symbol] :embed_content, :predict, or :batch_embed_contents
+      def endpoint
+        @endpoint ||=
+          if @client.vertex? && @model.start_with?("gemini-embedding-2")
+            :embed_content
+          elsif @client.vertex?
+            :predict
+          else
+            :batch_embed_contents
+          end
       end
 
       # @return [Context]
       def context
-        vertex? ? VERTEX_CONTEXT : DEFAULT_CONTEXT
+        case endpoint
+        when :embed_content then EMBED_CONTENT_CONTEXT
+        when :predict then PREDICT_CONTEXT
+        when :batch_embed_contents then BATCH_EMBED_CONTENTS_CONTEXT
+        end
       end
 
-      # @return [Array[Hash]]
-      def instances
-        arrayify(@input).map { |content| { content: } }
+      # @return [Hash]
+      def payload
+        case endpoint
+        when :embed_content then embed_content_payload
+        when :predict then predict_payload
+        when :batch_embed_contents then batch_embed_contents_payload
+        end
       end
 
-      # @return [Array[Hash]]
-      def requests
-        arrayify(@input).map do |text|
-          {
-            model: "models/#{@model}",
-            content: { parts: [{ text: }] },
-          }
-        end
+      # Builds payload for the Vertex embedContent endpoint (gemini-embedding-2-* models).
+      # @return [Hash]
+      def embed_content_payload
+        raise ArgumentError, "embedContent does not support batch input" if @input.is_a?(Array) && @input.length > 1
+
+        text = @input.is_a?(Array) ? @input.first : @input
+        result = { content: { parts: [{ text: }] } }
+        result[:taskType] = @options[:task_type] if @options[:task_type]
+        result
       end
 
+      # Builds payload for the Vertex predict endpoint (text-embedding and gemini-embedding-001 models).
       # @return [Hash]
-      def payload
-        vertex? ? { instances: } : { requests: }
+      def predict_payload
+        inputs = arrayify(@input)
+        { instances: inputs.map { |text| { content: text } } }
+      end
+
+      # Builds payload for the Google AI batchEmbedContents endpoint (non-Vertex).
+      # @return [Hash]
+      def batch_embed_contents_payload
+        inputs = arrayify(@input)
+        {
+          requests: inputs.map do |text|
+            request = {
+              model: "models/#{@model}",
+              content: { parts: [{ text: }] },
+            }
+            request[:taskType] = @options[:task_type] if @options[:task_type]
+            request
+          end,
+        }
       end
 
       # @return [Hash]
@@ -82,20 +140,15 @@ module OmniAI
         { key: (@client.api_key unless @client.credentials?) }.compact
       end
 
-      # @return [String]
-      def path
-        "/#{@client.path}/models/#{@model}:#{procedure}"
-      end
+      PROCEDURES = {
+        embed_content: "embedContent",
+        predict: "predict",
+        batch_embed_contents: "batchEmbedContents",
+      }.freeze
 
       # @return [String]
-      def procedure
-        vertex? ? "predict" : "batchEmbedContents"
-      end
-
-      # @param input [Object]
-      # @return [Array]
-      def arrayify(input)
-        input.is_a?(Array) ? input : [input]
+      def path
+        "/#{@client.path}/models/#{@model}:#{PROCEDURES[endpoint]}"
       end
     end
   end

data/lib/omniai/google/version.rb

@@ -2,6 +2,6 @@
 
 module OmniAI
   module Google
-    VERSION = "3.6.0"
+    VERSION = "3.7.1"
   end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: omniai-google
 version: !ruby/object:Gem::Version
-  version: 3.6.0
+  version: 3.7.1
 platform: ruby
 authors:
 - Kevin Sylvestre