ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/providers/vertexai/anthropic.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      # The Anthropic protocol over Vertex AI rawPredict endpoints.
+      class Anthropic < Protocols::Anthropic
+        ANTHROPIC_VERSION = 'vertex-2023-10-16'
+
+        def completion_url
+          "#{@provider.model_path(@model.id, publisher: 'anthropic')}:rawPredict"
+        end
+
+        def stream_url
+          "#{@provider.model_path(@model.id, publisher: 'anthropic')}:streamRawPredict"
+        end
+
+        def render_payload(messages, **)
+          payload = super
+          payload.delete(:model)
+          payload.merge(anthropic_version: ANTHROPIC_VERSION)
+        end
+
+        def count_tokens(*, **)
+          raise Error, "#{@provider.name} doesn't support token counting for Claude models"
+        end
+
+        def supports_provider_file_references?
+          false
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/capabilities.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      # Feature capability gaps not represented in upstream model catalogs.
+      module Capabilities
+        def self.augment(capabilities, model_id:, modalities:)
+          return capabilities if model_id.include?('embedding')
+
+          additions = []
+          additions << 'tool_choice' if model_id == 'gemini-2.5-flash'
+          additions << 'transcription' if modalities[:input].include?('audio') && modalities[:output].include?('text')
+          capabilities | additions
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/chat_completions/batches.rb

@@ -0,0 +1,54 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      class ChatCompletions
+        # Vertex AI MaaS batch prediction rows using OpenAI JSONL shape.
+        module Batches
+          include Protocols::VertexAI::BatchPrediction
+
+          private
+
+          def vertex_batch_request(request)
+            {
+              custom_id: request[:custom_id],
+              method: 'POST',
+              url: '/v1/chat/completions',
+              body: batch_payload(request)
+            }
+          end
+
+          def validate_batch_requests!(requests)
+            return if requests.all? { |request| chat_completion_payload?(request.fetch(:payload)) }
+
+            raise Error, 'vertexai MaaS batch requests require chat completion payloads'
+          end
+
+          def chat_completion_payload?(payload)
+            payload.key?(:messages) || payload.key?('messages')
+          end
+
+          def vertex_batch_model_path(model)
+            publisher, name = model.split('/', 2)
+            raise Error, 'vertexai MaaS batch requests require publisher/model ids' unless publisher && name
+
+            @provider.model_path(name, publisher:)
+          end
+
+          def parse_vertex_batch_result(line, fallback_index)
+            index = vertex_batch_result_index(line, fallback_index)
+            response = line['response']
+            body = response.is_a?(Hash) ? response['body'] || response : response
+
+            if body
+              [index, parse_completion_body(body, raw: body)]
+            else
+              [index, nil, batch_failure(index, line.dig('status', 'message') || batch_error_message(line))]
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/chat_completions.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      # MaaS models (publisher-prefixed ids like meta/llama-3.3-70b-instruct-maas)
+      # speak Chat Completions through Vertex AI's OpenAI-compatible endpoint.
+      class ChatCompletions < Protocols::ChatCompletions
+        def completion_url
+          "#{@provider.location_path}/endpoints/openapi/chat/completions"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/embed_content.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      # Gemini Embedding 2 over Vertex AI's single-content embedding endpoint.
+      class EmbedContent < Protocols::Gemini
+        def embedding_url(model:)
+          "#{@provider.model_path(model)}:embedContent"
+        end
+
+        def render_embedding(text, **options)
+          return super unless text.is_a?(Array)
+
+          { requests: text.map { |value| render_embedding_payload(value, **options) } }
+        end
+
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, with: [],
+                                     provider_options: {})
+          if text.is_a?(Array) && text.size != 1
+            raise ArgumentError, 'Vertex AI embedContent accepts one text at a time'
+          end
+          raise ArgumentError, "#{model} takes task instructions and titles in the text" if task_type || title
+
+          payload = {
+            content: { parts: Protocols::Gemini::Media.format_content(text.is_a?(Array) ? text.first : text, with) },
+            outputDimensionality: dimensions
+          }.compact
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        def parse_embedding_response(response, model:, text:)
+          vectors = response.body.dig('embedding', 'values')
+          raise Error.new('Vertex AI returned no embedding', response:) if vectors.nil? || vectors.empty?
+
+          vectors = [vectors] if text.is_a?(Array)
+          Embedding.new(vectors:, model:, input_tokens: response.body.dig('usageMetadata', 'promptTokenCount'))
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/embeddings.rb

@@ -8,23 +8,38 @@ module RubyLLM
         module_function
 
         def embedding_url(model:)
-          "projects/#{@config.vertexai_project_id}/locations/#{@config.vertexai_location}/publishers/google/models/#{model}:predict" # rubocop:disable Layout/LineLength
+          "#{@provider.model_path(model)}:predict"
         end
 
-        def render_embedding_payload(text, model:, dimensions:) # rubocop:disable Lint/UnusedMethodArgument
-          {
-            instances: [text].flatten.map { |t| { content: t.to_s } }
-          }.tap do |payload|
-            payload[:parameters] = { outputDimensionality: dimensions } if dimensions
+        def supports_embedding_media?
+          false
+        end
+
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, provider_options: {})
+          instances = [text].flatten.map do |t|
+            { content: t.to_s, task_type: task_type, title: title }.compact
           end
+          payload = { instances: instances }
+          payload[:parameters] = { outputDimensionality: dimensions } if dimensions
+
+          Support::Utils.deep_merge(payload, provider_options)
         end
 
         def parse_embedding_response(response, model:, text:)
           predictions = response.body['predictions']
           vectors = predictions&.map { |p| p.dig('embeddings', 'values') }
+          input_tokens = embedding_input_tokens(predictions)
           vectors = vectors.first if vectors&.length == 1 && !text.is_a?(Array)
 
-          Embedding.new(vectors:, model:, input_tokens: 0)
+          Embedding.new(vectors:, model:, input_tokens:)
+        end
+
+        def embedding_input_tokens(predictions)
+          counts = Array(predictions).filter_map do |prediction|
+            prediction.dig('embeddings', 'statistics', 'token_count')
+          end
+          counts.sum unless counts.empty?
         end
       end
     end

data/lib/ruby_llm/providers/vertexai/gemini/batches.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      class Gemini
+        # Vertex AI Gemini batch prediction rows.
+        module Batches
+          include Protocols::VertexAI::BatchPrediction
+
+          private
+
+          def vertex_batch_request(request)
+            payload = RubyLLM::Support::Utils.deep_stringify_keys(batch_payload(request))
+            labels = payload.fetch('labels', {}).merge('ruby_llm_batch_id' => request[:custom_id])
+            { request: payload.merge('labels' => labels) }
+          end
+
+          def parse_vertex_batch_result(line, fallback_index)
+            index = vertex_batch_result_index(line, fallback_index)
+
+            if line['response']
+              body = line['response']
+              [index, parse_completion_body(body, raw: body)]
+            else
+              [index, nil, batch_failure(index, vertex_batch_status_message(line))]
+            end
+          end
+
+          # Gemini prediction rows carry status as a plain string, empty on
+          # success and the error text on failure.
+          def vertex_batch_status_message(line)
+            status = line['status']
+            message = status.is_a?(Hash) ? status['message'] : status
+
+            message.to_s.empty? ? batch_error_message(line) : message
+          end
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/gemini.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      # The Gemini protocol over Vertex AI endpoints.
+      class Gemini < Protocols::Gemini
+        include VertexAI::Embeddings
+        include VertexAI::Models
+        include VertexAI::Videos
+
+        SERVER_TOOL_ALIASES = Protocols::Gemini::SERVER_TOOL_ALIASES.merge(
+          file_search: lambda { |options|
+            { tool: { retrieval: { vertexAiSearch: Support::Utils.deep_symbolize_keys(options) } } }
+          }
+        ).freeze
+
+        def server_tool_aliases
+          SERVER_TOOL_ALIASES
+        end
+
+        def completion_url
+          "#{@provider.model_path(@model.id)}:generateContent"
+        end
+
+        def stream_url
+          "#{@provider.model_path(@model.id)}:streamGenerateContent?alt=sse"
+        end
+
+        def count_tokens_url
+          "#{@provider.model_path(@model.id)}:countTokens"
+        end
+
+        def render_count_tokens_payload(messages, **options)
+          count_tokens_request(messages, **options)
+        end
+
+        def caches_url
+          "#{@provider.location_path}/cachedContents"
+        end
+
+        def cache_name(name)
+          name = name.name if name.is_a?(CachedContent)
+          name.to_s.include?('/') ? name.to_s : "#{caches_url}/#{name}"
+        end
+
+        def cache_model_name(model_id)
+          @provider.model_path(model_id)
+        end
+
+        def images_url(with: nil, mask: nil) # rubocop:disable Lint/UnusedMethodArgument
+          id = model_id(@model)
+
+          "#{@provider.model_path(id)}:#{image_endpoint_action(id)}"
+        end
+
+        def speech_url(model:)
+          "#{@provider.model_path(model)}:generateContent"
+        end
+
+        private
+
+        def transcription_url(model)
+          "#{@provider.model_path(model)}:generateContent"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/live_transcription.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI
+      class LiveTranscription < Protocols::Gemini::LiveTranscription # :nodoc: all
+        def validate_transcription_request(...)
+          super
+          return if @config.vertexai_location == 'global'
+
+          raise ArgumentError, 'Vertex AI Live transcription requires vertexai_location = "global"'
+        end
+
+        def transcription_model_name(model)
+          @provider.model_path(model)
+        end
+
+        def websocket_service
+          'google.cloud.aiplatform.v1beta1.LlmBidiService/BidiGenerateContent'
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/mistral.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    class VertexAI < Provider
+      # Mistral's models speak their own dialect of Chat Completions over Vertex AI
+      # rawPredict endpoints. We reuse Mistral's dialect wholesale and only swap the URLs.
+      class Mistral < Protocols::ChatCompletions
+        include Providers::Mistral::Chat
+        include Providers::Mistral::Embeddings
+        include Providers::Mistral::Media
+        include Providers::Mistral::Models
+
+        # The Mistral model families Vertex AI serves directly. Shared by the
+        # registry (which models to list) and protocol_for (where to route them).
+        MODELS = /\A(mistral|ministral|codestral|magistral|mathstral|pixtral|devstral|voxtral)/
+
+        def completion_url
+          "#{@provider.model_path(@model.id, publisher: 'mistralai')}:rawPredict"
+        end
+
+        def stream_url
+          "#{@provider.model_path(@model.id, publisher: 'mistralai')}:streamRawPredict"
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/providers/vertexai/models.rb

@@ -5,75 +5,171 @@ module RubyLLM
     class VertexAI
       # Models methods for the Vertex AI integration
       module Models
-        # Gemini and other Google models that aren't returned by the API
+        def self.models_dev_alias(model_id, models_dev_by_key, _provider_model = nil)
+          source = models_dev_by_key["gemini:#{model_id}"]
+          Model.new(source.to_h.merge(provider: 'vertexai')) if source
+        end
+
+        # Google models the publisher catalog omits in some regions while still
+        # serving them there. Every id must be callable in at least one region;
+        # ids that answer nowhere do not belong here.
         KNOWN_GOOGLE_MODELS = %w[
           gemini-2.5-flash-lite
           gemini-2.5-pro
           gemini-2.5-flash
           gemini-2.0-flash-lite-001
           gemini-2.0-flash-001
-          gemini-2.0-flash
-          gemini-2.0-flash-exp
           gemini-1.5-pro-002
           gemini-1.5-pro
-          gemini-1.5-flash-002
-          gemini-1.5-flash
-          gemini-1.5-flash-8b
           gemini-pro
           gemini-pro-vision
-          gemini-exp-1206
-          gemini-exp-1121
           gemini-embedding-001
           text-embedding-005
           text-embedding-004
           text-multilingual-embedding-002
         ].freeze
 
+        # Every publisher with models Vertex AI serves as a service. The rest
+        # of the Model Garden is deploy-it-yourself and not callable directly.
+        PUBLISHERS = %w[google anthropic mistralai meta deepseek-ai qwen openai moonshotai zai-org].freeze
+
+        # Vertex AI serves a different slice of the catalog in each location and
+        # neither of these is a superset of the other, so a listing unions them.
+        CATALOG_LOCATIONS = %w[global us-central1].freeze
+
+        # The configured location has to answer: without it we would report a
+        # catalog the caller cannot reach. A supplementary location is a bonus,
+        # so a failure there is a warning and the rest of the union stands.
         def list_models
-          all_models = []
-          page_token = nil
+          fetched = []
+          counts = {}
+
+          catalog_connections.each do |location, connection|
+            models = location_models(connection)
+            counts[location] = models.size
+            fetched.concat(models)
+          rescue StandardError => e
+            raise if location == configured_location
+
+            RubyLLM.logger.warn "Skipping the Vertex AI catalog at #{location}: #{e.class}: #{e.message}"
+          end
+
+          models = fetched.uniq(&:id)
+          log_catalog(counts, models)
+
+          models + build_known_models(models.map(&:id))
+        end
+
+        private
+
+        def configured_location
+          @config.vertexai_location.to_s
+        end
+
+        # The location lives in the host, so locations sharing one api base,
+        # as they do behind a custom vertexai_api_base, are one catalog.
+        def catalog_connections
+          [configured_location, *CATALOG_LOCATIONS]
+            .uniq { |location| @provider.api_base_for(location) }
+            .to_h { |location| [location, connection_for(location)] }
+        end
+
+        def connection_for(location)
+          return @connection if location == configured_location
+
+          Transport::Connection.new(@provider, @config, api_base: @provider.api_base_for(location))
+        end
+
+        def log_catalog(counts, models)
+          per_location = counts.map { |location, count| "#{location} (#{count})" }.join(', ')
+          RubyLLM.logger.info "Fetched the Vertex AI catalog from #{per_location}: #{models.size} models"
+        end
+
+        # A publisher with nothing to offer in a region answers 200 with an
+        # empty list, so any error here is infrastructure, not an empty
+        # catalog. Reporting a partial catalog as a success would drop the
+        # missing publishers from the registry.
+        def location_models(connection)
+          failures = []
+          models = PUBLISHERS.flat_map do |publisher|
+            publisher_models(publisher, connection)
+          rescue StandardError => e
+            failures << "#{publisher} (#{e.class}: #{e.message})"
+            []
+          end
+
+          raise Error, "Could not fetch the Vertex AI catalog for #{failures.join(', ')}" if failures.any?
+
+          models
+        end
+
+        def publisher_models(publisher, connection)
+          catalog(publisher, connection).filter_map { |model_data| build_publisher_model(publisher, model_data) }
+        end
 
-          all_models.concat(build_known_models)
+        # MaaS models are called as publisher/name through the OpenAI-compatible
+        # endpoint; directly served models by their bare catalog name.
+        def build_publisher_model(publisher, model_data)
+          name = model_data['name'].split('/').last
+          return if deployable?(model_data)
+
+          if name.end_with?('-maas')
+            build_model_from_api_data(model_data, "#{publisher}/#{name}")
+          elsif served_directly?(publisher, name)
+            build_model_from_api_data(model_data, name)
+          end
+        end
+
+        # Deploy-it-yourself Model Garden cards expose deploy actions; managed
+        # services Vertex AI serves on our behalf never do.
+        def deployable?(model_data)
+          actions = model_data['supportedActions'] || {}
+          actions.key?('deploy') || actions.key?('multiDeployVertex') || actions.key?('deployGke')
+        end
+
+        # Among the managed models, which publishers we route by bare name, and
+        # for Google (whose catalog is a grab-bag of vision, media, and AutoML
+        # products) which of those names are chat or embedding models.
+        def served_directly?(publisher, name)
+          case publisher
+          when 'google' then name.match?(/\Agemini|embedding/)
+          when 'anthropic' then true
+          when 'mistralai' then VertexAI::Mistral::MODELS.match?(name)
+          else false
+          end
+        end
+
+        def catalog(publisher, connection)
+          models = []
+          page_token = nil
 
           loop do
-            response = @connection.get('publishers/google/models') do |req|
+            response = connection.get("publishers/#{publisher}/models") do |req|
               req.headers['x-goog-user-project'] = @config.vertexai_project_id
               req.params = { pageSize: 100 }
               req.params[:pageToken] = page_token if page_token
             end
 
-            publisher_models = response.body['publisherModels'] || []
-            publisher_models.each do |model_data|
-              next if model_data['launchStage'] == 'DEPRECATED'
-
-              model_id = extract_model_id_from_path(model_data['name'])
-              all_models << build_model_from_api_data(model_data, model_id)
-            end
-
+            models.concat(response.body['publisherModels'] || [])
             page_token = response.body['nextPageToken']
             break unless page_token
           end
 
-          all_models
-        rescue StandardError => e
-          RubyLLM.logger.debug { "Error fetching Vertex AI models: #{e.message}" }
-          build_known_models
+          models.reject { |model_data| model_data['launchStage'] == 'DEPRECATED' }
         end
 
-        private
-
-        def build_known_models
-          KNOWN_GOOGLE_MODELS.map do |model_id|
-            Model::Info.new(
+        def build_known_models(fetched_ids)
+          (KNOWN_GOOGLE_MODELS - fetched_ids).map do |model_id|
+            Model.new(
               id: model_id,
               name: model_id,
-              provider: slug,
+              provider: @provider.slug,
               family: determine_model_family(model_id),
               created_at: nil,
               context_window: nil,
               max_output_tokens: nil,
               modalities: nil,
-              capabilities: %w[streaming function_calling],
+              capabilities: extract_capabilities(model_id),
               pricing: nil,
               metadata: {
                 source: 'known_models'
@@ -83,16 +179,16 @@ module RubyLLM
         end
 
         def build_model_from_api_data(model_data, model_id)
-          Model::Info.new(
+          Model.new(
             id: model_id,
             name: model_id,
-            provider: slug,
+            provider: @provider.slug,
             family: determine_model_family(model_id),
             created_at: nil,
             context_window: nil,
             max_output_tokens: nil,
             modalities: nil,
-            capabilities: extract_capabilities(model_data),
+            capabilities: extract_capabilities(model_data['name']),
             pricing: nil,
             metadata: {
               version_id: model_data['versionId'],
@@ -104,12 +200,21 @@ module RubyLLM
           )
         end
 
-        def extract_model_id_from_path(path)
-          path.split('/').last
-        end
-
         def determine_model_family(model_id)
           case model_id
+          when /^claude.*haiku/ then 'claude-haiku'
+          when /^claude.*sonnet/ then 'claude-sonnet'
+          when /^claude.*opus/ then 'claude-opus'
+          when /^claude/ then 'claude'
+          when %r{^meta/} then 'llama'
+          when %r{^deepseek-ai/} then 'deepseek'
+          when %r{^qwen/} then 'qwen'
+          when %r{^moonshotai/} then 'kimi'
+          when %r{^zai-org/} then 'glm'
+          when %r{^openai/} then 'gpt-oss'
+          when %r{^google/} then 'gemma'
+          when /^codestral/ then 'codestral'
+          when /^mi(ni)?stral/ then 'mistral'
           when /^gemini-2\.\d+/ then 'gemini-2'
           when /^gemini-1\.\d+/ then 'gemini-1.5'
           when /^text-embedding/ then 'text-embedding'
@@ -118,11 +223,8 @@ module RubyLLM
           end
         end
 
-        def extract_capabilities(model_data)
-          capabilities = ['streaming']
-          model_name = model_data['name']
-          capabilities << 'function_calling' if model_name.include?('gemini')
-          capabilities.uniq
+        def extract_capabilities(name)
+          name.match?(/ocr|embedding/) ? %w[streaming] : %w[streaming function_calling]
         end
       end
     end