# 第14課：RAG進階 (Part 4)

# query_expander.py - LLMベースのマルチクエリ拡張

import logging
import time
from typing import List, Dict, Any, Optional
from openai import OpenAI
from sentence_transformers import SentenceTransformer

# ロギング設定
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class MultiQueryExpander:
    """LLMを使用したマルチクエリ拡張クラス"""
    
    def __init__(
        self,
        llm_api_key: Optional[str] = None,
        llm_model: str = "gpt-3.5-turbo",
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        num_queries: int = 4,
        similarity_threshold: float = 0.8,
        use_embeddings_filter: bool = True,
        device: str = "cpu"
    ):
        """マルチクエリ拡張器の初期化
        
        Args:
            llm_api_key: LLM API キー
            llm_model: 使用するLLMモデル
            embedding_model: 埋め込みモデル名
            num_queries: 生成するクエリ数
            similarity_threshold: 類似性フィルタリングのしきい値
            use_embeddings_filter: 埋め込みでフィルタリングするかどうか
            device: 使用するデバイス
        """
        self.llm_model = llm_model
        self.num_queries = num_queries
        self.similarity_threshold = similarity_threshold
        self.use_embeddings_filter = use_embeddings_filter
        
        # OpenAIクライアントの初期化
        self.client = OpenAI(api_key=llm_api_key)
        
        # 埋め込みモデルの初期化（フィルタリングに使用）
        if use_embeddings_filter:
            logger.info(f"埋め込みモデルを初期化: {embedding_model}")
            self.embedding_model = SentenceTransformer(embedding_model, device=device)
    
    def expand_query(
        self,
        query: str,
        domain: Optional[str] = None,
        query_type: Optional[str] = None
    ) -> List[str]:
        """元のクエリから複数の拡張クエリを生成
        
        Args:
            query: 元のクエリ
            domain: クエリのドメイン（オプション）
            query_type: クエリのタイプ（オプション）
            
        Returns:
            拡張クエリのリスト
        """
        logger.info(f"クエリの拡張中: '{query}'")
        
        # プロンプトの準備
        system_prompt = "あなたは高度な検索システムの一部として機能する検索クエリ拡張スペシャリストです。"
        system_prompt += "与えられたクエリから、類似した意味を持つが表現の異なる代替クエリを生成してください。"
        system_prompt += "拡張クエリは多様な視点をカバーし、元のクエリと同等の意図を持つものでなければなりません。"
        
        if domain:
            system_prompt += f"\n特定のドメイン '{domain}' に関連するクエリを生成してください。"
        
        if query_type:
            system_prompt += f"\nこのクエリは '{query_type}' タイプのクエリです。この特性を考慮してください。"
        
        user_prompt = f"以下のクエリを拡張して、同じ意図を持つ{self.num_queries}つの代替クエリを生成してください。"
        user_prompt += f"\n\n元のクエリ: {query}\n\n"
        user_prompt += "クエリはそれぞれ異なる表現を使用し、元の意図を維持しながら検索範囲を広げるものであるべきです。"
        user_prompt += "各クエリは別々の行に、クエリのみをプレーンテキストで出力してください。番号や追加の説明は不要です。"
        
        # LLMを使ってクエリを生成
        try:
            start_time = time.time()
            
            response = self.client.chat.completions.create(
                model=self.llm_model,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.7,
                max_tokens=150 * self.num_queries,
                n=1
            )
            
            # レスポンスから拡張クエリを抽出
            content = response.choices[0].message.content
            expanded_queries = [line.strip() for line in content.split('\n') if line.strip()]
            
            # 元のクエリを追加
            all_queries = [query] + expanded_queries
            
            # 重複を削除
            unique_queries = list(dict.fromkeys(all_queries))
            
            elapsed_time = time.time() - start_time
            logger.info(f"クエリ拡張完了: {len(expanded_queries)}個のクエリを生成 ({elapsed_time:.2f}秒)")
            
            # 埋め込みベースのフィルタリング（有効な場合）
            if self.use_embeddings_filter and len(unique_queries) > 2:
                filtered_queries = self._filter_similar_queries(unique_queries)
                logger.info(f"フィルタリング後: {len(filtered_queries)}個のクエリ")
                return filtered_queries
            
            return unique_queries
            
        except Exception as e:
            logger.error(f"クエリ拡張エラー: {str(e)}")
            # エラーの場合は元のクエリのみを返す
            return [query]
    
    def _filter_similar_queries(self, queries: List[str]) -> List[str]:
        """埋め込みベースの類似度でクエリをフィルタリング
        
        Args:
            queries: フィルタリングするクエリリスト
            
        Returns:
            フィルタリングされたクエリリスト
        """
        # クエリが少なすぎる場合はそのまま返す
        if len(queries) <= 2:
            return queries
        
        # クエリの埋め込みを計算
        embeddings = self.embedding_model.encode(queries)
        
        # 元のクエリを常に保持
        filtered_queries = [queries[0]]
        original_embedding = embeddings[0]
        
        # 残りのクエリを処理
        for i, query in enumerate(queries[1:], 1):
            # 元のクエリとの類似度を計算
            similarity = self._cosine_similarity(embeddings[i], original_embedding)
            
            # 類似度が閾値を超えているがあまりに似ていない場合は追加
            if similarity < 0.95 and similarity > 0.5:
                # すでに選択されたクエリとの類似度もチェック
                too_similar = False
                for selected_idx, selected_query in enumerate(filtered_queries[1:], 1):
                    selected_embedding = self.embedding_model.encode([selected_query])[0]
                    if self._cosine_similarity(embeddings[i], selected_embedding) > self.similarity_threshold:
                        too_similar = True
                        break
                
                if not too_similar:
                    filtered_queries.append(query)
            
            # 最大クエリ数に達したら終了
            if len(filtered_queries) >= self.num_queries:
                break
        
        return filtered_queries
    
    def _cosine_similarity(self, a, b):
        """コサイン類似度を計算"""
        return float(a @ b) / (float(sum(a*a) ** 0.5) * float(sum(b*b) ** 0.5))