lex-llm 0.3.1 → 0.4.2

data/lib/legion/extensions/llm/provider.rb

@@ -62,6 +62,16 @@ module Legion
         end
 
         # rubocop:disable Metrics/ParameterLists
+        def chat(messages:, model:, tools: [], temperature: nil, params: {}, headers: {}, schema: nil, thinking: nil,
+                 tool_prefs: nil)
+          complete(messages, tools:, temperature:, model:, params:, headers:, schema:, thinking:, tool_prefs:)
+        end
+
+        def stream_chat(messages:, model:, tools: [], temperature: nil, params: {}, headers: {}, schema: nil,
+                        thinking: nil, tool_prefs: nil, &)
+          complete(messages, tools:, temperature:, model:, params:, headers:, schema:, thinking:, tool_prefs:, &)
+        end
+
         def complete(messages, tools:, temperature:, model:, params: {}, headers: {}, schema: nil, thinking: nil,
                      tool_prefs: nil, &)
           normalized_temperature = maybe_normalize_temperature(temperature, model)
@@ -88,14 +98,55 @@ module Legion
         end
         # rubocop:enable Metrics/ParameterLists
 
-        def list_models
+        def list_models(live: false, **filters)
+          _ = [live, filters]
           response = @connection.get models_url
           parse_list_models_response response, slug, capabilities
         end
 
-        def embed(text, model:, dimensions:)
-          payload = render_embedding_payload(text, model:, dimensions:)
-          response = @connection.post(embedding_url(model:), payload)
+        def discover_offerings(live: false, **filters)
+          return filter_cached_offerings(Array(@cached_offerings), filters) unless live
+
+          provider_health = health(live:)
+          @cached_offerings = Array(list_models(live:, **filters)).filter_map do |model|
+            next unless model_matches_filters?(model, filters)
+
+            offering_from_model(model, health: provider_health)
+          end
+          @cached_offerings
+        end
+
+        def health(live: false)
+          readiness_data = readiness(live:)
+          raw_health = readiness_data[:health] || readiness_data['health'] || {}
+          status = health_status(readiness_data, raw_health)
+          {
+            provider: slug.to_sym,
+            instance_id: provider_instance_id,
+            status:,
+            ready: readiness_data[:ready] == true || readiness_data['ready'] == true,
+            circuit_state: status == 'healthy' ? 'closed' : 'open',
+            latency_ms: raw_health[:latency_ms] || raw_health['latency_ms'],
+            raw: raw_health
+          }.compact
+        rescue StandardError => e
+          handle_exception(e, level: :warn, handled: true, operation: 'llm.provider.health')
+          {
+            provider: slug.to_sym,
+            instance_id: provider_instance_id,
+            status: 'unhealthy',
+            ready: false,
+            circuit_state: 'open',
+            error: e.class.name,
+            message: e.message
+          }
+        end
+
+        def embed(text:, model:, dimensions: nil, params: {}, headers: {})
+          payload = Utils.deep_merge(render_embedding_payload(text, model:, dimensions:), params)
+          response = @connection.post(embedding_url(model:), payload) do |req|
+            req.headers = headers.merge(req.headers) unless headers.empty?
+          end
           parse_embedding_response(response, model:, text:)
         end
 
@@ -112,6 +163,18 @@ module Legion
           parse_image_response(response, model:)
         end
 
+        def image(prompt:, model:, size:, with: nil, mask: nil, params: {}) # rubocop:disable Metrics/ParameterLists
+          paint(prompt, model:, size:, with:, mask:, params:)
+        end
+
+        def count_tokens(messages:, model:, params: {})
+          _ = [model, params]
+          Array(messages).sum do |message|
+            content = message.respond_to?(:content) ? message.content : message[:content] || message['content']
+            estimate_text_tokens(content)
+          end
+        end
+
         def transcribe(audio_file, model:, language:, **)
           file_part = build_audio_file_part(audio_file)
           payload = render_transcription_payload(file_part, model:, language:, **)
@@ -319,6 +382,12 @@ module Legion
           end
         end
 
+        def provider_instance_id
+          return config.instance_id.to_sym if config.respond_to?(:instance_id) && config.instance_id
+
+          :default
+        end
+
         class << self
           def name
             to_s.split('::').last
@@ -369,6 +438,131 @@ module Legion
           raise UnsupportedAttachmentError, "#{name} does not support image references in paint"
         end
 
+        def offering_from_model(model, health: {})
+          Routing::ModelOffering.new(
+            provider_family: slug.to_sym,
+            provider_instance: model.instance || provider_instance_id,
+            transport: offering_transport,
+            tier: offering_tier,
+            model: model.id,
+            canonical_model_alias: model.name,
+            model_family: model.family,
+            usage_type: offering_usage_type(model),
+            capabilities: model.capabilities,
+            limits: offering_limits(model),
+            health:,
+            metadata: offering_metadata(model)
+          )
+        end
+
+        def offering_transport
+          local? ? :local : :http
+        end
+
+        def offering_tier
+          local? ? :local : :direct
+        end
+
+        def offering_usage_type(model)
+          model.embedding? ? :embedding : :inference
+        end
+
+        def offering_limits(model)
+          {
+            context_window: model.context_length,
+            max_output_tokens: model.max_output_tokens
+          }.compact
+        end
+
+        def offering_metadata(model)
+          {
+            raw_model: model.id,
+            parameter_count: model.parameter_count,
+            parameter_size: model.parameter_size,
+            quantization: model.quantization,
+            size_bytes: model.size_bytes,
+            modalities_input: model.modalities_input,
+            modalities_output: model.modalities_output
+          }.merge(model.metadata || {}).compact
+        end
+
+        def model_matches_filters?(model, filters)
+          return true if filters.empty?
+
+          filters.all? do |key, value|
+            blank_filter_value?(value) || model_matches_filter?(model, key, value)
+          end
+        end
+
+        def blank_filter_value?(value)
+          value.nil? || (value.respond_to?(:empty?) && value.empty?)
+        end
+
+        def model_matches_filter?(model, key, value)
+          case key.to_sym
+          when :capability, :capabilities
+            Array(value).all? { |capability| model.supports?(capability) }
+          when :type, :usage_type, :purpose
+            offering_usage_type(model).to_s == value.to_s || model.type.to_s == value.to_s
+          when :model, :id, :name
+            [model.id, model.name].map(&:to_s).include?(value.to_s)
+          when :instance, :instance_id, :provider_instance
+            provider_instance_id.to_s == value.to_s || model.instance.to_s == value.to_s
+          else
+            true
+          end
+        end
+
+        def filter_cached_offerings(offerings, filters)
+          return offerings if filters.empty?
+
+          offerings.select do |offering|
+            filters.all? do |key, value|
+              blank_filter_value?(value) || offering_matches_filter?(offering, key, value)
+            end
+          end
+        end
+
+        def offering_matches_filter?(offering, key, value)
+          case key.to_sym
+          when :provider, :provider_family
+            offering.provider_family.to_s == value.to_s
+          when :capability, :capabilities
+            Array(value).all? { |capability| offering.supports?(capability) }
+          when :type, :usage_type, :purpose
+            offering.usage_type.to_s == value.to_s
+          when :model, :id, :name
+            [offering.model, offering.canonical_model_alias].compact.map(&:to_s).include?(value.to_s)
+          when :instance, :instance_id, :provider_instance
+            [offering.provider_instance, offering.instance_id].compact.map(&:to_s).include?(value.to_s)
+          else
+            true
+          end
+        end
+
+        def health_status(readiness_data, raw_health)
+          return 'healthy' if readiness_data[:ready] == true || readiness_data['ready'] == true
+
+          status = raw_health[:status] || raw_health['status'] || raw_health[:state] || raw_health['state']
+          return 'healthy' if %w[ok ready healthy running].include?(status.to_s.downcase)
+
+          'unhealthy'
+        end
+
+        def estimate_text_tokens(content)
+          text = case content
+                 when Content
+                   [content.text, *content.attachments.map(&:to_s)].compact.join(' ')
+                 when Array
+                   content.map do |part|
+                     part.respond_to?(:[]) ? part[:text] || part['text'] || part.to_s : part.to_s
+                   end.join(' ')
+                 else
+                   content.to_s
+                 end
+          [(text.length / 4.0).ceil, 1].max
+        end
+
         def build_audio_file_part(file_path)
           expanded_path = File.expand_path(file_path)
           mime_type = Marcel::MimeType.for(Pathname.new(expanded_path))

data/lib/legion/extensions/llm/provider_contract.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      # Documents the canonical public provider method signatures shared by provider gems.
+      module ProviderContract
+        REQUIRED_SIGNATURES = {
+          chat: [%i[keyreq messages], %i[keyreq model]],
+          stream_chat: [%i[keyreq messages], %i[keyreq model]],
+          embed: [%i[keyreq text], %i[keyreq model]],
+          image: [%i[keyreq prompt], %i[keyreq model]],
+          list_models: [%i[key live], %i[keyrest filters]],
+          discover_offerings: [%i[key live], %i[keyrest filters]],
+          health: [%i[key live]],
+          count_tokens: [%i[keyreq messages], %i[keyreq model], %i[key params]]
+        }.freeze
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/provider_settings.rb

@@ -7,7 +7,10 @@ module Legion
       module ProviderSettings
         module_function
 
-        def build(family:, instance: {}, enabled: true, discovery: {}, instances: {})
+        def build(family:, instance: {}, enabled: true, discovery: {}, instances: {}, fleet: nil, gateways: nil, # rubocop:disable Metrics/ParameterLists
+                  **legacy_settings)
+          validate_provider_defaults!(fleet:, gateways:, legacy_settings:)
+
           deep_merge(
             Legion::Extensions::Llm.default_settings,
             {
@@ -27,6 +30,20 @@ module Legion
           )
         end
 
+        def validate_provider_defaults!(fleet:, gateways:, legacy_settings:)
+          if fleet
+            raise ArgumentError,
+                  'Provider fleet defaults must be nested under an instance, for example instances.default.fleet'
+          end
+
+          raise ArgumentError, 'Provider gateways settings are no longer supported; use instances instead' if gateways
+
+          return if legacy_settings.empty?
+
+          invalid_keys = legacy_settings.keys.map(&:to_sym)
+          raise ArgumentError, "Unsupported top-level provider settings: #{invalid_keys.join(', ')}"
+        end
+
         def deep_dup(value)
           case value
           when Hash

data/lib/legion/extensions/llm/responses/chat_response.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Responses
+        # Normalized non-streaming chat provider response.
+        class ChatResponse
+          attr_reader :content, :thinking, :metadata, :model, :tool_calls, :tokens, :raw, :internal_metadata
+
+          def initialize(content:, thinking: nil, metadata: {}, model: nil, tool_calls: nil, tokens: nil, raw: nil) # rubocop:disable Metrics/ParameterLists
+            extraction = ThinkingExtractor.extract(content, metadata: metadata)
+
+            @content = extraction.content
+            @thinking = thinking || extraction.thinking
+            @metadata = extraction.metadata
+            @internal_metadata = metadata.to_h
+            @model = model
+            @tool_calls = tool_calls
+            @tokens = tokens
+            @raw = raw
+
+            freeze
+          end
+
+          def to_h
+            {
+              content: content,
+              metadata: metadata,
+              model: model,
+              tool_calls: tool_calls,
+              tokens: tokens
+            }.compact
+          end
+
+          def to_internal_h
+            to_h.merge(thinking: thinking, metadata: internal_metadata, raw: raw).compact
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/responses/embedding_response.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Responses
+        # Normalized embedding provider response.
+        class EmbeddingResponse
+          attr_reader :vectors, :model, :tokens, :metadata, :raw
+
+          def initialize(vectors:, model:, tokens: nil, metadata: {}, raw: nil)
+            @vectors = vectors
+            @model = model
+            @tokens = tokens
+            @metadata = ThinkingExtractor.extract(nil, metadata: metadata).metadata
+            @internal_metadata = metadata.to_h
+            @raw = raw
+
+            freeze
+          end
+
+          def to_h
+            {
+              vectors: vectors,
+              model: model,
+              tokens: tokens,
+              metadata: metadata
+            }.compact
+          end
+
+          def to_internal_h
+            to_h.merge(metadata: @internal_metadata, raw: raw).compact
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/responses/stream_chunk.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Responses
+        # Normalized streaming provider response chunk.
+        class StreamChunk
+          attr_reader :content, :thinking, :metadata, :model, :tool_calls, :tokens, :raw, :internal_metadata
+
+          def initialize(content: nil, thinking: nil, metadata: {}, model: nil, tool_calls: nil, tokens: nil, raw: nil) # rubocop:disable Metrics/ParameterLists
+            extraction = ThinkingExtractor.extract(content, metadata: metadata)
+
+            @content = extraction.content
+            @thinking = thinking || extraction.thinking
+            @metadata = extraction.metadata
+            @internal_metadata = metadata.to_h
+            @model = model
+            @tool_calls = tool_calls
+            @tokens = tokens
+            @raw = raw
+
+            freeze
+          end
+
+          def to_h
+            {
+              content: content,
+              metadata: metadata,
+              model: model,
+              tool_calls: tool_calls,
+              tokens: tokens
+            }.compact
+          end
+
+          def to_internal_h
+            to_h.merge(thinking: thinking, metadata: internal_metadata, raw: raw).compact
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/responses/thinking_extractor.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module Legion
+  module Extensions
+    module Llm
+      module Responses
+        # Separates provider thinking markup from caller-visible text.
+        module ThinkingExtractor
+          Extraction = Struct.new(:content, :thinking, :signature, :metadata, keyword_init: true)
+
+          THINK_OPEN = '<think>'
+          THINK_CLOSE = '</think>'
+          THINK_PATTERN = %r{<think>(.*?)</think>}m
+          THINKING_METADATA_KEYS = %i[
+            reasoning_content reasoning thinking thinking_text thinking_signature reasoning_signature thought_signature
+          ].freeze
+          RAW_METADATA_KEYS = %i[
+            raw raw_response response_body provider_body provider_response
+          ].freeze
+
+          module_function
+
+          def extract(content, metadata: {})
+            metadata = normalized_metadata(metadata)
+            content, extracted_thinking = extract_from_content(content)
+            metadata_thinking = extract_metadata_thinking(metadata)
+            metadata_signature = extract_metadata_signature(metadata)
+
+            Extraction.new(
+              content: content,
+              thinking: compact_thinking([metadata_thinking, extracted_thinking]),
+              signature: metadata_signature,
+              metadata: scrub_metadata(metadata)
+            )
+          end
+
+          def extract_from_content(content)
+            return [content, nil] unless content.is_a?(String)
+
+            clean = +''
+            thinking_parts = []
+            remaining = content.dup
+
+            remaining = consume_next_segment(remaining, clean, thinking_parts) until remaining.empty?
+
+            [clean.strip, compact_thinking(thinking_parts)]
+          end
+          private_class_method :extract_from_content
+
+          def consume_next_segment(remaining, clean, thinking_parts)
+            close_index = remaining.index(THINK_CLOSE)
+            open_index = remaining.index(THINK_OPEN)
+
+            if close_index && (open_index.nil? || close_index < open_index)
+              thinking_parts << remaining.slice(0, close_index)
+              remaining.slice((close_index + THINK_CLOSE.length)..).to_s.sub(/\A[[:space:]]+/, '')
+            elsif open_index
+              consume_open_think_segment(remaining, open_index, clean, thinking_parts)
+            else
+              clean << remaining
+              +''
+            end
+          end
+          private_class_method :consume_next_segment
+
+          def consume_open_think_segment(remaining, open_index, clean, thinking_parts)
+            clean << remaining.slice(0, open_index)
+            after_open = remaining.slice((open_index + THINK_OPEN.length)..).to_s
+            close_index = after_open.index(THINK_CLOSE)
+            unless close_index
+              thinking_parts << after_open
+              return +''
+            end
+
+            thinking_parts << after_open.slice(0, close_index)
+            after_open.slice((close_index + THINK_CLOSE.length)..).to_s
+          end
+          private_class_method :consume_open_think_segment
+
+          def extract_metadata_thinking(metadata)
+            compact_thinking(
+              [
+                metadata[:reasoning_content],
+                metadata[:reasoning],
+                metadata[:thinking],
+                metadata[:thinking_text]
+              ]
+            )
+          end
+          private_class_method :extract_metadata_thinking
+
+          def extract_metadata_signature(metadata)
+            [
+              metadata[:thinking_signature],
+              metadata[:reasoning_signature],
+              metadata[:thought_signature]
+            ].compact.map { |signature| signature.to_s.strip }.find { |signature| !signature.empty? }
+          end
+          private_class_method :extract_metadata_signature
+
+          def scrub_metadata(metadata)
+            metadata.each_with_object({}) do |(key, value), scrubbed|
+              normalized_key = normalize_metadata_key(key)
+              next if THINKING_METADATA_KEYS.include?(normalized_key) || RAW_METADATA_KEYS.include?(normalized_key)
+
+              scrubbed[normalized_key] = scrub_metadata_value(value)
+            end
+          end
+          private_class_method :scrub_metadata
+
+          def normalize_metadata_key(key)
+            key.to_s
+               .gsub(/([a-z\d])([A-Z])/, '\1_\2')
+               .tr('-', '_')
+               .downcase
+               .to_sym
+          end
+          private_class_method :normalize_metadata_key
+
+          def scrub_metadata_value(value)
+            case value
+            when Hash
+              scrub_metadata(normalized_metadata(value))
+            when Array
+              value.map { |item| scrub_metadata_value(item) }
+            when String
+              extract_from_content(value).first
+            else
+              value
+            end
+          end
+          private_class_method :scrub_metadata_value
+
+          def normalized_metadata(metadata)
+            return {} if metadata.nil?
+
+            metadata.to_h.transform_keys { |key| normalize_metadata_key(key) }
+          end
+          private_class_method :normalized_metadata
+
+          def compact_thinking(parts)
+            text = parts.compact.map { |part| part.to_s.strip }.reject(&:empty?).join
+            blank_to_nil(text)
+          end
+          private_class_method :compact_thinking
+
+          def blank_to_nil(value)
+            value.nil? || value.empty? ? nil : value
+          end
+          private_class_method :blank_to_nil
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/stream_accumulator.rb

@@ -210,8 +210,11 @@ module Legion
         end
 
         def consume_non_think_content(remaining, start_tag, output)
+          unmatched_close = remaining.index('</think>')
           start_index = remaining.index(start_tag)
-          if start_index
+          if unmatched_close && (start_index.nil? || unmatched_close < start_index)
+            consume_unmatched_think_close(remaining, unmatched_close)
+          elsif start_index
             output << remaining.slice(0, start_index)
             @inside_think_tag = true
             remaining.slice((start_index + start_tag.length)..) || +''
@@ -223,6 +226,14 @@ module Legion
           end
         end
 
+        def consume_unmatched_think_close(remaining, close_index)
+          end_tag = '</think>'
+          thinking = remaining.slice(0, close_index)
+          @thinking_text << thinking
+          @last_thinking_delta << thinking
+          remaining.slice((close_index + end_tag.length)..).to_s.sub(/\A[[:space:]]+/, '')
+        end
+
         def longest_suffix_prefix(text, tag)
           max = [text.length, tag.length - 1].min
           max.downto(1) do |len|

data/lib/legion/extensions/llm/transport/exchanges/fleet.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+require 'legion/transport'
+
+module Legion
+  module Extensions
+    module Llm
+      module Transport
+        module Exchanges
+          # Shared topic exchange for live LLM fleet requests and replies.
+          class Fleet < ::Legion::Transport::Exchange
+            def exchange_name
+              'llm.fleet'
+            end
+
+            def default_type
+              'topic'
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/legion/extensions/llm/transport/messages/fleet_error.rb

@@ -0,0 +1,64 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require_relative '../../fleet/default_exchange_reply'
+require_relative '../../fleet/envelope_validation'
+require_relative '../../fleet/protocol'
+require_relative '../exchanges/fleet'
+
+module Legion
+  module Extensions
+    module Llm
+      module Transport
+        module Messages
+          # Correlated protocol-v2 error envelope for fleet reply queues.
+          class FleetError < ::Legion::Transport::Message
+            include Fleet::DefaultExchangeReply
+            include Fleet::EnvelopeValidation
+
+            def type = Fleet::Protocol::ERROR_TYPE
+            def app_id = @options[:app_id] || 'lex-llm'
+            def reply_to = @options[:reply_to]
+            def correlation_id = @options[:correlation_id]
+            def message_id = @options[:message_id] ||= "llm_fleet_err_#{SecureRandom.uuid}"
+
+            def routing_key
+              @options[:reply_to] || raise(ArgumentError, 'reply_to is required')
+            end
+
+            def validate
+              reject_legacy_options!
+              require_option!(:request_id)
+              require_option!(:correlation_id)
+              require_option!(:reply_to)
+              require_option!(:code)
+              require_protocol_version!
+              @valid = true
+            end
+
+            def message
+              super.merge(
+                protocol_version: @options[:protocol_version] || Fleet::Protocol::VERSION,
+                request_id: @options[:request_id],
+                correlation_id: correlation_id,
+                idempotency_key: @options[:idempotency_key],
+                operation: @options[:operation],
+                provider: @options[:provider],
+                provider_instance: @options[:provider_instance] || @options[:instance],
+                model: @options[:model],
+                reply_to: reply_to,
+                message_context: @options[:message_context],
+                trace_context: @options[:trace_context],
+                code: @options[:code],
+                message: @options[:message],
+                error_class: @options[:error_class],
+                retryable: @options[:retryable],
+                metadata: @options[:metadata] || {}
+              ).compact
+            end
+          end
+        end
+      end
+    end
+  end
+end