Class: EmbeddingUtil::CLI

Inherits:
Thor
  • Object
show all
Defined in:
lib/embedding_util/cli.rb

Constant Summary collapse

CONFIG_OPTIONS = 
{
  profile: :to_sym.to_proc,
  runtime: ->(value) { RuntimeCommand.normalize_runtime(value) },
  endpoint: ->(value) { value },
  embedding_endpoint: ->(value) { value },
  reranker_endpoint: ->(value) { value },
  timeout: ->(value) { value },
  startup_timeout: ->(value) { value },
  shutdown_idle: :to_i.to_proc,
  verbose: ->(value) { value }
}.freeze

Instance Method Summary collapse

Instance Method Details

#embed(text) ⇒ Object 



51
52
53
54
55
56
 
# File 'lib/embedding_util/cli.rb', line 51

def embed(text)
  configure_embedding_util
  puts JSON.generate(EmbeddingUtil.embed(text))
rescue Error => e
  abort e.message
end

#profiles ⇒ Object 



42
43
44
45
46
47
48
 
# File 'lib/embedding_util/cli.rb', line 42

def profiles
  EmbeddingUtil.profiles.each do |profile|
    puts profile.name
    puts "  embedding: #{profile.embedding.fetch(:repo)} / #{profile.embedding.fetch(:file)}"
    puts "  reranker:  #{profile.reranker.fetch(:repo)} / #{profile.reranker.fetch(:file)}"
  end
end

#rerank(query, *documents) ⇒ Object 



59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
 
# File 'lib/embedding_util/cli.rb', line 59

def rerank(query, *documents)
  configure_embedding_util
  raise Error, "provide at least one document to rerank" if documents.empty?

  results = EmbeddingUtil.rerank(query, documents).map do |result|
    {
      index: result.index,
      document: result.document,
      score: result.score,
      metadata: result.metadata
    }
  end
  puts JSON.pretty_generate(results)
rescue Error => e
  abort e.message
end

#serve ⇒ Object 



81
82
83
84
85
86
87
88
89
90
 
# File 'lib/embedding_util/cli.rb', line 81

def serve
  configure_embedding_util
  ServerManager.new(config: EmbeddingUtil.configuration).serve(
    model: options[:model],
    runtime: options[:runtime] || EmbeddingUtil.configuration.runtime,
    shutdown_idle: options[:shutdown_idle]&.to_i,
    host: options[:host],
    port: options[:port]&.to_i
  )
end

#support ⇒ Object 



31
32
33
34
35
36
37
38
39
 
# File 'lib/embedding_util/cli.rb', line 31

def support
  configure_embedding_util
  EmbeddingUtil.support.each do |item|
    status = item.fetch(:supported) ? "supported" : "not supported"
    puts "#{item.fetch(:provider)}: #{status}"
    puts "  embedding_endpoint: #{item.fetch(:embedding_endpoint)}" if item[:embedding_endpoint]
    puts "  reranker_endpoint: #{item.fetch(:reranker_endpoint)}" if item[:reranker_endpoint]
  end
end