omniai-google 3.6.0 → 3.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bf6bab2505b42bcb8b61ca4e428b77bbffb9f961e183a759a5ad4c4caa6faded
-  data.tar.gz: 46a4e0576c2b311ef25bb97384e268d9ae831d9c5eab192bd6192c42578fdffd
+  metadata.gz: 81bb938d26c6804a6942a17e57e1b1393f0e77925467a60659443f28c76d3752
+  data.tar.gz: ce70db23296a73089d2cbc6eaf344ae269af945c112dfc0942a17e31dc116dfe
 SHA512:
-  metadata.gz: b076336eef0202b7a489003f8ac1acf6a75149e5f4c279c94452e3735075f15ef13e1849506fcd4e4bcd47471b1713affdbe794ddd69884c7144e656058fffb5
-  data.tar.gz: 2fa90885dddc31135e74b6985ac089b2e016150c83c40b1470b29b4f754c90c23387643a902cbaae9951c502e1598284fff84198d022e98d1d8731c619c87279
+  metadata.gz: 0f5101592e2603341f545b3402a0c4b3b3516dd3fee36532d2a5691803e11610124db98d89af4e532bd16ccfff2663a37308d999dee79aea700a3066df559771
+  data.tar.gz: 78109e536527bc9853e84111e664c25c3574a0f7c952fadbaf2a025478be0172c7cd1a203c4bdf88aa185b242f2fb74f9b88f6d88205a0ec8c479bfe2391b6a0

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, **options)
+        Embed.process!(input, model:, client: self, **options)
       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,124 @@ 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: prompt_tokens, total_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 [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 instances
-        arrayify(@input).map { |content| { content: } }
+      # 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: text }] } }
+        result[:taskType] = @options[:task_type] if @options[:task_type]
+        result
       end
 
-      # @return [Array[Hash]]
-      def requests
-        arrayify(@input).map do |text|
-          {
-            model: "models/#{@model}",
-            content: { parts: [{ text: }] },
-          }
-        end
+      # Builds payload for the Vertex predict endpoint (text-embedding and gemini-embedding-001 models).
+      # @return [Hash]
+      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 payload
-        vertex? ? { instances: } : { requests: }
+      def batch_embed_contents_payload
+        inputs = arrayify(@input)
+        {
+          requests: inputs.map do |text|
+            request = {
+              model: "models/#{@model}",
+              content: { parts: [{ text: text }] },
+            }
+            request[:taskType] = @options[:task_type] if @options[:task_type]
+            request
+          end
+        }
       end
 
       # @return [Hash]
@@ -84,18 +141,13 @@ module OmniAI
 
       # @return [String]
       def path
-        "/#{@client.path}/models/#{@model}:#{procedure}"
-      end
-
-      # @return [String]
-      def procedure
-        vertex? ? "predict" : "batchEmbedContents"
-      end
+        procedure = case endpoint
+                    when :embed_content then "embedContent"
+                    when :predict then "predict"
+                    when :batch_embed_contents then "batchEmbedContents"
+                    end
 
-      # @param input [Object]
-      # @return [Array]
-      def arrayify(input)
-        input.is_a?(Array) ? input : [input]
+        "/#{@client.path}/models/#{@model}:#{procedure}"
       end
     end
   end

data/lib/omniai/google/version.rb

@@ -2,6 +2,6 @@
 
 module OmniAI
   module Google
-    VERSION = "3.6.0"
+    VERSION = "3.7.0"
   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.0
 platform: ruby
 authors:
 - Kevin Sylvestre