make neighbor chainable (#219)

Co-authored-by: Vijay Swamidass <[email protected]>

Author: vswamidass-sfdc

app/controllers/base_documents_controller.rb

@@ -56,20 +56,17 @@ def index
       end
     end
 
+    # Apply text search before similarity search to maintain ActiveRecord relation
+    @documents = @documents.search_by_title_and_document(params[:contains]) if params[:contains].present?
+
     if params[:similar_to].present?
       embedding = get_embedding(params[:similar_to])
       # Get similar documents but preserve existing filters
-      similar_docs = related_documents_from_embedding_by_libraries(embedding, library_id)
-
-      # Sort by neighbor_distance (ascending - closest first) and convert to array
-      # This maintains the similarity order without needing raw SQL
-      @documents = similar_docs.sort_by(&:neighbor_distance)
+      @documents = @documents.related_by_embedding(embedding)
 
-      # Convert to Kaminari-compatible array for pagination
-      @documents = Kaminari.paginate_array(@documents)
+      # Sort by neighbor_distance using SQL to maintain ActiveRecord relation
+      @documents = @documents.order('neighbor_distance ASC')
     end
-
-    @documents = @documents.search_by_title_and_document(params[:contains]) if params[:contains].present?
     @documents = @documents.page(params[:page]).per(params[:per_page] || 10)
   end
 

app/controllers/concerns/neighbor_concern.rb

@@ -11,13 +11,8 @@ def related_documents_from_embedding(embedding)
     Document.nearest_neighbors(:embedding, embedding, distance: 'euclidean')
   end
 
-  # _limit is the number of documents to return.  Returning fewer is better since the most relevant documents are at the top.
-  # Because of the ordering, having too many documents may cause the most relevant documents to be lost.
-  # TODO implement a better way to get the most relevant documents.
-  def related_documents_from_embedding_by_libraries(_embedding, _library_ids, _limit = nil)
-    _limit ||= ENV.fetch('RELATED_DOCUMENTS_LIMIT', 25).to_i
-    scope = related_documents_from_embedding(_embedding)
-    scope = scope.where(library_id: _library_ids) if _library_ids.present?
-    scope.order(updated_at: :desc).limit(_limit)
+  # Delegate to the Document model scope for better organization
+  def related_documents_from_embedding_by_libraries(_embedding, __limit = nil)
+    Document.related_by_embedding(_embedding, limit: _limit)
   end
 end

app/models/document.rb

@@ -15,6 +15,15 @@ class Document < ApplicationRecord
                   }
 
   has_neighbors :embedding
+
+  # Scope to find related documents by embedding with optional library filtering
+  # _limit is the number of documents to return. Returning fewer is better since the most relevant documents are at the top.
+  # Because of the ordering, having too many documents may cause the most relevant documents to be lost.
+  scope :related_by_embedding, lambda { |embedding, limit = nil|
+    limit ||= ENV.fetch('RELATED_DOCUMENTS_LIMIT', 25).to_i
+    scope = nearest_neighbors(:embedding, embedding, distance: 'euclidean')
+    scope.order(updated_at: :desc).limit(limit)
+  }
   belongs_to :library, counter_cache: true
   belongs_to :user