"""Учебный поиск по трём документам. Требуется запущенный Ollama и embeddinggemma."""
import json
import math
import urllib.request

MODEL = "embeddinggemma"
DOCUMENTS = [
    ("A", "Пароль можно восстановить через ссылку на странице входа."),
    ("B", "Счёт за подписку находится в разделе Оплата."),
    ("C", "Пригласить коллегу можно в настройках команды."),
]
QUERY = "Не помню пароль, как снова войти?"


def cosine(a, b):
    if not a or len(a) != len(b):
        raise ValueError("Векторы пусты или имеют разную длину")
    if not all(math.isfinite(x) for x in a + b):
        raise ValueError("Вектор содержит некорректное число")
    denominator = math.sqrt(sum(x*x for x in a) * sum(x*x for x in b))
    if not denominator:
        raise ValueError("Нельзя сравнить нулевой вектор")
    return sum(x*y for x, y in zip(a, b)) / denominator


def main():
    payload = {"model": MODEL, "input": [t for _, t in DOCUMENTS] + [QUERY], "truncate": False}
    request = urllib.request.Request(
        "http://127.0.0.1:11434/api/embed",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(request, timeout=120) as response:
        result = json.load(response)
    vectors = result["embeddings"]
    if len(vectors) != len(DOCUMENTS) + 1:
        raise ValueError("Количество векторов не соответствует входным текстам")
    ranked = sorted(
        [(cosine(v, vectors[-1]), key, text)
         for (key, text), v in zip(DOCUMENTS, vectors[:-1])], reverse=True
    )
    print("Модель:", result.get("model", MODEL), "Размерность:", len(vectors[-1]))
    for score, key, text in ranked:
        print(f"{key}: {score:.4f} — {text}")
    print("Редакционный ориентир: документ A должен быть первым; это нужно проверить на запуске.")


if __name__ == "__main__":
    main()