Newer
Older
vmk-360_data_mcp / src / vmk_data_mcp / embedder.py
"""Клиент для генерации эмбеддингов через Ollama API."""

import httpx

from vmk_data_mcp.config import settings

_OLLAMA_CLIENT: httpx.AsyncClient | None = None


def _get_client() -> httpx.AsyncClient:
    """Возвращает (или создаёт) асинхронный HTTP-клиент."""
    global _OLLAMA_CLIENT
    if _OLLAMA_CLIENT is None:
        _OLLAMA_CLIENT = httpx.AsyncClient(
            base_url=settings.ollama_base_url,
            timeout=httpx.Timeout(settings.ollama_request_timeout),
        )
    return _OLLAMA_CLIENT


async def get_embedding(text: str) -> list[float]:
    """Генерирует вектор эмбеддинга для текста через Ollama.

    Args:
        text: Текст для эмбеддинга (должен быть на украинском).

    Returns:
        Список float размерностью `ollama_embed_dimensions`.

    Raises:
        httpx.HTTPError: при ошибке сети или HTTP.
        ValueError: если размерность ответа не совпадает с ожидаемой.
    """
    client = _get_client()
    payload = {
        "model": settings.ollama_embed_model,
        "input": text,
    }

    response = await client.post("/api/embed", json=payload)
    response.raise_for_status()
    data = response.json()

    # Ollama /api/embed возвращает "embeddings": [[0.1, 0.2, ...]]
    embeddings = data.get("embeddings")
    if not embeddings or not isinstance(embeddings, list):
        raise ValueError(f"Unexpected Ollama response format: {data}")

    vector = embeddings[0]
    if len(vector) != settings.ollama_embed_dimensions:
        raise ValueError(
            f"Embedding dimension mismatch: expected {settings.ollama_embed_dimensions}, "
            f"got {len(vector)}"
        )

    return list(vector)


async def close_client() -> None:
    """Закрывает HTTP-клиент."""
    global _OLLAMA_CLIENT
    if _OLLAMA_CLIENT is not None:
        await _OLLAMA_CLIENT.aclose()
        _OLLAMA_CLIENT = None