ragnar-cli 0.1.0.pre.1

data/lib/ragnar/query_processor.rb

@@ -0,0 +1,398 @@
+require 'json'
+require 'singleton'
+require 'set'
+require 'digest'
+
+module Ragnar
+  class QueryProcessor
+    attr_reader :database, :embedder, :rewriter, :reranker
+    
+    def initialize(db_path: Ragnar::DEFAULT_DB_PATH)
+      @database = Database.new(db_path)
+      @embedder = Embedder.new
+      @llm_manager = LLMManager.instance
+      @umap_service = UmapTransformService.instance
+      @rewriter = QueryRewriter.new(llm_manager: @llm_manager)
+      @reranker = nil # Will initialize when needed
+    end
+    
+    def query(user_query, top_k: 3, verbose: false)
+      puts "Processing query: #{user_query}" if verbose
+      
+      # Step 1: Rewrite and analyze the query
+      puts "\n#{'-'*60}" if verbose
+      puts "STEP 1: Query Analysis & Rewriting" if verbose
+      puts "-"*60 if verbose
+      
+      rewritten = @rewriter.rewrite(user_query)
+      
+      if verbose
+        puts "\nOriginal Query: #{user_query}"
+        puts "\nRewritten Query Analysis:"
+        puts "  Clarified Intent: #{rewritten['clarified_intent']}"
+        puts "  Query Type: #{rewritten['query_type']}"
+        puts "  Context Needed: #{rewritten['context_needed']}"
+        puts "\nGenerated Sub-queries (#{rewritten['sub_queries'].length}):"
+        rewritten['sub_queries'].each_with_index do |sq, idx|
+          puts "  #{idx + 1}. #{sq}"
+        end
+        if rewritten['key_terms'] && !rewritten['key_terms'].empty?
+          puts "\nKey Terms Identified:"
+          puts "  #{rewritten['key_terms'].join(', ')}"
+        end
+      end
+      
+      # Step 2: Retrieve candidates using RRF
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "STEP 2: Document Retrieval with RRF"
+        puts "-"*60
+      end
+      
+      candidates = retrieve_with_rrf(
+        rewritten['sub_queries'],
+        k: 20,
+        verbose: verbose
+      )
+      
+      if candidates.empty?
+        return {
+          query: user_query,
+          clarified: rewritten['clarified_intent'],
+          answer: "No relevant documents found in the database.",
+          sources: []
+        }
+      end
+      
+      if verbose
+        puts "\nRetrieval Summary:"
+        puts "  Total candidates found: #{candidates.size}"
+        puts "  Unique sources: #{candidates.map { |c| c[:file_path] }.uniq.size}"
+      end
+      
+      # Step 3: Rerank candidates
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "STEP 3: Document Reranking"
+        puts "-"*60
+      end
+      
+      reranked = rerank_documents(
+        query: rewritten['clarified_intent'],
+        documents: candidates,
+        top_k: top_k * 2  # Get more than we need for context
+      )
+      
+      if verbose && reranked.any?
+        puts "\nTop Reranked Documents:"
+        reranked[0..2].each_with_index do |doc, idx|
+          full_text = (doc[:chunk_text] || doc[:text] || "").gsub(/\s+/, ' ')
+          puts "  #{idx + 1}. [#{File.basename(doc[:file_path] || 'unknown')}]"
+          puts "     Score: #{doc[:score]&.round(4) if doc[:score]}"
+          puts "     Full chunk (#{full_text.length} chars):"
+          puts "     \"#{full_text}\""
+          puts ""
+        end
+      end
+      
+      # Step 4: Prepare context with neighboring chunks
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "STEP 4: Context Preparation"
+        puts "-"*60
+      end
+      
+      context_docs = prepare_context(reranked[0...top_k], rewritten['context_needed'])
+      
+      if verbose
+        puts "\nContext Documents Selected: #{context_docs.length}"
+        puts "Context strategy: #{rewritten['context_needed']}"
+      end
+      
+      # Step 5: Repack context for optimal LLM consumption
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "STEP 5: Context Repacking"
+        puts "-"*60
+      end
+      
+      repacked_context = ContextRepacker.repack(
+        context_docs,
+        rewritten['clarified_intent']
+      )
+      
+      if verbose
+        original_size = context_docs.sum { |d| (d[:chunk_text] || "").length }
+        puts "\nContext Optimization:"
+        puts "  Original size: #{original_size} chars"
+        puts "  Repacked size: #{repacked_context.length} chars"
+        puts "  Compression ratio: #{(100.0 * repacked_context.length / original_size).round(1)}%"
+        puts "\nFull Repacked Context:"
+        puts "-" * 40
+        puts repacked_context
+        puts "-" * 40
+      end
+      
+      # Step 6: Generate response
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "STEP 6: Response Generation"
+        puts "-"*60
+      end
+      response = generate_response(
+        query: rewritten['clarified_intent'],
+        repacked_context: repacked_context,
+        query_type: rewritten['query_type']
+      )
+      
+      if verbose
+        puts "\nGenerated Response:"
+        puts "-" * 40
+        puts response
+        puts "-" * 40
+      end
+      
+      result = {
+        query: user_query,
+        clarified: rewritten['clarified_intent'],
+        answer: response,
+        sources: context_docs.map { |d| 
+          {
+            source_file: d[:file_path] || d[:source_file],
+            chunk_index: d[:chunk_index]
+          }
+        },
+        sub_queries: rewritten['sub_queries'],
+        confidence: calculate_confidence(reranked[0...top_k])
+      }
+      
+      if verbose
+        puts "\n#{'-'*60}"
+        puts "FINAL RESULTS"
+        puts "-"*60
+        puts "\nConfidence Score: #{result[:confidence]}%"
+        puts "\nSources Used:"
+        result[:sources].each_with_index do |source, idx|
+          puts "  #{idx + 1}. #{source[:source_file]} (chunk #{source[:chunk_index]})"
+        end
+      end
+      
+      result
+    end
+    
+    private
+    
+    def retrieve_with_rrf(queries, k: 20, verbose: false)
+      all_results = []
+      
+      queries.each_with_index do |query, idx|
+        if verbose
+          puts "\nSub-query #{idx + 1}: \"#{query}\""
+          puts "  Generating embedding..."
+        end
+        
+        # Generate embedding for the query
+        query_embedding = @embedder.embed_text(query)
+        
+        if verbose
+          puts "  Embedding dimensions: #{query_embedding.length}"
+          puts "  Searching vector database..."
+        end
+        
+        # Check if we have reduced embeddings available for more efficient search
+        stats = @database.get_stats
+        use_reduced = stats[:with_reduced_embeddings] > 0
+        
+        # Prepare the search embedding (either full or reduced)
+        search_embedding = query_embedding
+        
+        if use_reduced
+          # Check if UMAP model is available
+          model_path = "./umap_model.bin"
+          
+          if @umap_service.model_available?(model_path)
+            if verbose
+              puts "  Transforming query to reduced space (#{stats[:reduced_dims]}D)"
+            end
+            
+            begin
+              # Transform the query embedding to reduced space
+              search_embedding = @umap_service.transform_query(query_embedding, model_path)
+              
+              if verbose
+                puts "  ✓ Query transformed to #{search_embedding.size}D"
+                puts "  Searching with reduced embeddings..."
+              end
+            rescue => e
+              puts "  ⚠️  Failed to transform query: #{e.message}" if verbose
+              puts "  Falling back to full embeddings" if verbose
+              use_reduced = false
+            end
+          else
+            if verbose
+              puts "  Note: Reduced embeddings available but UMAP model not found"
+              puts "  Falling back to full embeddings"
+            end
+            use_reduced = false
+          end
+        end
+        
+        vector_results = @database.search_similar(
+          search_embedding,
+          k: k,
+          use_reduced: use_reduced
+        )
+        
+        if verbose
+          puts "  Found #{vector_results.length} matches"
+          if vector_results.any?
+            best = vector_results.first
+            puts "  Best match: [#{File.basename(best[:file_path] || 'unknown')}] (distance: #{best[:distance]&.round(3)})"
+          end
+        end
+        
+        # Add query index for RRF
+        vector_results.each do |result|
+          result[:query_idx] = idx
+          result[:retrieval_method] = :vector
+        end
+        
+        all_results.concat(vector_results)
+      end
+      
+      if verbose
+        puts "\nApplying Reciprocal Rank Fusion..."
+        puts "  Total results before fusion: #{all_results.length}"
+      end
+      
+      # Apply Reciprocal Rank Fusion
+      fused = apply_rrf(all_results, k: k)
+      
+      if verbose
+        puts "  Results after RRF: #{fused.length}"
+      end
+      
+      fused
+    end
+    
+    def apply_rrf(results, k: 60)
+      # Group by document ID
+      doc_scores = {}
+      
+      results.each do |result|
+        doc_id = result[:id]
+        doc_scores[doc_id] ||= {
+          score: 0.0,
+          document: result
+        }
+        
+        # RRF formula: 1 / (k + rank)
+        # Using distance as a proxy for rank (lower distance = better rank)
+        rank = result[:distance] * 100  # Scale distance to rank-like values
+        doc_scores[doc_id][:score] += 1.0 / (k + rank)
+      end
+      
+      # Sort by RRF score and return documents
+      doc_scores.values
+        .sort_by { |item| -item[:score] }
+        .map { |item| item[:document] }
+    end
+    
+    def rerank_documents(query:, documents:, top_k:)
+      # Deduplicate documents based on chunk_text before reranking
+      seen_texts = Set.new
+      unique_docs = []
+      
+      documents.each do |doc|
+        text = doc[:chunk_text] || doc[:text] || ""
+        text_hash = Digest::SHA256.hexdigest(text)
+        
+        unless seen_texts.include?(text_hash)
+          seen_texts.add(text_hash)
+          unique_docs << doc
+        end
+      end
+      
+      if documents.length > unique_docs.length && @verbose
+        puts "  Deduplicated: #{documents.length} -> #{unique_docs.length} documents"
+      end
+      
+      # Initialize reranker if not already done
+      @reranker ||= Candle::Reranker.from_pretrained(
+        "cross-encoder/ms-marco-MiniLM-L-12-v2"
+      )
+      
+      # Prepare document texts - use chunk_text field
+      texts = unique_docs.map { |doc| doc[:chunk_text] || doc[:text] || "" }
+      
+      # Rerank - returns array of {doc_id:, score:, text:}
+      reranked = @reranker.rerank(query, texts)
+      
+      # Map back to original documents with scores
+      reranked.map do |result|
+        doc_idx = result[:doc_id]
+        unique_docs[doc_idx].merge(score: result[:score])
+      end.sort_by { |doc| -doc[:score] }.first(top_k)
+    rescue => e
+      puts "Warning: Reranking failed (#{e.message}), using original order"
+      unique_docs.first(top_k)
+    end
+    
+    def prepare_context(documents, context_needed)
+      # For now, just return the documents
+      # In the future, we could fetch neighboring chunks for more context
+      context_size = case context_needed
+                     when "extensive" then 5
+                     when "moderate" then 3
+                     else 2
+                     end
+      
+      documents.first(context_size)
+    end
+    
+    def generate_response(query:, repacked_context:, query_type:)
+      # Get cached LLM from manager
+      llm = @llm_manager.default_llm
+      
+      # Create prompt with repacked context
+      prompt = build_prompt(query, repacked_context, query_type)
+      
+      # Generate response using default config
+      llm.generate(prompt)
+    rescue => e
+      # Fallback to returning the repacked context
+      puts "Warning: LLM generation failed (#{e.message})"
+      "Based on the retrieved information:\n\n#{repacked_context[0..500]}..."
+    end
+    
+    def build_prompt(query, context, query_type)
+      base_prompt = <<~PROMPT
+        <|system|>
+        You are a helpful assistant. Answer questions based ONLY on the provided context.
+        If the answer is not in the context, say "I don't have enough information to answer that question."
+        </s>
+        <|user|>
+        Context:
+        #{context}
+        
+        Question: #{query}
+        </s>
+        <|assistant|>
+      PROMPT
+      
+      base_prompt
+    end
+    
+    def calculate_confidence(documents)
+      return 0.0 if documents.empty?
+      
+      # Simple confidence based on average similarity
+      avg_distance = documents.map { |d| d[:distance] }.sum / documents.size
+      
+      # Convert distance to confidence (0-1 scale)
+      # Assuming distances are typically 0-2
+      confidence = [1.0 - (avg_distance / 2.0), 0.0].max
+      (confidence * 100).round(1)
+    end
+  end
+end

data/lib/ragnar/query_rewriter.rb

@@ -0,0 +1,75 @@
+module Ragnar
+  class QueryRewriter
+    def initialize(llm_manager: nil)
+      @llm_manager = llm_manager || LLMManager.instance
+    end
+    
+    def rewrite(query)
+      # Get the cached LLM
+      model = @llm_manager.default_llm
+      
+      # Define the JSON schema for structured output
+      schema = {
+        type: "object",
+        properties: {
+          clarified_intent: {
+            type: "string",
+            description: "A clear, specific statement of what the user is looking for"
+          },
+          query_type: {
+            type: "string",
+            enum: ["factual", "conceptual", "procedural", "comparative", "analytical"],
+            description: "The type of query"
+          },
+          sub_queries: {
+            type: "array",
+            items: { type: "string" },
+            minItems: 2,
+            maxItems: 5,
+            description: "Simpler, focused queries that together answer the main query"
+          },
+          key_terms: {
+            type: "array",
+            items: { type: "string" },
+            description: "Important terms and their synonyms for searching"
+          },
+          context_needed: {
+            type: "string",
+            enum: ["minimal", "moderate", "extensive"],
+            description: "How much context is likely needed to answer this query"
+          }
+        },
+        required: ["clarified_intent", "query_type", "sub_queries", "key_terms", "context_needed"]
+      }
+      
+      prompt = <<~PROMPT
+        Analyze the following user query and break it down for retrieval-augmented generation.
+        Focus on understanding the user's intent and creating effective sub-queries for searching.
+        
+        User Query: #{query}
+        
+        Provide a structured analysis that will help retrieve the most relevant documents.
+      PROMPT
+      
+      begin
+        # Use structured generation with schema
+        result = model.generate_structured(
+          prompt,
+          schema: schema
+        )
+        
+        # The result should already be a JSON string
+        JSON.parse(result)
+      rescue => e
+        # Fallback to simple rewriting if structured generation fails
+        {
+          "clarified_intent" => query,
+          "query_type" => "general",
+          "sub_queries" => [query],
+          "key_terms" => query.split(/\s+/).select { |w| w.length > 3 },
+          "context_needed" => "moderate"
+        }
+      end
+    end
+  end
+end

data/lib/ragnar/topic_modeling/engine.rb

@@ -0,0 +1,221 @@
+require 'json'
+
+module Ragnar
+  module TopicModeling
+    class Engine
+      attr_reader :topics, :clusterer, :term_extractor
+      
+      def initialize(
+        min_cluster_size: 5,
+        min_samples: 3,
+        clustering_backend: nil,
+        reduce_dimensions: true,
+        n_components: 50,
+        labeling_method: :hybrid,
+        llm_client: nil,
+        verbose: false
+      )
+        @min_cluster_size = min_cluster_size
+        @min_samples = min_samples
+        @reduce_dimensions = reduce_dimensions
+        @n_components = n_components
+        @labeling_method = labeling_method
+        @verbose = verbose
+        
+        @clusterer = clustering_backend || build_default_clusterer
+        @term_extractor = TermExtractor.new
+        @labeler = TopicLabeler.new(method: labeling_method, llm_client: llm_client)
+        @topics = []
+      end
+      
+      def fit(embeddings:, documents:, metadata: nil)
+        raise ArgumentError, "Embeddings and documents must have same length" unless embeddings.length == documents.length
+        
+        @embeddings = embeddings
+        @documents = documents
+        @metadata = metadata || Array.new(documents.length) { {} }
+        
+        puts "Starting topic extraction..." if @verbose
+        
+        # Step 1: Optionally reduce dimensions for better clustering
+        working_embeddings = @embeddings
+        if @reduce_dimensions && @embeddings.first.length > @n_components
+          puts "  Reducing dimensions from #{@embeddings.first.length} to #{@n_components}..." if @verbose
+          working_embeddings = reduce_dimensions(@embeddings)
+        end
+        
+        # Step 2: Cluster embeddings
+        puts "  Clustering #{working_embeddings.length} documents..." if @verbose
+        cluster_ids = @clusterer.fit_predict(working_embeddings)
+        
+        # Step 3: Build topics from clusters
+        puts "  Building topics..." if @verbose
+        @topics = build_topics(cluster_ids)
+        
+        # Step 4: Extract terms for each topic
+        puts "  Extracting distinctive terms..." if @verbose
+        extract_topic_terms
+        
+        # Step 5: Generate labels
+        puts "  Generating topic labels..." if @verbose
+        generate_topic_labels
+        
+        puts "Found #{@topics.length} topics (plus #{count_outliers(cluster_ids)} outliers)" if @verbose
+        
+        @topics
+      end
+      
+      def transform(embeddings:, documents: nil)
+        # Assign new documents to existing topics
+        raise "Must call fit before transform" if @topics.empty?
+        
+        # Use approximate prediction if available
+        if @clusterer.respond_to?(:approximate_predict)
+          @clusterer.approximate_predict(embeddings)
+        else
+          # Fallback: assign to nearest topic centroid
+          assign_to_nearest_topic(embeddings)
+        end
+      end
+      
+      def get_topic(topic_id)
+        @topics.find { |t| t.id == topic_id }
+      end
+      
+      def outliers
+        @outliers ||= @documents.each_with_index.select { |_, idx| 
+          @cluster_ids && @cluster_ids[idx] == -1 
+        }.map(&:first)
+      end
+      
+      def save(path)
+        data = {
+          topics: @topics.map(&:to_h),
+          config: {
+            min_cluster_size: @min_cluster_size,
+            min_samples: @min_samples,
+            reduce_dimensions: @reduce_dimensions,
+            n_components: @n_components,
+            labeling_method: @labeling_method
+          }
+        }
+        File.write(path, JSON.pretty_generate(data))
+      end
+      
+      def self.load(path)
+        data = JSON.parse(File.read(path), symbolize_names: true)
+        engine = new(**data[:config])
+        # Reconstruct topics
+        engine.instance_variable_set(:@topics, data[:topics].map { |t| Topic.from_h(t) })
+        engine
+      end
+      
+      private
+      
+      def build_default_clusterer
+        begin
+          require 'clusterkit'
+          ClusterKit::Clustering::HDBSCAN.new(
+            min_cluster_size: @min_cluster_size,
+            min_samples: @min_samples,
+            metric: 'euclidean'
+          )
+        rescue LoadError
+          raise "ClusterKit required for topic modeling. Add 'gem \"clusterkit\"' to your Gemfile."
+        end
+      end
+      
+      def reduce_dimensions(embeddings)
+        require 'clusterkit'
+        
+        umap = ClusterKit::Dimensionality::UMAP.new(
+          n_components: @n_components,
+          n_neighbors: 15,
+          random_seed: 42  # For reproducibility
+        )
+        
+        # Convert to format UMAP expects
+        umap.fit_transform(embeddings)
+      rescue LoadError
+        puts "Warning: Dimensionality reduction requires ClusterKit. Using original embeddings." if @verbose
+        embeddings
+      end
+      
+      def build_topics(cluster_ids)
+        @cluster_ids = cluster_ids
+        
+        # Group documents by cluster
+        clusters = {}
+        cluster_ids.each_with_index do |cluster_id, doc_idx|
+          next if cluster_id == -1  # Skip outliers
+          clusters[cluster_id] ||= []
+          clusters[cluster_id] << doc_idx
+        end
+        
+        # Create Topic objects
+        clusters.map do |cluster_id, doc_indices|
+          Topic.new(
+            id: cluster_id,
+            document_indices: doc_indices,
+            documents: doc_indices.map { |i| @documents[i] },
+            embeddings: doc_indices.map { |i| @embeddings[i] },
+            metadata: doc_indices.map { |i| @metadata[i] }
+          )
+        end.sort_by(&:id)
+      end
+      
+      def extract_topic_terms
+        # Extract distinctive terms for each topic
+        all_docs_text = @documents.join(" ")
+        
+        @topics.each do |topic|
+          topic_docs_text = topic.documents.join(" ")
+          
+          # Use c-TF-IDF to find distinctive terms
+          terms = @term_extractor.extract_distinctive_terms(
+            topic_docs: topic.documents,
+            all_docs: @documents,
+            top_n: 20
+          )
+          
+          topic.set_terms(terms)
+        end
+      end
+      
+      def generate_topic_labels
+        @topics.each do |topic|
+          result = @labeler.generate_label(
+            topic: topic,
+            terms: topic.terms,
+            documents: topic.documents.first(3)  # Use top 3 representative docs
+          )
+          
+          # Set both label and description if available
+          topic.set_label(result[:label])
+          topic.instance_variable_set(:@description, result[:description]) if result[:description]
+          topic.instance_variable_set(:@label_confidence, result[:confidence])
+          topic.instance_variable_set(:@themes, result[:themes]) if result[:themes]
+        end
+      end
+      
+      def count_outliers(cluster_ids)
+        cluster_ids.count { |id| id == -1 }
+      end
+      
+      def assign_to_nearest_topic(embeddings)
+        # Simple nearest centroid assignment
+        topic_centroids = @topics.map(&:centroid)
+        
+        embeddings.map do |embedding|
+          distances = topic_centroids.map do |centroid|
+            # Euclidean distance
+            Math.sqrt(embedding.zip(centroid).map { |a, b| (a - b) ** 2 }.sum)
+          end
+          
+          min_idx = distances.index(distances.min)
+          @topics[min_idx].id
+        end
+      end
+    end
+  end
+end