query_transform.py

"""Query transformation techniques for improved retrieval.

This module implements advanced query transformation techniques that significantly
improve retrieval quality by bridging the query-document vocabulary gap:

1. HyDE (Hypothetical Document Embeddings): Generates a hypothetical answer to the
   query and embeds that instead of the raw query. Since hypothetical answers are
   more similar to actual documents, this improves retrieval accuracy.

2. Multi-Query Expansion: Generates multiple variations of the query from different
   perspectives and combines results using reciprocal rank fusion.

3. Combined Approach: Generates multiple query variations AND creates hypothetical
   documents for each, combining the benefits of both techniques.

4. Bilingual Expansion: Translates key concepts/terms in the query to additional
   languages (e.g., Greek) to find content written in those languages. Lightweight
   translation of nouns and concepts rather than full query translation.

5. LRU Caching: Query expansion results are cached to avoid redundant LLM calls
   for repeated or similar queries, reducing latency and API costs by 30-50%.
"""

import logging
import threading
from collections import OrderedDict
from typing import List, Union, Optional, Tuple
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.embeddings import BaseEmbedding

logger = logging.getLogger(__name__)


class QueryExpansionCache:
    """LRU cache for query expansion results to avoid redundant LLM calls.

    Caches results from HyDE, multi-query, and bilingual expansion to reduce
    latency and API costs for repeated or similar queries.
    """

    def __init__(self, max_size: int = 100):
        """Initialize cache with maximum size.

        Args:
            max_size: Maximum number of entries to cache (default: 100)
        """
        self._cache: OrderedDict[str, Union[str, List[str]]] = OrderedDict()
        self._max_size = max_size
        self._hits = 0
        self._misses = 0

    @staticmethod
    def normalize_query(query: str) -> str:
        """Normalize query for consistent cache keys.

        Args:
            query: Raw query string

        Returns:
            Normalized query (lowercase, collapsed whitespace, trimmed punctuation)
        """
        # Lowercase and collapse whitespace
        normalized = ' '.join(query.lower().split())
        # Remove trailing punctuation
        normalized = normalized.rstrip('?!.')
        return normalized

    def _make_key(self, method: str, query: str, *args) -> str:
        """Create cache key from method, normalized query, and extra args.

        Args:
            method: Expansion method name (hyde, multi_query, bilingual)
            query: User query
            *args: Additional parameters (e.g., num_queries, languages)

        Returns:
            Cache key string
        """
        normalized = self.normalize_query(query)
        args_str = ':'.join(str(a) for a in args) if args else ''
        return f"{method}:{normalized}:{args_str}"

    def get(
        self,
        method: str,
        query: str,
        *args
    ) -> Optional[Union[str, List[str]]]:
        """Get cached expansion result.

        Args:
            method: Expansion method name
            query: User query
            *args: Additional parameters

        Returns:
            Cached result or None if not found
        """
        key = self._make_key(method, query, *args)
        if key in self._cache:
            # Move to end (most recently used)
            self._cache.move_to_end(key)
            self._hits += 1
            logger.debug(f"Cache HIT for {method}: {query[:50]}...")
            return self._cache[key]
        self._misses += 1
        logger.debug(f"Cache MISS for {method}: {query[:50]}...")
        return None

    def set(
        self,
        method: str,
        query: str,
        result: Union[str, List[str]],
        *args
    ) -> None:
        """Store expansion result in cache.

        Args:
            method: Expansion method name
            query: User query
            result: Expansion result to cache
            *args: Additional parameters
        """
        key = self._make_key(method, query, *args)
        if key in self._cache:
            self._cache.move_to_end(key)
        else:
            if len(self._cache) >= self._max_size:
                # Remove oldest entry
                evicted_key, _ = self._cache.popitem(last=False)
                logger.debug(f"Cache evicted oldest entry: {evicted_key[:50]}...")
        self._cache[key] = result

    def stats(self) -> dict:
        """Get cache statistics.

        Returns:
            Dict with hits, misses, size, and hit rate
        """
        total = self._hits + self._misses
        hit_rate = (self._hits / total * 100) if total > 0 else 0.0
        return {
            "hits": self._hits,
            "misses": self._misses,
            "size": len(self._cache),
            "max_size": self._max_size,
            "hit_rate": f"{hit_rate:.1f}%"
        }

    def clear(self) -> None:
        """Clear all cached entries."""
        self._cache.clear()
        logger.info("Query expansion cache cleared")


# Global cache instance (shared across QueryTransformer instances)
_expansion_cache = QueryExpansionCache(max_size=100)
_inflight_hyde_lock = threading.Lock()
_inflight_hyde_events: dict[str, threading.Event] = {}


def get_expansion_cache() -> QueryExpansionCache:
    """Get the global query expansion cache instance."""
    return _expansion_cache


class QueryTransformer:
    """Query transformer implementing HyDE and Multi-Query expansion techniques.

    Includes LRU caching to avoid redundant LLM calls for repeated queries.
    """

    def __init__(
        self,
        llm: GoogleGenAI,
        embed_model: Optional[BaseEmbedding] = None,
        enable_cache: bool = True,
        hyde_temperature: float = 0.7
    ):
        """
        Initialize query transformer.

        Args:
            llm: Language model for generating transformations
            embed_model: Optional embedding model (for future use)
            enable_cache: Whether to use LRU cache for expansion results (default: True)
            hyde_temperature: Temperature for HyDE hypothetical document generation (default: 0.7)
        """
        self.llm = llm
        self.embed_model = embed_model
        self.enable_cache = enable_cache
        self._cache = _expansion_cache

        # Create separate HyDE LLM with higher temperature for diverse hypothetical docs
        self.hyde_llm = self._create_hyde_llm(hyde_temperature)
        logger.info(f"QueryTransformer initialized (cache={'enabled' if enable_cache else 'disabled'}, hyde_temp={hyde_temperature})")

    def _create_hyde_llm(self, temperature: float):
        """Create a separate LLM instance for HyDE with higher temperature."""
        try:
            import os
            hyde_llm = GoogleGenAI(
                model=self.llm.model if hasattr(self.llm, 'model') else "gemini-3-flash-preview",
                api_key=os.getenv("GOOGLE_API_KEY", ""),
                temperature=temperature,
                max_tokens=2048,
            )
            logger.info(f"Created HyDE LLM with temperature={temperature}")
            return hyde_llm
        except Exception as e:
            logger.warning(f"Failed to create HyDE LLM (temp={temperature}), using main LLM: {e}")
            return self.llm

    def hyde_transform(self, query: str) -> str:
        """
        Generate hypothetical document for query using HyDE technique.

        HyDE works by generating a hypothetical answer that resembles the structure
        and vocabulary of actual documents. This bridges the query-document gap since
        documents are more similar to other documents than to queries.

        Results are cached to avoid redundant LLM calls for repeated queries.

        Args:
            query: Original user query

        Returns:
            Hypothetical document text
        """
        logger.debug(f"Applying HyDE transformation to query: {query[:100]}...")

        # Check cache first
        if self.enable_cache:
            cached = self._cache.get("hyde", query)
            if cached is not None:
                logger.info(f"HyDE cache hit - skipping LLM call")
                return cached

        # Prevent duplicate concurrent HyDE generations for the same query.
        # In federated retrieval, multiple sources may call HyDE in parallel.
        inflight_key = QueryExpansionCache.normalize_query(query)
        is_owner = False
        with _inflight_hyde_lock:
            event = _inflight_hyde_events.get(inflight_key)
            if event is None:
                event = threading.Event()
                _inflight_hyde_events[inflight_key] = event
                is_owner = True

        if not is_owner:
            logger.debug("HyDE generation in-flight for query; waiting for cached result")
            event.wait(timeout=30)
            if self.enable_cache:
                cached = self._cache.get("hyde", query)
                if cached is not None:
                    logger.info("HyDE cache hit after in-flight wait - reusing transformed query")
                    return cached

        # Craft a prompt that generates a document-like response
        # Using thinking mode for better quality hypothetical documents
        prompt = f"""You are helping to search a personal knowledge base (Obsidian vault).

For the following question, write a detailed, informative passage that would appear in someone's notes and would answer this question. Write as if you're composing a note entry, not as if you're answering the question directly.

Question: {query}

Write a comprehensive note passage (2-3 paragraphs) that would contain the answer to this question. Use the style and vocabulary typical of knowledge base notes:"""

        try:
            # Generate hypothetical document (using higher-temperature HyDE LLM)
            response = self.hyde_llm.complete(prompt)
            hypothetical_doc = response.text.strip()

            logger.debug(f"Generated hypothetical document ({len(hypothetical_doc)} chars)")
            logger.debug(f"Hypothetical doc preview: {hypothetical_doc[:200]}...")

            # Cache the result
            if self.enable_cache:
                self._cache.set("hyde", query, hypothetical_doc)

            return hypothetical_doc

        except Exception as e:
            logger.warning(f"HyDE transformation failed: {e}. Falling back to original query.")
            return query
        finally:
            if is_owner:
                with _inflight_hyde_lock:
                    done_event = _inflight_hyde_events.pop(inflight_key, None)
                    if done_event:
                        done_event.set()

    def multi_query_expand(self, query: str, num_queries: int = 3) -> List[str]:
        """
        Generate multiple query variations from different perspectives.

        Multi-query expansion generates semantically similar queries that approach
        the information need from different angles. This increases recall by
        capturing documents that might match one variation better than others.

        Results are cached to avoid redundant LLM calls for repeated queries.

        Args:
            query: Original user query
            num_queries: Number of variations to generate (default: 3)

        Returns:
            List of query variations including the original
        """
        logger.debug(f"Generating {num_queries} query variations for: {query[:100]}...")

        # Check cache first
        if self.enable_cache:
            cached = self._cache.get("multi_query", query, num_queries)
            if cached is not None:
                logger.info(f"Multi-query cache hit - skipping LLM call")
                return cached

        # Craft a prompt that generates diverse query reformulations
        prompt = f"""You are helping to improve search in a personal knowledge base.

Generate {num_queries} different versions of the following search query. Each version should:
- Approach the topic from a slightly different angle
- Use alternative vocabulary and phrasing
- Maintain the same core information need
- Be concise (1-2 sentences each)

Original query: {query}

Generate {num_queries} alternative search queries, one per line, without numbering:"""

        try:
            # Generate query variations
            response = self.llm.complete(prompt)
            variations_text = response.text.strip()

            # Parse the variations (split by newlines and clean)
            variations = [
                line.strip()
                for line in variations_text.split('\n')
                if line.strip() and not line.strip().startswith(('#', '-', '*'))
            ]

            # Remove any numbering or bullets that might be present
            cleaned_variations = []
            for var in variations:
                # Remove common prefixes like "1.", "1)", "-", "*", etc.
                cleaned = var
                if var and len(var) > 2:
                    # Remove leading numbers and punctuation
                    for prefix in ['1.', '2.', '3.', '4.', '5.', '1)', '2)', '3)', '4)', '5)', '- ', '* ']:
                        if cleaned.startswith(prefix):
                            cleaned = cleaned[len(prefix):].strip()
                cleaned_variations.append(cleaned)

            # Take only the requested number of variations
            variations = cleaned_variations[:num_queries]

            # Always include the original query first
            all_queries = [query] + variations

            logger.info(f"Generated {len(variations)} query variations (+ original)")
            for i, q in enumerate(all_queries):
                logger.debug(f"  Query {i+1}: {q[:100]}...")

            # Cache the result
            if self.enable_cache:
                self._cache.set("multi_query", query, all_queries, num_queries)

            return all_queries

        except Exception as e:
            logger.warning(f"Multi-query expansion failed: {e}. Using original query only.")
            return [query]

    # Language name mapping for prompts
    LANGUAGE_NAMES = {
        "el": "Greek (Ελληνικά)",
        "es": "Spanish (Español)",
        "de": "German (Deutsch)",
        "fr": "French (Français)",
        "it": "Italian (Italiano)",
        "pt": "Portuguese (Português)",
        "nl": "Dutch (Nederlands)",
        "ru": "Russian (Русский)",
        "zh": "Chinese (中文)",
        "ja": "Japanese (日本語)",
        "ko": "Korean (한국어)",
        "ar": "Arabic (العربية)"
    }

    def bilingual_expand(
        self,
        query: str,
        target_languages: List[str] = None
    ) -> List[str]:
        """
        Expand query with translations of key terms to target languages.

        This is a lightweight expansion that translates key nouns and concepts
        rather than the full query, enabling retrieval of documents written
        in other languages that discuss the same topics.

        Results are cached to avoid redundant LLM calls for repeated queries.

        Args:
            query: Original user query (in English or any language)
            target_languages: List of language codes (e.g., ["el", "es"]).
                            Defaults to ["el"] (Greek).

        Returns:
            List of queries including original + translated versions
        """
        if target_languages is None:
            target_languages = ["el"]

        logger.debug(f"Bilingual expansion for query: {query[:100]}... to {target_languages}")

        # Check cache first (convert list to tuple for hashability)
        langs_key = tuple(sorted(target_languages))
        if self.enable_cache:
            cached = self._cache.get("bilingual", query, langs_key)
            if cached is not None:
                logger.info(f"Bilingual cache hit - skipping LLM call")
                return cached

        # Build language names for the prompt
        lang_names = [self.LANGUAGE_NAMES.get(lang, lang) for lang in target_languages]
        languages_str = ", ".join(lang_names)

        prompt = f"""You are helping to search a bilingual knowledge base containing notes in English and other languages.

For the following query, translate the KEY CONCEPTS and NOUNS to the target language(s). This is NOT a full translation - focus on:
- Main nouns and topics (e.g., "habits" → "συνήθειες")
- Key concepts (e.g., "productivity" → "παραγωγικότητα")
- Important terms the user might have written notes about in the target language

Original query: {query}

Target language(s): {languages_str}

For EACH target language, provide a search query that includes the translated key terms. The translated query should be useful for finding documents written in that language about the same topic.

Output one translated query per line, without numbering or prefixes. If the original query already contains terms in a target language, still output a version with any additional English terms translated:"""

        try:
            response = self.llm.complete(prompt)
            translations_text = response.text.strip()

            # Parse translations
            translations = [
                line.strip()
                for line in translations_text.split('\n')
                if line.strip() and not line.strip().startswith(('#', '-', '*'))
            ]

            # Clean up any numbering or bullets
            cleaned_translations = []
            for trans in translations:
                cleaned = trans
                if len(trans) > 2:
                    for prefix in ['1.', '2.', '3.', '1)', '2)', '3)', '- ', '* ', '• ']:
                        if cleaned.startswith(prefix):
                            cleaned = cleaned[len(prefix):].strip()
                if cleaned:
                    cleaned_translations.append(cleaned)

            # Take one translation per target language (limit to len(target_languages))
            translations = cleaned_translations[:len(target_languages)]

            # Always include original query first
            all_queries = [query] + translations

            logger.info(f"Bilingual expansion: {len(translations)} translated queries generated")
            # Log full queries for debugging
            for i, q in enumerate(all_queries):
                label = "Original" if i == 0 else f"Translated {i}"
                logger.info(f"  [{label}]: {q}")

            # Cache the result
            if self.enable_cache:
                self._cache.set("bilingual", query, all_queries, langs_key)

            return all_queries

        except Exception as e:
            logger.warning(f"Bilingual expansion failed: {e}. Using original query only.")
            return [query]

    def transform_query(
        self,
        query: str,
        method: str = "hyde",
        num_queries: int = 3
    ) -> Union[str, List[str]]:
        """
        Apply selected query transformation method.

        Args:
            query: Original user query
            method: Transformation method - "hyde", "multi_query", or "both"
            num_queries: Number of query variations for multi_query (default: 3)

        Returns:
            Transformed query (string for hyde, list for multi_query/both)
        """
        logger.info(f"Transforming query with method: {method}")

        if method == "hyde":
            return self.hyde_transform(query)

        elif method == "multi_query":
            return self.multi_query_expand(query, num_queries)

        elif method == "both":
            # Generate query variations first
            variations = self.multi_query_expand(query, num_queries)

            # Generate hypothetical document for each variation
            logger.debug("Applying HyDE to each query variation...")
            hypothetical_docs = []
            for i, variation in enumerate(variations):
                logger.debug(f"HyDE transform {i+1}/{len(variations)}")
                hyde_doc = self.hyde_transform(variation)
                hypothetical_docs.append(hyde_doc)

            logger.info(f"Generated {len(hypothetical_docs)} hypothetical documents from {len(variations)} queries")
            return hypothetical_docs

        elif method == "none" or method == "disabled":
            # No transformation
            logger.debug("Query transformation disabled, returning original query")
            return query

        else:
            logger.warning(f"Unknown transformation method: {method}. Using original query.")
            return query

    def get_embedding_queries(
        self,
        query: str,
        method: str = "hyde",
        num_queries: int = 3
    ) -> List[str]:
        """
        Get list of queries to embed for retrieval.

        This is a convenience method that ensures the output is always a list,
        making it easier to integrate with retrievers.

        Args:
            query: Original user query
            method: Transformation method
            num_queries: Number of variations for multi_query

        Returns:
            List of queries/documents to embed
        """
        result = self.transform_query(query, method, num_queries)

        # Ensure we always return a list
        if isinstance(result, str):
            return [result]
        elif isinstance(result, list):
            return result
        else:
            logger.warning(f"Unexpected transform result type: {type(result)}")
            return [query]

    def decompose_query(self, query: str) -> list[str]:
        """Decompose a complex multi-part query into atomic sub-queries.

        Detects queries that ask about multiple topics (e.g., "Compare X and Y")
        and splits them into independent retrieval queries. Returns [original]
        if the query is already atomic.

        Args:
            query: User query to potentially decompose

        Returns:
            List of sub-queries (length 1 if atomic)
        """
        # Check cache
        if self.enable_cache:
            cached = self._cache.get("decompose", query)
            if cached is not None:
                logger.info("Decomposition cache hit")
                return cached

        prompt = f"""Analyze whether this search query should be split into separate sub-queries for better retrieval.

A query should be decomposed when it asks about MULTIPLE distinct topics that would be found in different documents.
Examples:
- "Compare HyDE and multi-query retrieval" -> ["How does HyDE retrieval work?", "How does multi-query retrieval work?"]
- "What are the benefits and risks of intermittent fasting?" -> ["What are the benefits of intermittent fasting?", "What are the risks of intermittent fasting?"]
- "How does sleep affect memory?" -> NO DECOMPOSITION (single topic)

Query: {query}

If this query should be decomposed, output each sub-query on a separate line (2-4 sub-queries).
If the query is already focused on a single topic, output just: ATOMIC

Output:"""

        try:
            response = self.llm.complete(prompt)
            text = response.text.strip()

            if "ATOMIC" in text.upper() or not text:
                result = [query]
            else:
                sub_queries = [
                    line.strip()
                    for line in text.split('\n')
                    if line.strip() and not line.strip().upper().startswith('ATOMIC')
                ]
                # Clean numbering
                cleaned = []
                for sq in sub_queries:
                    for prefix in ['1.', '2.', '3.', '4.', '1)', '2)', '3)', '4)', '- ', '* ']:
                        if sq.startswith(prefix):
                            sq = sq[len(prefix):].strip()
                    if sq:
                        cleaned.append(sq)
                result = cleaned[:4] if len(cleaned) >= 2 else [query]

            logger.info(f"Query decomposition: {len(result)} sub-queries")
            if self.enable_cache:
                self._cache.set("decompose", query, result)
            return result

        except Exception as e:
            logger.warning(f"Query decomposition failed: {e}")
            return [query]

    def cache_stats(self) -> dict:
        """Get query expansion cache statistics.

        Returns:
            Dict with hits, misses, size, max_size, and hit_rate
        """
        return self._cache.stats()

    def clear_cache(self) -> None:
        """Clear the query expansion cache."""
        self._cache.clear()