embedding_util 0.1.0

data/lib/embedding_util/providers/endpoint.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+require "json"
+require "net/http"
+require "uri"
+
+require_relative "../provider"
+require_relative "../result"
+
+module EmbeddingUtil
+  module Providers
+    class Endpoint < Provider
+      NETWORK_ERRORS = [
+        Errno::ECONNREFUSED,
+        Errno::ECONNRESET,
+        Errno::EHOSTUNREACH,
+        Errno::ENETUNREACH,
+        EOFError,
+        IOError,
+        Net::OpenTimeout,
+        Net::ReadTimeout,
+        SocketError,
+        Timeout::Error
+      ].freeze
+
+      def supported?
+        !!(config.embedding_endpoint_url || config.reranker_endpoint_url)
+      end
+
+      def support
+        {
+          provider: provider_name,
+          supported: supported?,
+          embedding_endpoint: config.embedding_endpoint_url,
+          reranker_endpoint: config.reranker_endpoint_url
+        }
+      end
+
+      def embed(texts, profile: config.resolved_profile)
+        endpoint = require_endpoint(config.embedding_endpoint_url, "embedding")
+        response = post_json(endpoint, "/v1/embeddings", {
+                               input: texts,
+                               model: profile.embedding.fetch(:model)
+                             })
+
+        data = Array(response.fetch("data"))
+        embeddings = data.sort_by { |item| item.fetch("index", data.index(item) || 0) }.map { |item| item.fetch("embedding") }
+        EmbeddingResult.new(
+          embedding: embeddings,
+          model: response["model"],
+          profile: profile.name,
+          provider: provider_name,
+          metadata: { usage: response["usage"] }.compact
+        )
+      end
+
+      def rerank(query, documents, profile: config.resolved_profile)
+        endpoint = require_endpoint(config.reranker_endpoint_url, "reranker")
+        response = begin
+          post_json(endpoint, "/v1/rerank", rerank_payload(query, documents, profile))
+        rescue EndpointNotFoundError => e
+          raise unless fallback_rerank_not_found?(e)
+
+          post_json(endpoint, "/rerank", rerank_payload(query, documents, profile))
+        end
+
+        RerankResult.new(
+          results: ranked_documents(response, documents),
+          model: response["model"],
+          profile: profile.name,
+          provider: provider_name,
+          metadata: { usage: response["usage"] }.compact
+        )
+      end
+
+      private
+
+      def rerank_payload(query, documents, profile)
+        {
+          query: query,
+          documents: documents,
+          model: profile.reranker.fetch(:model)
+        }
+      end
+
+      def ranked_documents(response, documents)
+        Array(response.fetch("results")).map do |item|
+          index = item.fetch("index")
+          RankedDocument.new(
+            index: index,
+            document: item["document"] || fetch_document_at(documents, index),
+            score: item.fetch("relevance_score") { item.fetch("score") },
+            metadata: item.reject { |key, _value| %w[index document relevance_score score].include?(key) }
+          )
+        end
+      end
+
+      def fetch_document_at(documents, index)
+        documents.fetch(index)
+      rescue IndexError
+        raise EndpointError, "server returned out-of-range document index #{index.inspect} (#{documents.size} documents sent)"
+      end
+
+      def require_endpoint(endpoint, capability)
+        raise UnsupportedProviderError, "no #{capability} endpoint configured" unless endpoint
+
+        endpoint
+      end
+
+      def post_json(endpoint, path, payload)
+        uri = endpoint_uri(endpoint, path)
+        request = Net::HTTP::Post.new(uri)
+        request["Content-Type"] = "application/json"
+        request.body = JSON.generate(payload)
+
+        response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == "https", read_timeout: config.timeout, open_timeout: config.timeout) do |http|
+          http.request(request)
+        end
+
+        raise EndpointNotFoundError.new(uri, path: path, body: response.body) if response.code.to_i == 404 && route_missing_response?(response.body)
+        raise EndpointError, "#{uri} returned #{response.code}: #{response.body}" unless response.is_a?(Net::HTTPSuccess)
+
+        JSON.parse(response.body)
+      rescue JSON::ParserError => e
+        raise EndpointError, "invalid JSON response from #{uri}: #{e.message}"
+      rescue URI::InvalidURIError => e
+        raise EndpointError, "invalid endpoint URL #{endpoint.inspect}: #{e.message}"
+      rescue *NETWORK_ERRORS => e
+        raise EndpointError, "could not reach #{uri}: #{e.message}"
+      end
+
+      def endpoint_uri(endpoint, path)
+        uri = URI(endpoint)
+        segments = [uri.path, path].map { |part| part.to_s.gsub(%r{\A/+|/+\z}, "") }.reject(&:empty?)
+        uri.path = "/#{segments.join('/')}"
+        uri
+      end
+
+      def route_missing_response?(body)
+        return true if body.to_s.strip.empty?
+
+        JSON.parse(body)
+        false
+      rescue JSON::ParserError
+        true
+      end
+
+      def fallback_rerank_not_found?(error)
+        error.path == "/v1/rerank"
+      end
+    end
+  end
+end

data/lib/embedding_util/providers/self_hosted.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require_relative "../provider"
+require_relative "../server_manager"
+require_relative "endpoint"
+
+module EmbeddingUtil
+  module Providers
+    class SelfHosted < Provider
+      def supported?
+        ServerManager.supported?(config)
+      end
+
+      def support
+        {
+          provider: provider_name,
+          supported: supported?,
+          runtime: RuntimeCommand.resolve(config.runtime),
+          shutdown_idle: config.shutdown_idle,
+          state_dir: config.state_dir
+        }
+      end
+
+      def embed(texts, profile: config.resolved_profile)
+        endpoint = ServerManager.new(config: config).ensure_server(:embedding, profile: profile)
+        endpoint_provider(embedding_endpoint: endpoint).embed(texts, profile: profile)
+      end
+
+      def rerank(query, documents, profile: config.resolved_profile)
+        endpoint = ServerManager.new(config: config).ensure_server(:reranker, profile: profile)
+        endpoint_provider(reranker_endpoint: endpoint).rerank(query, documents, profile: profile)
+      end
+
+      private
+
+      def endpoint_provider(embedding_endpoint: nil, reranker_endpoint: nil)
+        endpoint_config = config.dup
+        endpoint_config.embedding_endpoint = embedding_endpoint
+        endpoint_config.reranker_endpoint = reranker_endpoint
+        Endpoint.new(config: endpoint_config)
+      end
+    end
+  end
+end

data/lib/embedding_util/result.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module EmbeddingUtil
+  EmbeddingResult = Data.define(:embedding, :model, :profile, :provider, :metadata)
+  RankedDocument = Data.define(:index, :document, :score, :metadata)
+  RerankResult = Data.define(:results, :model, :profile, :provider, :metadata)
+end

data/lib/embedding_util/runtime_command.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+module EmbeddingUtil
+  class RuntimeCommand
+    attr_reader :runtime, :server_model, :host, :port
+
+    def initialize(runtime:, server_model:, host:, port:)
+      @runtime = self.class.normalize_runtime(runtime)
+      @server_model = server_model
+      @host = host
+      @port = port
+    end
+
+    def self.available?(runtime)
+      case normalize_runtime(runtime)
+      when :auto
+        available?(:ramalama) || available?(:llama_server)
+      when :ramalama
+        !!command_path("ramalama")
+      when :llama_server
+        !!command_path("llama-server")
+      else
+        false
+      end
+    end
+
+    def self.resolve(runtime)
+      requested = normalize_runtime(runtime)
+      return requested unless requested == :auto
+
+      return :ramalama if available?(:ramalama)
+      return :llama_server if available?(:llama_server)
+
+      :auto
+    end
+
+    def self.command_path(command)
+      ENV.fetch("PATH", "").split(File::PATH_SEPARATOR).map { |dir| File.join(dir, command) }.find { |path| File.executable?(path) && !File.directory?(path) }
+    end
+
+    def self.normalize_runtime(runtime)
+      runtime.to_s.tr("-", "_").to_sym
+    end
+
+    def argv
+      case runtime
+      when :ramalama then ramalama_argv
+      when :llama_server then llama_server_argv
+      else raise UnsupportedProviderError, "no supported local runtime found; install ramalama or llama-server"
+      end
+    end
+
+    def label
+      runtime == :llama_server ? "llama-server" : runtime.to_s
+    end
+
+    private
+
+    def ramalama_argv
+      [
+        "ramalama", "--runtime=llama.cpp", "serve",
+        "--host", host,
+        "--port", port.to_s,
+        "--runtime-args=#{server_model.settings.fetch(:server_flags).join(' ')}",
+        huggingface_model
+      ]
+    end
+
+    def llama_server_argv
+      [
+        "llama-server",
+        "--host", host,
+        "--port", port.to_s,
+        "-hf", server_model.settings.fetch(:repo),
+        "-hff", server_model.settings.fetch(:file),
+        *server_model.settings.fetch(:server_flags)
+      ]
+    end
+
+    def huggingface_model
+      "hf://#{server_model.settings.fetch(:repo)}/#{server_model.settings.fetch(:file)}"
+    end
+  end
+end

data/lib/embedding_util/server_manager.rb

@@ -0,0 +1,258 @@
+# frozen_string_literal: true
+
+require "fileutils"
+require "json"
+require "net/http"
+require "open3"
+require "rbconfig"
+require "socket"
+require "time"
+require "uri"
+
+module EmbeddingUtil
+  class ServerManager
+    attr_reader :config
+
+    def initialize(config: EmbeddingUtil.configuration)
+      @config = config
+    end
+
+    def self.supported?(config = EmbeddingUtil.configuration)
+      RuntimeCommand.available?(config.runtime)
+    end
+
+    def ensure_server(capability, profile: config.resolved_profile)
+      server_model = ServerModel.for(capability, profile)
+      log_path = server_log_path(server_model)
+
+      with_lock(server_model) do
+        state = read_state(server_model)
+        log_path = start_background(server_model) unless healthy_state?(state) || running_server_state?(state)
+      end
+
+      wait_for_healthy(server_model, log_path: log_path)
+    end
+
+    def serve(model:, runtime: config.runtime, shutdown_idle: config.shutdown_idle, host: config.host, port: nil)
+      server_model = model.is_a?(ServerModel) ? model : ServerModel.parse(model)
+      resolved_runtime = RuntimeCommand.resolve(runtime)
+      selected_port = selected_port_for(server_model, host: host, port: port)
+      command = RuntimeCommand.new(runtime: resolved_runtime, server_model: server_model, host: host, port: selected_port)
+      last_output_at = Time.now
+
+      FileUtils.mkdir_p(config.state_dir)
+      puts "starting #{server_model.name} with #{command.label} on http://#{host}:#{selected_port}"
+      puts "shutdown idle: #{shutdown_idle}s" if shutdown_idle&.positive?
+
+      Open3.popen2e(*command.argv) do |_stdin, output, wait_thread|
+        write_state(server_model, pid: wait_thread.pid, url: "http://#{host}:#{selected_port}", runtime: command.label, port: selected_port)
+        watchdog = start_watchdog(wait_thread.pid, shutdown_idle) { last_output_at }
+
+        output.each_line do |line|
+          last_output_at = Time.now
+          print line
+        end
+
+        watchdog&.kill
+        delete_state(server_model)
+        wait_thread.value.exitstatus
+      end
+    end
+
+    private
+
+    def start_background(server_model)
+      FileUtils.mkdir_p(config.state_dir)
+      log_path = server_log_path(server_model)
+      selected_port = selected_port_for(server_model, host: config.host)
+      argv = [
+        RbConfig.ruby, executable_path, "serve",
+        "--model", server_model.name,
+        "--runtime", config.runtime.to_s,
+        "--host", config.host,
+        "--port", selected_port.to_s
+      ]
+      argv.push("--shutdown-idle", config.shutdown_idle.to_s) unless config.shutdown_idle.nil?
+      warn "starting #{server_model.name} in background: #{argv.join(' ')}" if config.verbose
+      warn "#{server_model.name} log: #{log_path}" if config.verbose
+      pid = Process.spawn(*argv, out: [log_path, "a"], err: %i[child out], pgroup: true)
+      write_state(server_model, pid: pid, url: "http://#{config.host}:#{selected_port}", runtime: "starting", port: selected_port)
+      Process.detach(pid)
+      log_path
+    end
+
+    def server_log_path(server_model)
+      File.join(config.state_dir, "#{server_model.name}.log")
+    end
+
+    def executable_path
+      local_path = File.expand_path("../../exe/embedding_util", __dir__)
+      return local_path if File.exist?(local_path)
+
+      Gem.bin_path("embedding_util", "embedding_util")
+    end
+
+    def selected_port_for(server_model, host:, port: nil)
+      return required_port(host, port) if port
+
+      available_port(host, server_model.default_port(config))
+    end
+
+    def required_port(host, port)
+      return port if port_available?(host, port)
+
+      raise UnsupportedProviderError, "port #{host}:#{port} is already in use"
+    end
+
+    def available_port(host, preferred_port)
+      (preferred_port...(preferred_port + 100)).find { |candidate| port_available?(host, candidate) } || raise(
+        UnsupportedProviderError,
+        "no free port found for #{host} starting at #{preferred_port}"
+      )
+    end
+
+    def port_available?(host, port)
+      # Advisory only: the child runtime performs the real bind after this process releases the port.
+      # with_lock serializes callers within a single process, but a cross-process race window
+      # still exists between the probe socket closing here and the child process binding.
+      server = TCPServer.new(host, port)
+      true
+    rescue Errno::EADDRINUSE, Errno::EACCES, SocketError
+      false
+    ensure
+      server&.close
+    end
+
+    def wait_for_healthy(server_model, log_path: nil)
+      deadline = Time.now + config.startup_timeout
+      loop do
+        state = read_state(server_model)
+        return state.fetch("url") if healthy_state?(state)
+        raise UnsupportedProviderError, process_exited_message(server_model, log_path) if tracked_process_exited?(state)
+        raise UnsupportedProviderError, startup_timeout_message(server_model, log_path) if Time.now >= deadline
+
+        sleep 0.25
+      end
+    end
+
+    def start_watchdog(pid, shutdown_idle)
+      return unless shutdown_idle&.positive?
+
+      Thread.new do
+        loop do
+          sleep [shutdown_idle / 5.0, 1].max
+          next if Time.now - yield < shutdown_idle
+
+          terminate_idle_process(pid)
+        rescue Errno::ESRCH
+          break
+        end
+      end
+    end
+
+    def terminate_idle_process(pid)
+      Process.kill("TERM", pid)
+      sleep 5
+      Process.kill("KILL", pid) if process_running?(pid)
+    end
+
+    def startup_timeout_message(server_model, log_path)
+      message = "timed out after #{config.startup_timeout}s waiting for #{server_model.name} to become healthy"
+      return message unless log_path
+
+      lines = log_tail(log_path)
+      message += "\nlog: #{log_path}"
+      message += "\nlast log lines:\n#{lines}" unless lines.empty?
+      message
+    end
+
+    def process_exited_message(server_model, log_path)
+      message = "#{server_model.name} server process exited before becoming healthy"
+      return message unless log_path
+
+      lines = log_tail(log_path)
+      message += "\nlog: #{log_path}"
+      message += "\nlast log lines:\n#{lines}" unless lines.empty?
+      message
+    end
+
+    def log_tail(log_path)
+      return "" unless File.exist?(log_path)
+
+      File.readlines(log_path).last(20).join
+    rescue Errno::ENOENT, Errno::EACCES, IOError
+      ""
+    end
+
+    def healthy_state?(state)
+      return false unless state && state["url"] && state["pid"]
+      return false unless process_running?(state.fetch("pid"))
+
+      healthy_url?(state.fetch("url"))
+    end
+
+    def running_server_state?(state)
+      state && state["url"] && state["pid"] && process_running?(state.fetch("pid"))
+    end
+
+    def tracked_process_exited?(state)
+      state && state["url"] && state["pid"] && !process_running?(state.fetch("pid"))
+    end
+
+    def healthy_url?(url)
+      uri = URI.join(url.end_with?("/") ? url : "#{url}/", "health")
+      response = Net::HTTP.start(uri.hostname, uri.port, use_ssl: uri.scheme == "https", read_timeout: 2, open_timeout: 2) { |http| http.get(uri) }
+      response.is_a?(Net::HTTPSuccess)
+    rescue StandardError
+      false
+    end
+
+    def process_running?(pid)
+      Process.kill(0, Integer(pid))
+      true
+    rescue Errno::ESRCH, ArgumentError
+      false
+    rescue Errno::EPERM
+      # Process exists but belongs to a different user; treat as running.
+      true
+    end
+
+    def with_lock(server_model)
+      FileUtils.mkdir_p(config.state_dir)
+      File.open(lock_path(server_model), File::RDWR | File::CREAT, 0o644) do |file|
+        file.flock(File::LOCK_EX)
+        yield
+      end
+    end
+
+    def write_state(server_model, pid:, url:, runtime:, port:)
+      File.write(state_path(server_model), JSON.pretty_generate({
+                                                                  pid: pid,
+                                                                  url: url,
+                                                                  profile: server_model.profile.name,
+                                                                  capability: server_model.capability,
+                                                                  runtime: runtime,
+                                                                  port: port,
+                                                                  updated_at: Time.now.utc.iso8601
+                                                                }))
+    end
+
+    def read_state(server_model)
+      JSON.parse(File.read(state_path(server_model)))
+    rescue Errno::ENOENT, JSON::ParserError
+      nil
+    end
+
+    def delete_state(server_model)
+      FileUtils.rm_f(state_path(server_model))
+    end
+
+    def state_path(server_model)
+      File.join(config.state_dir, "#{server_model.name}.json")
+    end
+
+    def lock_path(server_model)
+      File.join(config.state_dir, "#{server_model.name}.lock")
+    end
+  end
+end

data/lib/embedding_util/server_model.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+module EmbeddingUtil
+  SERVER_MODEL_PREFIXES = {
+    "embedding" => :embedding,
+    "reranker" => :reranker,
+    "rerank" => :reranker
+  }.freeze
+
+  ServerModel = Data.define(:capability, :profile) do
+    def self.parse(value)
+      text = value.to_s
+      prefix, profile_name = text.split("-", 2)
+      capability = SERVER_MODEL_PREFIXES[prefix]
+      raise ArgumentError, "unknown server model #{value.inspect}; expected embedding-PROFILE or reranker-PROFILE" unless capability && profile_name
+
+      new(capability: capability, profile: Profiles.fetch(profile_name))
+    end
+
+    def self.for(capability, profile)
+      new(capability: capability.to_sym, profile: profile)
+    end
+
+    def name
+      "#{capability_name}-#{profile.name}"
+    end
+
+    def settings
+      case capability
+      when :embedding then profile.embedding
+      when :reranker then profile.reranker
+      else raise ArgumentError, "unknown server capability: #{capability.inspect}"
+      end
+    end
+
+    def default_port(config)
+      capability == :embedding ? config.embedding_port : config.reranker_port
+    end
+
+    private
+
+    def capability_name
+      capability == :reranker ? "reranker" : capability.to_s
+    end
+  end
+end

data/lib/embedding_util/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module EmbeddingUtil
+  VERSION = "0.1.0"
+end