hlamingo/hlamingo/ollama.py

"""Клиент Ollama: чат и эмбеддинги.

Копия нужной части emilia/ollama.py, а не импорт из соседнего проекта: Эмилия
активно развивается, и общий модуль сделал бы hlamingo заложником её рефакторингов.

Адресов два, а не один. Разбор фраз, теги и фото идут на `config.chat_host()` —
это либо Ollama, либо сервис deepseek-api, который держит вкладку веб-DeepSeek
и отдаёт наружу тот же протокол. Эмбеддинги всегда идут на OLLAMA_HOST: веб-чат
их не считает, и подменять их нечем.
"""

import io
import json
import time
import base64
import socket
import logging
import urllib.parse
import urllib.request

# Адрес и модели читаем через `config.X`, а не копией при импорте: их можно
# поменять из браузера, и это должно действовать сразу, без перезапуска.
from . import config
from .config import (CHAT_TIMEOUT, EMBED_TIMEOUT, MAX_OUTPUT_TOKENS,
                     EMBED_DOC_PREFIX, EMBED_QUERY_PREFIX)

logger = logging.getLogger(__name__)

# ============================================================
# Предохранитель
#
# Если машина с Ollama просто выключена, пакеты уходят в никуда, и каждый
# запрос честно ждёт свой таймаут. Замер: запись одной вещи занимала 92 секунды,
# разбор фразы — 154. Пользователь решит, что приложение зависло.
#
# Поэтому: сначала быстрая проверка TCP (полтора секунды вместо двух минут),
# а после неудачи — пауза, в течение которой к серверу вообще не ходим.
# ============================================================
# Бюджет ответа для моделей, которые всегда размышляют: размышления идут
# в счёт того же лимита, и обычных 512 токенов им не хватает даже на короткий
# ответ. Замер на qwen3-vl: при 200 токенах content пуст, при 400 — уже ответ.
THINKING_BUDGET = 1200

CONNECT_TIMEOUT = 1.5     # столько ждём само соединение
DOWN_COOLDOWN = 30.0      # столько не тревожим сервер после неудачи
UP_CACHE = 5.0            # столько верим, что сервер жив, без повторной проверки

# Предохранитель считается по каждому адресу отдельно: лежачая Ollama не должна
# гасить работающий deepseek-api и наоборот.
_down_until = {}
_up_until = {}


def _host_port(host):
    parsed = urllib.parse.urlparse(host)
    return parsed.hostname or "localhost", parsed.port or (443 if parsed.scheme == "https" else 80)


def available(host=None):
    """Стоит ли вообще пытаться. Быстро и без долгих ожиданий."""
    host = host or config.chat_host()
    now = time.time()
    if now < _up_until.get(host, 0.0):
        return True
    if now < _down_until.get(host, 0.0):
        return False
    try:
        socket.create_connection(_host_port(host), timeout=CONNECT_TIMEOUT).close()
        _mark_up(host)
        return True
    except Exception:
        _mark_down(host)
        return False


def _mark_down(host):
    """Запрос не удался — держим паузу, чтобы не ждать таймаут на каждом шаге."""
    _down_until[host] = time.time() + DOWN_COOLDOWN
    _up_until.pop(host, None)


def _mark_up(host):
    _up_until[host] = time.time() + UP_CACHE
    _down_until.pop(host, None)


def reset_breaker(host=None):
    """Сбросить паузу — например, когда адрес поменяли в настройках.

    Без аргумента забываем всё: движок могли переключить, и прошлые неудачи
    ничего не говорят о новом адресе.
    """
    if host is None:
        _down_until.clear()
        _up_until.clear()
    else:
        _down_until.pop(host, None)
        _up_until.pop(host, None)


def check(host=None):
    """Доступен ли сервер. Отвечает быстро: это дёргает баннер в UI."""
    host = host or config.chat_host()
    if not available(host):
        return False
    try:
        req = urllib.request.Request(host + "/api/tags")
        with urllib.request.urlopen(req, timeout=3) as r:
            if r.status == 200:
                _mark_up(host)
                return True
    except Exception:
        # Пока браузер с чатом не поднят, deepseek-api отвечает 503 — это тоже
        # «недоступен», и сюда оно приходит исключением HTTPError.
        _mark_down(host)
    return False


def models(host=None):
    """Список моделей на этом адресе (для диагностики и выпадающих списков)."""
    host = host or config.chat_host()
    if not available(host):
        return []
    try:
        req = urllib.request.Request(host + "/api/tags")
        with urllib.request.urlopen(req, timeout=5) as r:
            data = json.loads(r.read().decode("utf-8"))
        _mark_up(host)
        return [m.get("model", "") for m in data.get("models", [])]
    except Exception as e:
        _mark_down(host)
        logger.warning("Список моделей не получен: %s", e)
        return []


def chat(model, messages, timeout=None, num_predict=MAX_OUTPUT_TOKENS,
         think=False):
    """Один запрос к /api/chat. Возвращает content или None при любой беде.

    num_predict обязателен: без капа модель может уйти в длинную генерацию
    и подвесить ход на минуты — а нам нужен короткий JSON.

    think=False выключает режим размышлений у моделей-ризонеров (qwen3,
    deepseek-r1 и подобных). Без этого они тратят весь бюджет ответа на
    внутренний монолог и возвращают пустой content — проверено на qwen3:8b.
    Модели без такого режима поле просто игнорируют.
    """
    host = config.chat_host()
    timeout = timeout or config.chat_timeout()
    payload = {
        "model": model,
        "messages": messages,
        "stream": False,
        "think": think,
        "options": {"num_predict": num_predict, "temperature": 0},
    }
    if not available(host):
        return None          # сервер лежит: не ждём таймаут на каждой записи
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        host + "/api/chat", data=data,
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            resp = json.loads(r.read().decode("utf-8"))
        _mark_up(host)
    except Exception as e:
        _mark_down(host)
        logger.warning("Запрос к модели %s не удался: %s", model, e)
        return None

    message = resp.get("message") or {}
    content = message.get("content", "")

    # Модели вроде qwen3-vl размышляют всегда, и think=False их не выключает:
    # весь бюджет уходит в поле thinking, а content остаётся пустым при
    # done_reason="length". Один раз повторяем с запасом — иначе такая модель
    # выглядит нерабочей, хотя дело только в лимите.
    if not content.strip() and message.get("thinking") \
            and resp.get("done_reason") == "length" and num_predict < THINKING_BUDGET:
        logger.info("Модель %s ушла в размышления — повторяю с бюджетом %d",
                    model, THINKING_BUDGET)
        return chat(model, messages, timeout=timeout,
                    num_predict=THINKING_BUDGET, think=think)
    return content


def look(image_bytes, prompt, model=None, budget=None, timeout=180):
    """Показать модели картинку и получить текстовый ответ. None при неудаче.

    Картинка уменьшается перед отправкой: изображение разворачивается в токены,
    и снимок с телефона стоит их тысячами без пользы для чтения крупных надписей.
    """
    host = config.chat_host()
    model = model or config.VISION_MODEL
    if not model or not image_bytes or not available(host):
        return None

    small = _shrink(image_bytes, config.VISION_MAX_SIDE)
    payload = {
        "model": model,
        "stream": False,
        "messages": [{"role": "user", "content": prompt,
                      "images": [base64.b64encode(small).decode("ascii")]}],
        "options": {"num_predict": budget or config.VISION_BUDGET, "temperature": 0},
    }
    req = urllib.request.Request(
        host + "/api/chat",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            resp = json.loads(r.read().decode("utf-8"))
        _mark_up(host)
        return ((resp.get("message") or {}).get("content") or "").strip()
    except Exception as e:
        _mark_down(host)
        logger.warning("Не удалось распознать фото моделью %s: %s", model, e)
        return None


def _shrink(image_bytes, side):
    """Уменьшает картинку по длинной стороне. Без Pillow отдаёт как есть."""
    try:
        from PIL import Image
    except ImportError:
        return image_bytes
    try:
        img = Image.open(io.BytesIO(image_bytes))
        img = img.convert("RGB")
        if max(img.size) > side:
            ratio = side / float(max(img.size))
            img = img.resize((max(1, int(img.width * ratio)),
                              max(1, int(img.height * ratio))), Image.LANCZOS)
        out = io.BytesIO()
        img.save(out, format="JPEG", quality=85)
        return out.getvalue()
    except Exception as e:
        logger.warning("Не удалось уменьшить картинку: %s", e)
        return image_bytes


def embed(text, is_query=False):
    """Вектор текста или None. is_query=True для поисковых запросов.

    Префиксы обязательны: nomic-embed-text обучен с ними, и без разделения
    «документ / запрос» качество поиска заметно падает.
    """
    text = (text or "").strip()
    if not text:
        return None
    # Эмбеддинги всегда у настоящей Ollama, даже когда думает DeepSeek:
    # веб-чат векторов не отдаёт.
    host = config.OLLAMA_HOST
    if not available(host):
        return None
    prefix = EMBED_QUERY_PREFIX if is_query else EMBED_DOC_PREFIX
    data = json.dumps({"model": config.EMBED_MODEL, "prompt": prefix + text}).encode("utf-8")
    req = urllib.request.Request(
        host + "/api/embeddings", data=data,
        headers={"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(req, timeout=EMBED_TIMEOUT) as r:
            vector = json.loads(r.read().decode("utf-8")).get("embedding")
        _mark_up(host)
        return vector
    except Exception as e:
        _mark_down(host)
        logger.warning("Эмбеддинг не получен: %s", e)
        return None