jhonny-editor/code/engine/persistencia/busca_de_conteudo.py

"""
Busca de conteúdo falado no banco de análises.

Combina duas listas independentes sobre o mesmo material: a lexical, que acha
a palavra exata via FTS5, e a semântica, que acha a ideia via embeddings dos
enunciados. Nenhuma das duas basta sozinha — a lexical não encontra "preço"
quando a pessoa disse "quanto custa", e a semântica erra nomes próprios e
termos técnicos que precisam bater literalmente.

Toda resposta carrega ``clipe_id``, ``inicio`` e ``fim``. Um acerto sem
timecode não serve para cortar, e por isso não é considerado resultado válido
aqui.
"""

from __future__ import annotations

import math
import sqlite3
from dataclasses import dataclass

from ..integracoes.embeddings import ErroDeEmbedding, ProviderDeEmbeddings, desempacotar

# Constante da fusão recíproca de ranking (RRF). Amortece a diferença entre as
# escalas incomparáveis do FTS5 (bm25, menor é melhor) e do cosseno (maior é
# melhor): o que entra na conta é a posição em cada lista, não a nota.
CONSTANTE_DE_FUSAO = 60.0


class ErroDeBusca(ValueError):
    """Consulta vazia ou parâmetros de busca inválidos."""


@dataclass(frozen=True)
class ResultadoDeBusca:
    """
    Um trecho encontrado, sempre localizável na timeline.

    Atributos:
        video_id: Vídeo onde o trecho está.
        clipe_id: Clipe onde o trecho está.
        inicio: Início do trecho em segundos.
        fim: Fim do trecho em segundos.
        texto: Texto do trecho encontrado.
        falante: Falante do trecho, quando conhecido.
        pontuacao: Nota da fusão; maior é mais relevante.
        origem: Como o trecho foi encontrado — "lexical", "semantica" ou
            "ambas", para que o resultado seja explicável.
        falas: Ids das falas que compõem o trecho, para descer ao corte.
    """

    video_id: str
    clipe_id: str
    inicio: float
    fim: float
    texto: str
    falante: str | None
    pontuacao: float
    origem: str
    falas: tuple[int, ...]


class BuscaDeConteudo:
    """
    Busca falas e enunciados por texto exato, por sentido, ou pelos dois.

    A busca semântica é opcional: sem um provider de embeddings a classe
    continua servindo a busca lexical, em vez de falhar. Isso mantém o painel
    utilizável quando o Ollama não está rodando.

    Atributos:
        conexao: Conexão SQLite já aberta e com o esquema aplicado.
        provider: Provider de embeddings, ou None para busca só lexical.
    """

    def __init__(
        self,
        conexao: sqlite3.Connection,
        provider: ProviderDeEmbeddings | None = None,
    ) -> None:
        """
        Inicializa a busca sobre uma conexão e um provider opcional.

        Parâmetros:
            conexao: Conexão SQLite já aberta e com o esquema aplicado.
            provider: Provider de embeddings para a busca semântica. Quando
                ``None``, apenas a busca lexical fica disponível.
        """
        self.conexao = conexao
        self.provider = provider

    def buscar(
        self, consulta: str, video_id: str | None = None, limite: int = 10,
    ) -> list[ResultadoDeBusca]:
        """
        Busca um texto combinando as vias lexical e semântica.

        Parâmetros:
            consulta: Texto livre a procurar.
            video_id: Restringe a busca a um vídeo. Quando ``None``, procura
                em todo o acervo.
            limite: Máximo de resultados devolvidos.

        Retorna:
            Os trechos mais relevantes, do melhor para o pior.

        Pode gerar:
            ErroDeBusca: quando a consulta é vazia ou o limite não é positivo.
        """
        self._validar(consulta, limite)
        lexicais = self.buscar_lexical(consulta, video_id, limite * 2)
        semanticos = self.buscar_semantica(consulta, video_id, limite * 2)
        return self._fundir(lexicais, semanticos, limite)

    def buscar_lexical(
        self, consulta: str, video_id: str | None = None, limite: int = 10,
    ) -> list[ResultadoDeBusca]:
        """
        Busca a palavra exata nas falas, via índice FTS5.

        Parâmetros:
            consulta: Texto a procurar. Acentos são ignorados na comparação.
            video_id: Restringe a busca a um vídeo.
            limite: Máximo de resultados devolvidos.

        Retorna:
            As falas que contêm os termos, da mais relevante para a menos.

        Pode gerar:
            ErroDeBusca: quando a consulta é vazia ou o limite não é positivo.
        """
        self._validar(consulta, limite)
        filtro = "AND s.video_id = ?" if video_id else ""
        parametros: list[object] = [self._consulta_fts(consulta)]
        if video_id:
            parametros.append(video_id)
        parametros.append(limite)
        try:
            linhas = self.conexao.execute(
                f"""SELECT s.id, s.video_id, s.clipe_id, s.inicio, s.fim, s.texto, s.falante,
                           bm25(busca_de_falas) AS nota
                    FROM busca_de_falas
                    JOIN segmentos_de_transcricao s ON s.id = busca_de_falas.rowid
                    WHERE busca_de_falas MATCH ? {filtro}
                    ORDER BY nota
                    LIMIT ?""",
                parametros,
            ).fetchall()
        except sqlite3.OperationalError as erro:
            raise ErroDeBusca(f"Consulta lexical inválida: {erro}") from erro

        return [
            ResultadoDeBusca(
                video_id=linha["video_id"], clipe_id=linha["clipe_id"],
                inicio=linha["inicio"], fim=linha["fim"], texto=linha["texto"],
                falante=linha["falante"], pontuacao=-float(linha["nota"]),
                origem="lexical", falas=(int(linha["id"]),),
            )
            for linha in linhas
        ]

    def buscar_semantica(
        self, consulta: str, video_id: str | None = None, limite: int = 10,
    ) -> list[ResultadoDeBusca]:
        """
        Busca por sentido nos enunciados, comparando embeddings.

        Percorre os vetores em memória: no volume deste acervo (milhares de
        enunciados) isso custa milissegundos, e evita depender de uma extensão
        de banco vetorial. Se o acervo crescer uma ordem de grandeza, o ponto
        de troca por um índice aproximado é aqui, sem mexer no esquema.

        Parâmetros:
            consulta: Texto a procurar por sentido.
            video_id: Restringe a busca a um vídeo.
            limite: Máximo de resultados devolvidos.

        Retorna:
            Os enunciados mais próximos da consulta. Lista vazia quando não há
            provider configurado ou nenhum enunciado foi embedado.

        Pode gerar:
            ErroDeBusca: quando a consulta é vazia ou o limite não é positivo.
            ErroDeEmbedding: quando o provider existe mas falha ao responder.
        """
        self._validar(consulta, limite)
        if self.provider is None:
            return []

        vetor_da_consulta = self.provider.gerar_para_consulta(consulta)
        filtro = "WHERE e.video_id = ?" if video_id else ""
        parametros = (video_id,) if video_id else ()
        candidatos = []
        for linha in self.conexao.execute(
            f"""SELECT e.id, e.video_id, e.clipe_id, e.inicio, e.fim, e.texto, e.falante,
                       b.vetor
                FROM enunciados e
                JOIN embeddings_de_enunciado b ON b.enunciado_id = e.id
                {filtro}""",
            parametros,
        ):
            similaridade = self._cosseno(vetor_da_consulta, desempacotar(linha["vetor"]))
            candidatos.append((similaridade, linha))

        candidatos.sort(key=lambda item: item[0], reverse=True)
        return [
            ResultadoDeBusca(
                video_id=linha["video_id"], clipe_id=linha["clipe_id"],
                inicio=linha["inicio"], fim=linha["fim"], texto=linha["texto"],
                falante=linha["falante"], pontuacao=similaridade,
                origem="semantica", falas=self._falas_do_enunciado(int(linha["id"])),
            )
            for similaridade, linha in candidatos[:limite]
        ]

    def _falas_do_enunciado(self, enunciado_id: int) -> tuple[int, ...]:
        """Lê os ids das falas que compõem um enunciado, em ordem."""
        return tuple(
            int(linha["segmento_id"])
            for linha in self.conexao.execute(
                "SELECT segmento_id FROM falas_do_enunciado WHERE enunciado_id = ? ORDER BY ordem",
                (enunciado_id,),
            )
        )

    def _fundir(
        self,
        lexicais: list[ResultadoDeBusca],
        semanticos: list[ResultadoDeBusca],
        limite: int,
    ) -> list[ResultadoDeBusca]:
        """Funde as duas listas por posição (RRF), somando o peso de cada via."""
        acumulado: dict[tuple[str, str, int], tuple[float, set[str], ResultadoDeBusca]] = {}
        for lista in (lexicais, semanticos):
            for posicao, resultado in enumerate(lista):
                chave = (resultado.video_id, resultado.clipe_id, round(resultado.inicio * 100))
                peso = 1.0 / (CONSTANTE_DE_FUSAO + posicao + 1)
                if chave in acumulado:
                    nota, origens, guardado = acumulado[chave]
                    origens.add(resultado.origem)
                    acumulado[chave] = (nota + peso, origens, guardado)
                else:
                    acumulado[chave] = (peso, {resultado.origem}, resultado)

        ordenados = sorted(acumulado.values(), key=lambda item: item[0], reverse=True)
        return [
            ResultadoDeBusca(
                video_id=guardado.video_id, clipe_id=guardado.clipe_id,
                inicio=guardado.inicio, fim=guardado.fim, texto=guardado.texto,
                falante=guardado.falante, pontuacao=round(nota, 6),
                origem="ambas" if len(origens) > 1 else next(iter(origens)),
                falas=guardado.falas,
            )
            for nota, origens, guardado in ordenados[:limite]
        ]

    @staticmethod
    def _consulta_fts(consulta: str) -> str:
        """
        Monta a expressão do FTS5 a partir do texto digitado.

        Cada palavra vira um termo com prefixo, e os termos são exigidos
        juntos. Escapar em aspas evita que pontuação do usuário seja
        interpretada como operador do FTS5 e derrube a consulta.
        """
        termos = [palavra for palavra in consulta.replace('"', " ").split() if palavra]
        return " AND ".join(f'"{termo}"*' for termo in termos)

    @staticmethod
    def _cosseno(primeiro: tuple[float, ...], segundo: tuple[float, ...]) -> float:
        """Similaridade de cosseno entre dois vetores, ou 0 se algum for nulo."""
        if len(primeiro) != len(segundo):
            return 0.0
        produto = sum(a * b for a, b in zip(primeiro, segundo))
        norma_primeiro = math.sqrt(sum(a * a for a in primeiro))
        norma_segundo = math.sqrt(sum(b * b for b in segundo))
        if norma_primeiro == 0.0 or norma_segundo == 0.0:
            return 0.0
        return produto / (norma_primeiro * norma_segundo)

    @staticmethod
    def _validar(consulta: str, limite: int) -> None:
        """Recusa consulta vazia ou limite não positivo antes de tocar o banco."""
        if not isinstance(consulta, str) or not consulta.strip():
            raise ErroDeBusca("A consulta não pode ser vazia.")
        if limite <= 0:
            raise ErroDeBusca("O limite precisa ser maior que zero.")


class IndexadorSemantico:
    """
    Gera e grava os embeddings dos enunciados de um vídeo.

    Atributos:
        conexao: Conexão SQLite já aberta e com o esquema aplicado.
        provider: Provider que transforma o texto de cada enunciado em vetor.
    """

    def __init__(self, conexao: sqlite3.Connection, provider: ProviderDeEmbeddings) -> None:
        """
        Inicializa o indexador com a conexão e o provider de embeddings.

        Parâmetros:
            conexao: Conexão SQLite já aberta e com o esquema aplicado.
            provider: Provider que gera os vetores.
        """
        self.conexao = conexao
        self.provider = provider

    def indexar(self, video_id: str, refazer: bool = False) -> int:
        """
        Gera os embeddings faltantes dos enunciados de um vídeo.

        Parâmetros:
            video_id: Vídeo cujos enunciados serão indexados.
            refazer: Quando verdadeiro, regenera também os que já têm vetor —
                necessário ao trocar de modelo de embeddings.

        Retorna:
            A quantidade de enunciados indexados nesta execução.

        Pode gerar:
            ErroDeEmbedding: quando o provider falha ao gerar algum vetor.
        """
        from ..integracoes.embeddings import empacotar

        filtro = "" if refazer else "AND b.enunciado_id IS NULL"
        pendentes = self.conexao.execute(
            f"""SELECT e.id, e.texto FROM enunciados e
                LEFT JOIN embeddings_de_enunciado b ON b.enunciado_id = e.id
                WHERE e.video_id = ? {filtro}
                ORDER BY e.inicio""",
            (video_id,),
        ).fetchall()

        indexados = 0
        for linha in pendentes:
            if not linha["texto"].strip():
                continue
            vetor = self.provider.gerar_para_documento(linha["texto"])
            with self.conexao:
                self.conexao.execute(
                    """INSERT INTO embeddings_de_enunciado
                       (enunciado_id, modelo, dimensoes, vetor)
                       VALUES (?, ?, ?, ?)
                       ON CONFLICT (enunciado_id) DO UPDATE SET
                           modelo = excluded.modelo,
                           dimensoes = excluded.dimensoes,
                           vetor = excluded.vetor,
                           gerado_em = strftime('%Y-%m-%dT%H:%M:%fZ','now')""",
                    (linha["id"], self.provider.modelo, self.provider.dimensoes,
                     empacotar(vetor)),
                )
            indexados += 1
        return indexados


__all__ = [
    "BuscaDeConteudo",
    "ErroDeBusca",
    "ErroDeEmbedding",
    "IndexadorSemantico",
    "ResultadoDeBusca",
]