Class: Legion::Extensions::Llm::Vertex::Provider

Inherits:
Provider
  • Object
show all
Defined in:
lib/legion/extensions/llm/vertex/provider.rb

Overview

Google Cloud Vertex AI provider implementation for the Legion::Extensions::Llm contract.

Defined Under Namespace

Modules: Capabilities

Constant Summary collapse

DEFAULT_LOCATION =

rubocop:disable Metrics/ClassLength

'us-central1'
DEFAULT_PROJECT = 
'env://GOOGLE_CLOUD_PROJECT'
DEFAULT_PUBLISHER = 
'google'
STATIC_MODELS = 
[
  { model: 'gemini-2.5-flash', alias: 'gemini-flash', publisher: 'google', model_family: :gemini },
  { model: 'gemini-2.5-pro', alias: 'gemini-pro', publisher: 'google', model_family: :gemini },
  { model: 'gemini-embedding-001', alias: 'gemini-embedding', publisher: 'google',
    model_family: :gemini, usage_type: :embedding },
  { model: 'text-embedding-005', alias: 'text-embedding', publisher: 'google',
    model_family: :gemini, usage_type: :embedding },
  { model: 'claude-sonnet-4-5', alias: 'claude-sonnet', publisher: 'anthropic',
    model_family: :anthropic, api: :raw_predict },
  { model: 'mistral-medium-3', alias: 'mistral-medium', publisher: 'mistralai',
    model_family: :mistral, api: :raw_predict },
  { model: 'llama-4-maverick', alias: 'llama-4-maverick', publisher: 'meta',
    model_family: :meta, api: :raw_predict }
].freeze
ALIASES = 
STATIC_MODELS.to_h { |entry| [entry.fetch(:alias), entry.fetch(:model)] }.freeze
PUBLISHERS = 
STATIC_MODELS.to_h { |entry| [entry.fetch(:model), entry.fetch(:publisher)] }.freeze
API_MODES = 
STATIC_MODELS.to_h { |entry| [entry.fetch(:model), entry.fetch(:api, :generate_content)] }.freeze
MODEL_FAMILIES = 
STATIC_MODELS.to_h { |entry| [entry.fetch(:model), entry.fetch(:model_family)] }.freeze

Class Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Class Attribute Details

.registry_publisher ⇒ Object 



59
60
61
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 59

def registry_publisher
  @registry_publisher ||= RegistryPublisher.new
end

Class Method Details

.capabilities ⇒ Object 



57
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 57

def capabilities = Capabilities

.configuration_options ⇒ Object 



44
45
46
47
48
49
50
51
52
53
54
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 44

def configuration_options
  %i[
    vertex_project
    vertex_location
    vertex_api_base
    vertex_access_token
    vertex_credentials
    vertex_model_aliases
    vertex_discovery_live
  ]
end

.configuration_requirements ⇒ Object 



56
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 56

def configuration_requirements = []

.resolve_model_id(model_id, config: nil) ⇒ Object 



63
64
65
66
67
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 63

def resolve_model_id(model_id, config: nil)
  configured_aliases = config.respond_to?(:vertex_model_aliases) ? config.vertex_model_aliases : nil
  aliases = ALIASES.merge((configured_aliases || {}).transform_keys(&:to_s))
  aliases.fetch(model_id.to_s, model_id.to_s)
end

.slug ⇒ Object 



42
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 42

def slug = 'vertex'

Instance Method Details

#api_base ⇒ Object 



87
88
89
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 87

def api_base
  config.vertex_api_base || "https://#{location}-aiplatform.googleapis.com/v1"
end

#chat(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {}) ⇒ Object 



168
169
170
171
172
173
174
175
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 168

def chat(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {})
  model_id = model_id(model)
  @model = model_id
  payload = Utils.deep_merge(chat_payload(messages, model: model_id, temperature:, max_tokens:, tools:,
                                                    tool_prefs:, stream: false), params)
  response = connection.post(chat_url(model_id, stream: false), payload)
  parse_chat_response(response, model: model_id)
end

#complete(messages, tools:, temperature:, model:, params: {}, schema: nil, thinking: nil, tool_prefs: nil) ⇒ Object 



217
218
219
220
221
222
223
224
225
226
227
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 217

def complete(messages, tools:, temperature:, model:, params: {}, schema: nil, thinking: nil, tool_prefs: nil,
             &)
  payload = params.dup
  payload[:generationConfig] = Utils.deep_merge(payload[:generationConfig] || {},
                                                generation_config(temperature, schema, thinking))
  if block_given?
    stream(messages, model:, temperature:, tools:, tool_prefs:, params: payload, &)
  else
    chat(messages, model:, temperature:, tools:, tool_prefs:, params: payload)
  end
end

#completion_url ⇒ Object 



98
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 98

def completion_url = generate_content_url(model: @model || STATIC_MODELS.first.fetch(:model))

#count_tokens(messages, model:, params: {}) ⇒ Object 



188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 188

def count_tokens(messages, model:, params: {})
  model_id = model_id(model)
  unless generate_content_model?(model_id)
    return {
      supported: false,
      provider: :vertex,
      model: resource_name(model_id),
      reason: 'Vertex countTokens is standardized for generateContent publisher models'
    }
  end

  payload = Utils.deep_merge({ contents: format_messages(messages) }, params)
  response = connection.post(count_tokens_url(model: model_id), payload)
  { input_tokens: response.body['totalTokens'], raw: response.body }
end

#count_tokens_url(model:) ⇒ Object 



100
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 100

def count_tokens_url(model:) = "#{publisher_model_path(model)}:countTokens"

#discover_offerings(live: false, **filters) ⇒ Object 



116
117
118
119
120
121
122
123
124
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 116

def discover_offerings(live: false, **filters)
  return static_offerings(**filters) unless live

  response = connection.get(models_url)
  models = response.body['publisherModels'] || response.body['models'] || []
  models.map { |model| offering_from_live_model(model) }.tap do |offerings|
    self.class.registry_publisher.publish_offerings_async(offerings, readiness: readiness(live: false))
  end
end

#embed(text, model:, dimensions: nil, task_type: nil, title: nil, params: {}) ⇒ Object 



204
205
206
207
208
209
210
211
212
213
214
215
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 204

def embed(text, model:, dimensions: nil, task_type: nil, title: nil, params: {})
  model_id = model_id(model)
  unless Capabilities.embeddings?(model_id)
    raise NotImplementedError, "Vertex embedding payload for #{model_id} is not standardized"
  end

  instances = Array(text).map { |item| embedding_instance(item, task_type:, title:) }
  parameters = { outputDimensionality: dimensions }.compact
  payload = Utils.deep_merge({ instances: instances, parameters: parameters }, params)
  response = connection.post(embedding_url(model: model_id), payload)
  parse_embedding_response(response, model: model_id)
end

#embedding_url(model:) ⇒ Object 



101
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 101

def embedding_url(model:) = "#{publisher_model_path(model)}:predict"

#generate_content_url(model:) ⇒ Object 



103
104
105
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 103

def generate_content_url(model:)
  "#{publisher_model_path(model)}:generateContent"
end

#headers ⇒ Object 



91
92
93
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 91

def headers
  { 'Authorization' => bearer_token, 'Content-Type' => 'application/json; charset=utf-8' }.compact
end

#health(live: false) ⇒ Object 



143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 143

def health(live: false)
  baseline = {
    provider: :vertex,
    project: project,
    location: location,
    configured: configured?,
    ready: configured?,
    live: live,
    credentials: credential_source
  }
  return baseline.merge(checked: false) unless live

  connection.get(models_url)
  baseline.merge(checked: true)
rescue StandardError => e
  baseline.merge(checked: true, ready: false, error: e.class.name, message: e.message)
end

#location ⇒ Object 



96
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 96

def location = config.vertex_location || DEFAULT_LOCATION

#models_url ⇒ Object 



97
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 97

def models_url = publisher_parent

#offering_for(model:, model_family: nil, instance_id: :default, **metadata) ⇒ Object 



126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 126

def offering_for(model:, model_family: nil, instance_id: :default, **metadata)
  model_id = model_id(model)
  publisher = metadata.delete(:publisher) || publisher_for(model_id)
  family = model_family || metadata.delete(:model_family) || model_family_for(model_id, publisher)

  build_offering(
    model: resource_name(model_id, publisher:),
    alias_name: alias_for(model_id),
    model_family: family,
    instance_id: instance_id,
    publisher: publisher,
    usage_type: metadata.delete(:usage_type) || usage_type_for(model_id),
    api: metadata.delete(:api) || api_for(model_id),
    metadata: metadata
  )
end

#project ⇒ Object 



95
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 95

def project = config.vertex_project || ENV.fetch('GOOGLE_CLOUD_PROJECT', DEFAULT_PROJECT)

#raw_predict_url(model:, stream: false) ⇒ Object 



111
112
113
114
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 111

def raw_predict_url(model:, stream: false)
  suffix = stream ? 'streamRawPredict' : 'rawPredict'
  "#{publisher_model_path(model)}:#{suffix}"
end

#readiness(live: false) ⇒ Object 



161
162
163
164
165
166
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 161

def readiness(live: false)
  health(live:).merge(local: false, remote: true, api_base: api_base,
                      endpoints: endpoint_manifest).tap do |metadata|
    self.class.registry_publisher.publish_readiness_async(metadata) if live
  end
end

#stream(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {}) {|chunk| ... } ⇒ Object

Yields:

  • (chunk) 


177
178
179
180
181
182
183
184
185
186
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 177

def stream(messages, model:, temperature: nil, max_tokens: nil, tools: {}, tool_prefs: nil, params: {})
  model_id = model_id(model)
  @model = model_id
  payload = Utils.deep_merge(chat_payload(messages, model: model_id, temperature:, max_tokens:, tools:,
                                                    tool_prefs:, stream: true), params)
  response = connection.post(chat_url(model_id, stream: true), payload)
  chunk = build_chunk(response.body, model: model_id)
  yield chunk if block_given? && chunk.content
  parse_chat_response(response, model: model_id)
end

#stream_generate_content_url(model:) ⇒ Object 



107
108
109
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 107

def stream_generate_content_url(model:)
  "#{publisher_model_path(model)}:streamGenerateContent?alt=sse"
end

#stream_url ⇒ Object 



99
 
# File 'lib/legion/extensions/llm/vertex/provider.rb', line 99

def stream_url = stream_generate_content_url(model: @model || STATIC_MODELS.first.fetch(:model))