Support Embedding Dimensions (#73)

This PR implements the ability to specify custom dimensions when
generating embeddings, as requested in issue #47.

### What's included
- Added support for passing a dimensions parameter to the embed method
- Implemented dimensions handling in both OpenAI and Gemini providers
- Added tests to verify dimension param works correctly
- Optimized the Gemini provider's `embed` method to reduce unnecessary
API calls when embedding texts, resulting in lower token usage. From now
on, it uses `batchEmbedContents` endpoint within one request, for both
single and multiple text embeddings.
- Modernize Gemini embeddings following DIP principle, as implemented in
`openai/embeddings.rb`.
- The Gemini embeddings API response does not contain the
promptTokenCount attribute, so I have removed it.

### Implementation notes
I've decided to only implement the per-request dimension configuration
and not the global configuration option that was initially proposed in
the issue. This is because each embedding model has its own default
dimensions, making a global setting potentially confusing.

With this implementation, users can set the embedding dimensions like:
```ruby
embedding = RubyLLM.embed(
  "Ruby is a programmer's best friend",
  model: "text-embedding-3-small",
  dimensions: 512
)
```

### References
- OpenAI API docs:
https://platform.openai.com/docs/api-reference/embeddings
- Gemini API docs: https://ai.google.dev/api/embeddings

Resolves #47

---------

Co-authored-by: Carmine Paolino <[email protected]>

Author: papgmez

docs/guides/embeddings.md

@@ -92,6 +92,44 @@ end
 
 Refer to the [Working with Models Guide]({% link guides/models.md %}) for details on finding available embedding models and their capabilities.
 
+## Choosing Dimensions
+
+Each embedding model has its own default output dimensions. For example, OpenAI's `text-embedding-3-small` outputs 1536 dimensions by default, while `text-embedding-3-large` outputs 3072 dimensions. RubyLLM allows you to specify these dimensions per request:
+
+```ruby
+embedding = RubyLLM.embed(
+  "This is a test sentence",
+  model: "text-embedding-3-small",
+  dimensions: 512
+)
+```
+
+This is particularly useful when:
+- Working with vector databases that have specific dimension requirements
+- Ensuring consistent dimensionality across different requests
+- Optimizing storage and query performance in your vector database
+
+Note that not all models support custom dimensions. If you specify dimensions that aren't supported by the chosen model, RubyLLM will use the model's default dimensions.
+
+## Using Embedding Results
+
+### Vector Properties
+
+The embedding result contains useful information:
+
+```ruby
+embedding = RubyLLM.embed("Example text")
+
+# The vector representation
+puts embedding.vectors.class  # => Array
+puts embedding.vectors.first.class  # => Float
+
+# The vector dimensions
+puts embedding.vectors.first.length # => 1536
+
+# The model used
+puts embedding.model  # => "text-embedding-3-small"
+
 ## Using Embedding Results
 
 A primary use case for embeddings is measuring the semantic similarity between texts. Cosine similarity is a common metric.

lib/ruby_llm/embedding.rb

@@ -12,14 +12,14 @@ def initialize(vectors:, model:, input_tokens: 0)
       @input_tokens = input_tokens
     end
 
-    def self.embed(text, model: nil, provider: nil, context: nil)
+    def self.embed(text, model: nil, provider: nil, context: nil, dimensions: nil)
       config = context&.config || RubyLLM.config
       model_id = model || config.default_embedding_model
       Models.find(model_id, provider)
 
       provider = Provider.for(model_id)
       connection = context ? context.connection_for(provider) : provider.connection(config)
-      provider.embed(text, model: model_id, connection:)
+      provider.embed(text, model: model_id, connection:, dimensions:)
     end
   end
 end

lib/ruby_llm/provider.rb

@@ -31,10 +31,10 @@ def list_models(connection:)
         parse_list_models_response response, slug, capabilities
       end
 
-      def embed(text, model:, connection:)
-        payload = render_embedding_payload(text, model:)
-        response = connection.post embedding_url, payload
-        parse_embedding_response response
+      def embed(text, model:, connection:, dimensions:)
+        payload = render_embedding_payload(text, model:, dimensions:)
+        response = connection.post(embedding_url(model:), payload)
+        parse_embedding_response(response, model:)
       end
 
       def paint(prompt, model:, size:, connection:)

lib/ruby_llm/providers/gemini/embeddings.rb

@@ -5,47 +5,31 @@ module Providers
     module Gemini
       # Embeddings methods for the Gemini API integration
       module Embeddings
-        # Must be public for Provider module
-        def embed(text, model:, connection:) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength
-          payload = {
-            content: {
-              parts: format_text_for_embedding(text)
-            }
-          }
+        module_function
 
-          url = "models/#{model}:embedContent"
-          response = connection.post url, payload
+        def embedding_url(model:)
+          "models/#{model}:batchEmbedContents"
+        end
+
+        def render_embedding_payload(text, model:, dimensions:)
+          { requests: [text].flatten.map { |t| single_embedding_payload(t, model:, dimensions:) } }
+        end
 
-          if text.is_a?(Array)
-            # We need to make separate calls for each text with Gemini
-            embeddings = text.map do |t|
-              single_payload = { content: { parts: [{ text: t.to_s }] } }
-              single_response = connection.post url, single_payload
-              single_response.body.dig('embedding', 'values')
-            end
+        def parse_embedding_response(response, model:)
+          vectors = response.body['embeddings']&.map { |e| e['values'] }
+          vectors in [vectors]
 
-            Embedding.new(
-              vectors: embeddings,
-              model: model,
-              input_tokens: response.body.dig('usageMetadata', 'promptTokenCount') || 0
-            )
-          else
-            Embedding.new(
-              vectors: response.body.dig('embedding', 'values'),
-              model: model,
-              input_tokens: response.body.dig('usageMetadata', 'promptTokenCount') || 0
-            )
-          end
+          Embedding.new(vectors:, model:, input_tokens: 0)
         end
 
         private
 
-        def format_text_for_embedding(text)
-          if text.is_a?(Array)
-            text.map { |t| { text: t.to_s } }
-          else
-            [{ text: text.to_s }]
-          end
+        def single_embedding_payload(text, model:, dimensions:)
+          {
+            model: "models/#{model}",
+            content: { parts: [{ text: text.to_s }] },
+            outputDimensionality: dimensions
+          }.compact
         end
       end
     end

lib/ruby_llm/providers/openai/embeddings.rb

@@ -7,31 +7,27 @@ module OpenAI
       module Embeddings
         module_function
 
-        def embedding_url
+        def embedding_url(...)
           'embeddings'
         end
 
-        def render_embedding_payload(text, model:)
+        def render_embedding_payload(text, model:, dimensions:)
           {
             model: model,
-            input: text
-          }
+            input: text,
+            dimensions: dimensions
+          }.compact
         end
 
-        def parse_embedding_response(response)
+        def parse_embedding_response(response, model:)
           data = response.body
-          model_id = data['model']
           input_tokens = data.dig('usage', 'prompt_tokens') || 0
           vectors = data['data'].map { |d| d['embedding'] }
 
           # If we only got one embedding, return it as a single vector
-          vectors = vectors.first if vectors.size == 1
+          vectors in [vectors]
 
-          Embedding.new(
-            vectors: vectors,
-            model: model_id,
-            input_tokens: input_tokens
-          )
+          Embedding.new(vectors:, model:, input_tokens:)
         end
       end
     end