[Rails] Add Codebase Embeddings Search Service

Context

We need to introduce a service class that will perform a query on Code Embeddings given a target question.

This search service should allow a filter on the following fields:

  • project_id
  • path - we should allow prefix filtering on this field. IE, if the given path is app/models, then the search should include all embeddings with file_path = 'app/models/*'

References

Proposal

Create a new service class, e.g. CodeEmbeddings::SearchService, and follow the AI Abstraction Context Layer usage guide for finding documents closest to a given embedding.

Here is an example of how the search service class would look like:

(Note: this example code is based on the AI Abstraction Context Layer usage guide as of time of writing. Please make sure that the actual code follows the latest usage instructions.)

module CodeEmbeddings
  class SearchService
    def initialize(project_id:, path:, target_question:)
      @project_id = project_id
      @path = path
      @target_question = target_question
    end

    def execute
      target_embedding = ::ActiveContext::Embeddings.generate_embeddings(target_question)

      query = ActiveContext::Query.filter(project_id: 1).knn(target: "embeddings", vector: target_embedding)

      # or we may need to set limits to the result, e.g.:
      query = ActiveContext::Query.filter(project_id: 1).knn(target: "embeddings", vector: target_embedding, limit: 5)

      results = Ai::Context::Collections::MergeRequest.search(user: current_user, query: query)
      results.to_a
    end

    private

    attr_reader :project_id, :path, :target_question
  end
end
Edited by Pam Artiaga