embedding_util 0.1.0

data/lib/embedding_util.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+require_relative "embedding_util/version"
+
+module EmbeddingUtil
+  class Error < StandardError; end
+  class UnsupportedProviderError < Error; end
+  class EndpointError < Error; end
+
+  class EndpointNotFoundError < EndpointError
+    attr_reader :uri, :path, :body
+
+    def initialize(uri, path:, body: nil)
+      @uri = uri
+      @path = path
+      @body = body
+      super("#{uri} returned 404")
+    end
+  end
+
+  autoload :Configuration, "embedding_util/configuration"
+  autoload :CLI, "embedding_util/cli"
+  autoload :EmbeddingResult, "embedding_util/result"
+  autoload :Profile, "embedding_util/profile"
+  autoload :Profiles, "embedding_util/profiles"
+  autoload :Provider, "embedding_util/provider"
+  autoload :ProviderRegistry, "embedding_util/provider_registry"
+  autoload :RankedDocument, "embedding_util/result"
+  autoload :RerankResult, "embedding_util/result"
+  autoload :ServerManager, "embedding_util/server_manager"
+  autoload :ServerModel, "embedding_util/server_model"
+  autoload :RuntimeCommand, "embedding_util/runtime_command"
+
+  module Providers
+    autoload :Endpoint, "embedding_util/providers/endpoint"
+    autoload :SelfHosted, "embedding_util/providers/self_hosted"
+  end
+
+  module_function
+
+  def configuration
+    @configuration ||= Configuration.new
+  end
+
+  def configure
+    yield configuration
+  end
+
+  def reset_configuration!
+    @configuration = Configuration.new
+    @registry = nil
+  end
+
+  def registry
+    @registry ||= begin
+      registry = ProviderRegistry.new
+      registry.register(Providers::Endpoint)
+      registry.register(Providers::SelfHosted)
+      registry
+    end
+  end
+
+  def register_provider(provider_class)
+    registry.register(provider_class)
+  end
+
+  def support
+    registry.support(config: configuration)
+  end
+
+  def embed(text, **options)
+    embed_result(text, **options).embedding
+  end
+
+  def embed_many(texts, profile: configuration.resolved_profile, provider: nil, **_options)
+    selected_provider(provider).embed(normalize_texts(texts), profile: resolve_profile(profile)).embedding
+  end
+
+  def embed_result(input, profile: configuration.resolved_profile, provider: nil)
+    scalar = !input.is_a?(Array)
+    texts = normalize_texts(input)
+    result = selected_provider(provider).embed(texts, profile: resolve_profile(profile))
+    return result unless scalar
+
+    raise EndpointError, "server returned no embeddings for the given input" if result.embedding.empty?
+
+    EmbeddingResult.new(
+      embedding: result.embedding.fetch(0),
+      model: result.model,
+      profile: result.profile,
+      provider: result.provider,
+      metadata: result.metadata
+    )
+  end
+
+  def rerank(query, documents, **options)
+    rerank_result(query, documents, **options).results
+  end
+
+  def rerank_result(query, documents, profile: configuration.resolved_profile, provider: nil)
+    selected_provider(provider).rerank(query.to_s, Array(documents).map(&:to_s), profile: resolve_profile(profile))
+  end
+
+  def profiles
+    Profiles.all
+  end
+
+  def profile(name = configuration.profile)
+    resolve_profile(name)
+  end
+
+  def normalize_texts(input)
+    input.is_a?(Array) ? input.map(&:to_s) : [input.to_s]
+  end
+
+  def resolve_profile(value)
+    value.is_a?(Profile) ? value : Profiles.fetch(value)
+  end
+
+  def selected_provider(provider)
+    return registry.resolve(config: configuration) unless provider
+
+    local_config = configuration.dup
+    local_config.provider = provider
+    registry.resolve(config: local_config)
+  end
+end

metadata

@@ -0,0 +1,83 @@
+--- !ruby/object:Gem::Specification
+name: embedding_util
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- hmdne
+bindir: exe
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: thor
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+description: A small rbutils gem for computing embeddings and true reranking through
+  local embedding model runtimes.
+email:
+- 54514036+hmdne@users.noreply.github.com
+executables:
+- embedding_util
+extensions: []
+extra_rdoc_files: []
+files:
+- ".rspec"
+- ".rubocop.yml"
+- CHANGELOG.md
+- LICENSE.txt
+- README.md
+- Rakefile
+- exe/embedding_util
+- lib/embedding_util.rb
+- lib/embedding_util/cli.rb
+- lib/embedding_util/configuration.rb
+- lib/embedding_util/profile.rb
+- lib/embedding_util/profiles.rb
+- lib/embedding_util/provider.rb
+- lib/embedding_util/provider_registry.rb
+- lib/embedding_util/providers/endpoint.rb
+- lib/embedding_util/providers/self_hosted.rb
+- lib/embedding_util/result.rb
+- lib/embedding_util/runtime_command.rb
+- lib/embedding_util/server_manager.rb
+- lib/embedding_util/server_model.rb
+- lib/embedding_util/version.rb
+homepage: https://github.com/rbutils/embedding_util
+licenses:
+- MIT
+metadata:
+  allowed_push_host: https://rubygems.org
+  source_code_uri: https://github.com/rbutils/embedding_util
+  changelog_uri: https://github.com/rbutils/embedding_util/blob/master/CHANGELOG.md
+  documentation_uri: https://github.com/rbutils/embedding_util#readme
+  bug_tracker_uri: https://github.com/rbutils/embedding_util/issues
+  rubygems_mfa_required: 'true'
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.2.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 4.0.6
+specification_version: 4
+summary: Local-first text embeddings and reranking for Ruby
+test_files: []