src/query_expansion.py

"""
Expansão de Queries (Multi-Query Retrieval)

Gera múltiplas variações de uma query para melhorar cobertura da busca.
"""
import re
from typing import List, Dict, Any, Optional
from src.generation import GenerationManager


class QueryExpander:
    """Expande queries em múltiplas variações para melhor retrieval"""

    def __init__(self, generation_manager: GenerationManager):
        """
        Args:
            generation_manager: Gerenciador de geração de texto
        """
        self.generation_manager = generation_manager

    def expand_query(
        self,
        query: str,
        num_variations: int = 3,
        method: str = "llm"
    ) -> List[str]:
        """
        Expande query em múltiplas variações

        Args:
            query: Query original
            num_variations: Número de variações a gerar
            method: Método de expansão ("llm", "template", "paraphrase")

        Returns:
            Lista com query original + variações
        """
        if method == "llm":
            return self._expand_with_llm(query, num_variations)
        elif method == "template":
            return self._expand_with_templates(query, num_variations)
        elif method == "paraphrase":
            return self._expand_with_paraphrase(query, num_variations)
        else:
            return [query]

    def _expand_with_llm(self, query: str, num_variations: int) -> List[str]:
        """
        Usa LLM para gerar variações da query

        Estratégia: Pede ao LLM para reformular a pergunta de formas diferentes
        """
        prompt = f"""Você é um assistente que ajuda a reformular perguntas para melhorar buscas.

Pergunta original: "{query}"

Gere {num_variations} reformulações diferentes desta pergunta. Cada reformulação deve:
- Manter o mesmo significado e intenção
- Usar palavras e estruturas diferentes
- Ser igualmente específica

Formato de saída (uma por linha):
1. [primeira reformulação]
2. [segunda reformulação]
3. [terceira reformulação]

Reformulações:"""

        try:
            response = self.generation_manager.generate(
                prompt=prompt,
                max_tokens=200,
                temperature=0.7
            )

            # Extrai variações do response
            variations = self._parse_llm_variations(response)

            # Garante que temos pelo menos a query original
            if not variations:
                variations = [query]
            elif query not in variations:
                variations.insert(0, query)

            return variations[:num_variations + 1]  # +1 para incluir original

        except Exception as e:
            print(f"Erro ao expandir query com LLM: {e}")
            return [query]

    def _parse_llm_variations(self, response: str) -> List[str]:
        """
        Extrai variações do response do LLM

        Procura por linhas numeradas ou bullets
        """
        variations = []

        # Tenta extrair linhas numeradas: "1. texto", "2. texto"
        pattern = r'^\d+\.\s*(.+)$'
        for line in response.split('\n'):
            line = line.strip()
            match = re.match(pattern, line)
            if match:
                variation = match.group(1).strip()
                if variation:
                    variations.append(variation)

        # Se não encontrou numeradas, tenta bullets: "- texto", "* texto"
        if not variations:
            pattern = r'^[-*]\s*(.+)$'
            for line in response.split('\n'):
                line = line.strip()
                match = re.match(pattern, line)
                if match:
                    variation = match.group(1).strip()
                    if variation:
                        variations.append(variation)

        return variations

    def _expand_with_templates(self, query: str, num_variations: int) -> List[str]:
        """
        Usa templates fixos para expandir query

        Útil quando LLM não está disponível ou para casos simples
        """
        templates = [
            query,  # Original
            f"Explique sobre {query}",
            f"O que é {query}?",
            f"Como funciona {query}?",
            f"Qual a definição de {query}?",
            f"Informações sobre {query}",
        ]

        return templates[:num_variations + 1]

    def _expand_with_paraphrase(self, query: str, num_variations: int) -> List[str]:
        """
        Usa paraphrasing simples baseado em sinônimos

        Nota: Implementação básica. Para produção, considere usar
        modelo de paraphrase como T5 ou BART
        """
        # Implementação simplificada com algumas variações comuns
        variations = [query]

        # Substituições comuns em português
        substitutions = [
            ("o que é", "qual é"),
            ("como funciona", "qual o funcionamento de"),
            ("explique", "descreva"),
            ("diferença entre", "distinção entre"),
            ("vantagens", "benefícios"),
        ]

        for old, new in substitutions:
            if old in query.lower():
                variation = query.lower().replace(old, new).capitalize()
                if variation not in variations:
                    variations.append(variation)
                    if len(variations) > num_variations:
                        break

        return variations[:num_variations + 1]

    def get_expansion_info(self, method: str) -> Dict[str, Any]:
        """
        Retorna informações sobre método de expansão

        Args:
            method: Nome do método

        Returns:
            Dicionário com informações
        """
        info = {
            "llm": {
                "name": "LLM-based",
                "description": "Usa modelo de linguagem para gerar variações contextuais",
                "pros": "Variações de alta qualidade, contextuais",
                "cons": "Mais lento, requer LLM disponível",
                "best_for": "Queries complexas e conceituais"
            },
            "template": {
                "name": "Template-based",
                "description": "Usa templates fixos para reformular queries",
                "pros": "Rápido, determinístico, sem dependências",
                "cons": "Variações genéricas, pode não preservar nuances",
                "best_for": "Queries simples, prototipação rápida"
            },
            "paraphrase": {
                "name": "Paraphrase-based",
                "description": "Usa substituições de sinônimos e paráfrases",
                "pros": "Balanceado, mantém estrutura original",
                "cons": "Limitado por dicionário de sinônimos",
                "best_for": "Queries médias, quando LLM não está disponível"
            }
        }

        return info.get(method, {
            "name": method,
            "description": "Método desconhecido",
            "pros": "N/A",
            "cons": "N/A",
            "best_for": "N/A"
        })