# frozen_string_literal: true

# This tool searches for the most relevant *chunks* of text
# across all transcripts, or optionally within a specific transcript.
class ChunkSearch < RubyLLM::Tool
  description "Searches transcript *chunks* for specific information. Returns the text snippet and a precise source citation."
  param :query, desc: "The specific search query for finding relevant text snippets."

  def initialize(respondent_category: nil)
    @respondent_category = respondent_category
  end

  def execute(query:)
    model_map = FeatureModelMap.for_feature!("TRANSCRIPT_EMBEDDINGS")
    embedding = RubyLLM.embed(query, model: model_map.llm_model_name, provider: model_map.provider).vectors

    scope = Chunk.where(chunkable_type: "TranscriptSection",
                        chunkable_id: @respondent_category.transcript_sections.select(:id)) + Chunk.where(chunkable_type: "Respondent",
                                                                                                          chunkable_id: @respondent_category.respondents.select(:id))
    scope = Chunk.where(id: scope.collect(&:id))
    chunks = scope.nearest_neighbors(
      :embedding,
      embedding,
      distance: "euclidean"
    ).limit(5).includes(:chunkable)

    chunks.map do |chunk|
      <<~CONTEXT_BLOCK
        #{chunk.citation}
        Content:
        #{chunk.content}
      CONTEXT_BLOCK
    end.join("\n\n---\n\n")
  end
end