Class: Maglev::Retriever

Inherits:
Object
  • Object
show all
Defined in:
lib/maglev/retriever.rb

Instance Method Summary collapse

Constructor Details

#initialize(model_class, chunk_model: Chunk, embedding_adapter: Maglev.configuration.embedding_adapter, authorization: Authorization.new, vector_store: Maglev.configuration.vector_store) ⇒ Retriever

Returns a new instance of Retriever.



12
13
14
15
16
17
18
# File 'lib/maglev/retriever.rb', line 12

def initialize(model_class, chunk_model: Chunk, embedding_adapter: Maglev.configuration.embedding_adapter, authorization: Authorization.new, vector_store: Maglev.configuration.vector_store)
  @model_class = model_class
  @chunk_model = chunk_model
  @embedding_adapter = embedding_adapter || Adapters::RubyLLMEmbedding.new
  @authorization = authorization
  @vector_store = vector_store
end

Instance Method Details

#search(query, limit:, owner: nil, user: nil) ⇒ Object



20
21
22
23
24
25
26
27
28
29
30
31
# File 'lib/maglev/retriever.rb', line 20

def search(query, limit:, owner: nil, user: nil)
  embedding = ProviderCall.new.call(operation: "embed") { @embedding_adapter.embed(query) }
  return search_vector_store(embedding, limit: limit, owner: owner, user: user) if @vector_store

  scope = @chunk_model.where(owner_model_name: @model_class.name)
  scope = scope.where(owner: owner) if owner
  scope = apply_authorization_scope(scope, user) unless owner
  rows = scope
    .nearest_neighbors(:embedding, embedding, distance: "cosine")

  unique_owner_results(rows).first(limit)
end