Class: Legion::Extensions::Llm::Mlx::Provider

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

Overview

MLX provider implementation for local OpenAI-compatible servers.

Defined Under Namespace

Modules: Capabilities

Class Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Class Attribute Details

.registry_publisher ⇒ Object 



24
25
26
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 24

def registry_publisher
  @registry_publisher ||= Legion::Extensions::Llm::RegistryPublisher.new(provider_family: :mlx)
end

Class Method Details

.capabilities ⇒ Object 



22
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 22

def capabilities = Capabilities

.configuration_options ⇒ Object 



20
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 20

def configuration_options = %i[mlx_api_base mlx_api_key]

.configuration_requirements ⇒ Object 



21
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 21

def configuration_requirements = []

.default_tier ⇒ Object 



19
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 19

def default_tier = :local

.default_transport ⇒ Object 



18
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 18

def default_transport = :http

.local? ⇒ Boolean

Returns:

  • (Boolean) 


17
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 17

def local? = true

.slug ⇒ Object 



16
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 16

def slug = 'mlx'

Instance Method Details

#api_base ⇒ Object 



57
58
59
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 57

def api_base
  normalize_url(config.mlx_api_base || settings[:endpoint] || 'http://localhost:8000')
end

#circuit_state(status) ⇒ Object 



184
185
186
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 184

def circuit_state(status)
  status == 'healthy' ? 'closed' : 'open'
end

#embedding_model?(model_id) ⇒ Boolean

Returns:

  • (Boolean) 


154
155
156
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 154

def embedding_model?(model_id)
  model_id.to_s.match?(/embed|bge|e5|nomic/i)
end

#extract_catalog_capabilities(model_info) ⇒ Object 



145
146
147
148
149
150
151
152
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 145

def extract_catalog_capabilities(model_info)
  model_id = model_info.respond_to?(:id) ? model_info.id.to_s : model_info.to_s
  caps = {}
  caps[:embeddings] = true if model_id.match?(/embed|bge|e5|nomic/i)
  caps[:vision] = true if model_id.match?(/vlm|vision|llava|pixtral|qwen.*vl/i)
  caps[:streaming] = true unless caps[:embeddings]
  caps
end

#extract_real_capabilities(model_info) ⇒ Object 



135
136
137
138
139
140
141
142
143
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 135

def extract_real_capabilities(model_info)
  return {} unless model_info.respond_to?(:metadata)

  meta = model_info.metadata
  return {} unless meta.is_a?(Hash)

  caps = meta[:capabilities]
  caps.is_a?(Hash) ? caps : {}
end

#fetch_models_config ⇒ Object 



223
224
225
226
227
228
229
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 223

def fetch_models_config
  conf = config.models if config.respond_to?(:models)
  conf ||= config[:models] if config.respond_to?(:[])
  conf if conf.respond_to?(:to_h)
rescue StandardError
  nil
end

#headers ⇒ Object 



61
62
63
64
65
66
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 61

def headers
  hdrs = identity_headers
  token = config.mlx_api_key
  hdrs['Authorization'] = "Bearer #{token}" unless token.nil? || token.to_s.empty?
  hdrs
end

#health(live: false) ⇒ Object 



70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 70

def health(live: false)
  log.info("Checking MLX health live=#{live} at #{api_base}#{health_url}")
  raw = connection.get(health_url).body
  health_payload(raw)
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: 'mlx.provider.health')
  {
    provider: :mlx,
    instance_id: provider_instance_id,
    status: 'unhealthy',
    ready: false,
    circuit_state: 'open',
    error: e.class.name,
    message: e.message
  }
end

#health_payload(raw) ⇒ Object 



162
163
164
165
166
167
168
169
170
171
172
173
174
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 162

def health_payload(raw)
  ready = health_ready?(raw)
  status = health_status(ready)

  {
    provider: :mlx,
    instance_id: provider_instance_id,
    status: status,
    ready: ready,
    circuit_state: circuit_state(status),
    raw: raw
  }
end

#health_ready?(raw) ⇒ Boolean

Returns:

  • (Boolean) 


176
177
178
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 176

def health_ready?(raw)
  raw.is_a?(Hash) ? raw.fetch('ready', raw.fetch(:ready, true)) : true
end

#health_status(ready) ⇒ Object 



180
181
182
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 180

def health_status(ready)
  ready ? 'healthy' : 'unhealthy'
end

#health_url ⇒ Object 



68
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 68

def health_url = '/health'

#instance_capability_config ⇒ Object 



196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 196

def instance_capability_config
  cfg = config
  result = {}
  %i[capabilities enable_thinking enable_tools enable_streaming enable_vision enable_embeddings
     thinking_flag tools_flag streaming_flag vision_flag embedding_flag embeddings_flag
     tool_flag images_flag image_flag].each do |key|
    next unless cfg.respond_to?(key)

    val = cfg.send(key)
    result[key] = val unless val.nil?
  rescue StandardError
    next
  end
  result
end

#list_models ⇒ Object 



94
95
96
97
98
99
100
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 94

def list_models(**)
  log.info('Listing available MLX models')
  super.tap do |models|
    log.info("Discovered #{Array(models).size} MLX models")
    self.class.registry_publisher.publish_models_async(models, readiness: readiness(live: false))
  end
end

#model_capability_config(model_id) ⇒ Object 



212
213
214
215
216
217
218
219
220
221
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 212

def model_capability_config(model_id)
  models_conf = fetch_models_config
  return {} unless models_conf

  hash = models_conf.to_h
  hash[model_id.to_s] || hash[model_id.to_sym] || {}
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: 'mlx.model_capability_config')
  {}
end

#offering_from_model(model_info, health: {}) ⇒ Object

rubocop:disable Metrics/AbcSize



102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 102

def offering_from_model(model_info, health: {}) # rubocop:disable Metrics/AbcSize
  policy = resolve_capability_policy(model_info)
  ctx = model_info.respond_to?(:context_length) ? model_info.context_length : nil

  Legion::Extensions::Llm::Routing::ModelOffering.new(
    provider_family: :mlx,
    instance_id: config.respond_to?(:instance_id) ? config.instance_id : :default,
    transport: offering_transport,
    tier: offering_tier,
    model: model_info.id,
    canonical_model_alias: model_info.respond_to?(:name) ? model_info.name : nil,
    model_family: model_info.respond_to?(:family) ? model_info.family : nil,
    usage_type: embedding_model?(model_info.id) ? :embedding : :inference,
    capabilities: policy[:capabilities],
    capability_sources: policy[:sources],
    limits: { context_window: ctx }.compact,
    health: health,
    metadata: offering_metadata_for(model_info).merge(capability_sources: policy[:sources])
  )
end

#offering_metadata_for(model_info) ⇒ Object 



231
232
233
234
235
236
237
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 231

def offering_metadata_for(model_info)
  {
    raw_model: model_info.id,
    parameter_count: model_info.respond_to?(:parameter_count) ? model_info.parameter_count : nil,
    quantization: model_info.respond_to?(:quantization) ? model_info.quantization : nil
  }.compact
end

#provider_capability_config ⇒ Object 



188
189
190
191
192
193
194
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 188

def provider_capability_config
  conf = Legion::Extensions::Llm::CredentialSources.setting(:extensions, :llm, :mlx)
  conf.is_a?(Hash) ? conf.to_h.except(:instances, 'instances') : {}
rescue StandardError => e
  handle_exception(e, level: :warn, handled: true, operation: 'mlx.provider_capability_config')
  {}
end

#provider_envelope_capabilities ⇒ Object 



158
159
160
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 158

def provider_envelope_capabilities
  { streaming: true }
end

#readiness(live: false) ⇒ Object 



87
88
89
90
91
92
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 87

def readiness(live: false)
  log.info("Checking MLX readiness (live=#{live})")
  super.tap do |metadata|
    self.class.registry_publisher.publish_readiness_async(metadata) if live
  end
end

#resolve_capability_policy(model_info) ⇒ Object 



123
124
125
126
127
128
129
130
131
132
133
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 123

def resolve_capability_policy(model_info)
  Legion::Extensions::Llm::CapabilityPolicy.resolve(
    real: extract_real_capabilities(model_info),
    provider_catalog: extract_catalog_capabilities(model_info),
    probe: {},
    provider_envelope: provider_envelope_capabilities,
    provider_config: provider_capability_config,
    instance_config: instance_capability_config,
    model_config: model_capability_config(model_info.id)
  )
end

#settings ⇒ Object 



53
54
55
 
# File 'lib/legion/extensions/llm/mlx/provider.rb', line 53

def settings
  Mlx.default_settings
end