"""Клиент для генерации эмбеддингов через 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