jhonny-editor/code/engine/integracoes/embeddings/provider_ollama.py

"""
Geração de embeddings por um Ollama local.

O projeto já roda um Ollama nesta máquina com o ``nomic-embed-text`` — o mesmo
modelo que o RAG de código usa. Reaproveitá-lo mantém a análise de conteúdo
local: nenhum trecho de transcrição de cliente sai da máquina, e não há custo
por chamada.

O ``nomic-embed-text`` exige prefixos distintos para o que é indexado e para o
que é consultado. Sem eles a similaridade cai de forma silenciosa, então os
prefixos são aplicados aqui e não ficam a cargo de quem chama.
"""

from __future__ import annotations

import json
import struct
import urllib.error
import urllib.request
from typing import Protocol, Sequence

URL_PADRAO = "http://localhost:11434"
MODELO_PADRAO = "nomic-embed-text"
DIMENSOES_PADRAO = 768
PREFIXO_DE_DOCUMENTO = "search_document: "
PREFIXO_DE_CONSULTA = "search_query: "


class ErroDeEmbedding(RuntimeError):
    """Falha ao gerar um embedding — serviço fora do ar, modelo ausente ou resposta inválida."""


class ProviderDeEmbeddings(Protocol):
    """Contrato de quem transforma texto em vetor."""

    @property
    def modelo(self) -> str:
        """Nome do modelo que gerou os vetores."""
        ...

    @property
    def dimensoes(self) -> int:
        """Quantidade de dimensões dos vetores gerados."""
        ...

    def gerar_para_documento(self, texto: str) -> tuple[float, ...]:
        """Gera o vetor de um texto que será indexado."""
        ...

    def gerar_para_consulta(self, texto: str) -> tuple[float, ...]:
        """Gera o vetor de um texto que está sendo buscado."""
        ...


class ProviderDeEmbeddingsOllama:
    """
    Gera embeddings chamando um Ollama local por HTTP.

    Atributos:
        modelo: Nome do modelo de embeddings usado.
        dimensoes: Dimensões esperadas nos vetores devolvidos.
        url: Endereço base do serviço Ollama.
        tempo_limite: Segundos a esperar por cada resposta.
    """

    def __init__(
        self,
        modelo: str = MODELO_PADRAO,
        dimensoes: int = DIMENSOES_PADRAO,
        url: str = URL_PADRAO,
        tempo_limite: float = 60.0,
    ) -> None:
        """
        Inicializa o provider com o modelo e o endereço do serviço.

        Parâmetros:
            modelo: Nome do modelo de embeddings no Ollama.
            dimensoes: Dimensões esperadas — vetores de outro tamanho são
                recusados, para não misturar modelos no mesmo índice.
            url: Endereço base do Ollama.
            tempo_limite: Segundos a esperar por cada resposta.
        """
        self._modelo = modelo
        self._dimensoes = dimensoes
        self.url = url.rstrip("/")
        self.tempo_limite = tempo_limite

    @property
    def modelo(self) -> str:
        """Nome do modelo que gera os vetores."""
        return self._modelo

    @property
    def dimensoes(self) -> int:
        """Quantidade de dimensões dos vetores gerados."""
        return self._dimensoes

    def gerar_para_documento(self, texto: str) -> tuple[float, ...]:
        """
        Gera o vetor de um enunciado que será indexado.

        Parâmetros:
            texto: Texto do enunciado.

        Retorna:
            O vetor normalizado do texto.

        Pode gerar:
            ErroDeEmbedding: quando o Ollama não responde ou devolve um vetor
                com dimensões diferentes das esperadas.
        """
        return self._gerar(PREFIXO_DE_DOCUMENTO + texto)

    def gerar_para_consulta(self, texto: str) -> tuple[float, ...]:
        """
        Gera o vetor de uma consulta feita pelo usuário.

        Parâmetros:
            texto: Texto da consulta.

        Retorna:
            O vetor normalizado da consulta.

        Pode gerar:
            ErroDeEmbedding: quando o Ollama não responde ou devolve um vetor
                com dimensões diferentes das esperadas.
        """
        return self._gerar(PREFIXO_DE_CONSULTA + texto)

    def _gerar(self, texto: str) -> tuple[float, ...]:
        """Chama o Ollama e valida o vetor devolvido."""
        corpo = json.dumps({"model": self._modelo, "prompt": texto}).encode("utf-8")
        requisicao = urllib.request.Request(
            f"{self.url}/api/embeddings", data=corpo,
            headers={"Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(requisicao, timeout=self.tempo_limite) as resposta:
                dados = json.loads(resposta.read().decode("utf-8"))
        except urllib.error.URLError as erro:
            raise ErroDeEmbedding(
                f"Não foi possível falar com o Ollama em {self.url}: {erro}"
            ) from erro
        except json.JSONDecodeError as erro:
            raise ErroDeEmbedding(f"Resposta do Ollama não é JSON válido: {erro}") from erro

        vetor = dados.get("embedding")
        if not isinstance(vetor, list) or not vetor:
            raise ErroDeEmbedding(f"O Ollama não devolveu embedding para o modelo {self._modelo}.")
        if len(vetor) != self._dimensoes:
            raise ErroDeEmbedding(
                f"Embedding com {len(vetor)} dimensões, esperado {self._dimensoes}."
            )
        return tuple(float(valor) for valor in vetor)


def empacotar(vetor: Sequence[float]) -> bytes:
    """
    Serializa um vetor para o BLOB gravado no banco.

    Parâmetros:
        vetor: Valores do vetor.

    Retorna:
        Os valores como float de 32 bits em sequência.
    """
    return struct.pack(f"<{len(vetor)}f", *vetor)


def desempacotar(dados: bytes) -> tuple[float, ...]:
    """
    Reconstrói um vetor a partir do BLOB lido do banco.

    Parâmetros:
        dados: Bytes gravados por ``empacotar``.

    Retorna:
        Os valores do vetor.
    """
    return struct.unpack(f"<{len(dados) // 4}f", dados)