@restnfeel/agentc-starter-kit/src/rag/retrieval/hybrid-retriever.ts

한국어 기업용 CMS 모듈 - Task Master AI와 함께 빠르게 웹사이트를 구현할 수 있는 재사용 가능한 컴포넌트 시스템

import { SearchResult, VectorStore, EmbeddingModel } from "../types";
import { QueryProcessor, ProcessedQuery } from "./query-processor";

export interface HybridSearchConfig {
  vectorWeight: number; // 0-1, weight for vector search
  keywordWeight: number; // 0-1, weight for keyword search
  maxResults: number;
  minScore: number;
  enableReranking?: boolean;
}

export interface KeywordSearchResult {
  chunkId: string;
  score: number;
  matches: string[];
}

export class HybridRetriever {
  private vectorStore: VectorStore;
  private embeddingModel: EmbeddingModel;
  private queryProcessor: QueryProcessor;
  private config: HybridSearchConfig;
  private documentIndex: Map<string, string> = new Map(); // chunkId -> content

  constructor(
    vectorStore: VectorStore,
    embeddingModel: EmbeddingModel,
    config: Partial<HybridSearchConfig> = {}
  ) {
    this.vectorStore = vectorStore;
    this.embeddingModel = embeddingModel;
    this.queryProcessor = new QueryProcessor();

    this.config = {
      vectorWeight: 0.7,
      keywordWeight: 0.3,
      maxResults: 10,
      minScore: 0.1,
      enableReranking: true,
      ...config,
    };

    // Ensure weights sum to 1
    const totalWeight = this.config.vectorWeight + this.config.keywordWeight;
    if (Math.abs(totalWeight - 1.0) > 0.01) {
      this.config.vectorWeight = this.config.vectorWeight / totalWeight;
      this.config.keywordWeight = this.config.keywordWeight / totalWeight;
    }
  }

  async search(query: string, k?: number): Promise<SearchResult[]> {
    const maxResults = k || this.config.maxResults;

    // Step 1: Process the query
    const processedQuery = await this.queryProcessor.processQuery(query);

    // Step 2: Perform vector search
    const vectorResults = await this.performVectorSearch(
      processedQuery,
      maxResults * 2
    );

    // Step 3: Perform keyword search
    const keywordResults = await this.performKeywordSearch(
      processedQuery,
      maxResults * 2
    );

    // Step 4: Combine and score results
    const combinedResults = this.combineResults(
      vectorResults,
      keywordResults,
      processedQuery
    );

    // Step 5: Re-rank if enabled
    let finalResults = combinedResults;
    if (this.config.enableReranking) {
      finalResults = await this.rerankResults(combinedResults, processedQuery);
    }

    // Step 6: Filter by minimum score and limit results
    return finalResults
      .filter((result) => result.score >= this.config.minScore)
      .slice(0, maxResults);
  }

  private async performVectorSearch(
    processedQuery: ProcessedQuery,
    k: number
  ): Promise<SearchResult[]> {
    try {
      // Use the processed query for better semantic matching
      const searchQuery = processedQuery.processed || processedQuery.original;
      return await this.vectorStore.search(searchQuery, k);
    } catch (error) {
      console.warn("Vector search failed:", error);
      return [];
    }
  }

  private async performKeywordSearch(
    processedQuery: ProcessedQuery,
    k: number
  ): Promise<KeywordSearchResult[]> {
    try {
      const keywords = processedQuery.keywords;
      if (keywords.length === 0) {
        return [];
      }

      const results: KeywordSearchResult[] = [];

      // Simple keyword matching - in production, you'd use a proper search engine
      for (const [chunkId, content] of this.documentIndex) {
        const matches: string[] = [];
        let score = 0;

        const lowerContent = content.toLowerCase();

        for (const keyword of keywords) {
          const regex = new RegExp(`\\b${keyword}\\b`, "gi");
          const matchCount = (content.match(regex) || []).length;

          if (matchCount > 0) {
            matches.push(keyword);
            // TF-like scoring
            score += matchCount * Math.log(content.length / (matchCount + 1));
          }
        }

        if (matches.length > 0) {
          // Normalize score by content length
          score = score / Math.sqrt(content.length);
          results.push({ chunkId, score, matches });
        }
      }

      return results.sort((a, b) => b.score - a.score).slice(0, k);
    } catch (error) {
      console.warn("Keyword search failed:", error);
      return [];
    }
  }

  private combineResults(
    vectorResults: SearchResult[],
    keywordResults: KeywordSearchResult[],
    processedQuery: ProcessedQuery
  ): SearchResult[] {
    const combinedMap = new Map<string, SearchResult>();

    // Add vector search results
    for (const result of vectorResults) {
      const key = result.chunk.id;
      const weightedScore = result.score * this.config.vectorWeight;

      combinedMap.set(key, {
        ...result,
        score: weightedScore,
      });
    }

    // Add or update with keyword search results
    for (const keywordResult of keywordResults) {
      const key = keywordResult.chunkId;
      const weightedScore = keywordResult.score * this.config.keywordWeight;

      if (combinedMap.has(key)) {
        // Combine scores
        const existing = combinedMap.get(key)!;
        existing.score += weightedScore;
      } else {
        // Find the corresponding vector result or create a minimal one
        const vectorResult = vectorResults.find((vr) => vr.chunk.id === key);
        if (vectorResult) {
          combinedMap.set(key, {
            ...vectorResult,
            score: weightedScore,
          });
        }
      }
    }

    return Array.from(combinedMap.values()).sort((a, b) => b.score - a.score);
  }

  private async rerankResults(
    results: SearchResult[],
    processedQuery: ProcessedQuery
  ): Promise<SearchResult[]> {
    // Simple re-ranking based on multiple factors
    return results
      .map((result) => {
        let bonusScore = 0;

        // Bonus for exact phrase matches
        if (
          result.chunk.content
            .toLowerCase()
            .includes(processedQuery.original.toLowerCase())
        ) {
          bonusScore += 0.1;
        }

        // Bonus for keyword density
        const keywordCount = processedQuery.keywords.reduce(
          (count, keyword) => {
            const regex = new RegExp(`\\b${keyword}\\b`, "gi");
            return count + (result.chunk.content.match(regex) || []).length;
          },
          0
        );

        const keywordDensity =
          keywordCount / result.chunk.content.split(" ").length;
        bonusScore += Math.min(keywordDensity * 0.5, 0.2);

        // Bonus for recent documents (if timestamp available)
        if (result.document.metadata.updatedAt) {
          try {
            // Ensure updatedAt is a Date object
            const updatedAt =
              result.document.metadata.updatedAt instanceof Date
                ? result.document.metadata.updatedAt
                : new Date(result.document.metadata.updatedAt);

            const daysSinceUpdate =
              (Date.now() - updatedAt.getTime()) / (1000 * 60 * 60 * 24);
            if (daysSinceUpdate < 30) {
              bonusScore += 0.05;
            }
          } catch (dateError) {
            console.warn(
              "Failed to process updatedAt for bonus scoring:",
              dateError
            );
          }
        }

        // Bonus for language match
        if (processedQuery.language === result.document.metadata.language) {
          bonusScore += 0.05;
        }

        return {
          ...result,
          score: Math.min(result.score + bonusScore, 1.0),
        };
      })
      .sort((a, b) => b.score - a.score);
  }

  // Method to update the document index for keyword search
  updateDocumentIndex(chunkId: string, content: string): void {
    this.documentIndex.set(chunkId, content);
  }

  // Method to remove from document index
  removeFromIndex(chunkId: string): void {
    this.documentIndex.delete(chunkId);
  }

  // Method to update search configuration
  updateConfig(config: Partial<HybridSearchConfig>): void {
    this.config = { ...this.config, ...config };

    // Ensure weights sum to 1
    const totalWeight = this.config.vectorWeight + this.config.keywordWeight;
    if (Math.abs(totalWeight - 1.0) > 0.01) {
      this.config.vectorWeight = this.config.vectorWeight / totalWeight;
      this.config.keywordWeight = this.config.keywordWeight / totalWeight;
    }
  }

  // Method to get search statistics
  getSearchStats(): {
    indexSize: number;
    config: HybridSearchConfig;
  } {
    return {
      indexSize: this.documentIndex.size,
      config: this.config,
    };
  }

  // Method to clear the index
  clearIndex(): void {
    this.documentIndex.clear();
  }
}