pi-qmem/gateway/tests/test_embed_chain.py

"""Test della catena di fallback per gli embedding e del wrapper Qdrant resilient."""
from __future__ import annotations

import asyncio
import json

import httpx
import pytest

import embed as embed_mod
import state
from test_api import auth_headers, make_record

CHAIN = json.dumps(
    [
        {"name": "primario", "url": "http://primario:9001", "api": "llamacpp", "key": "k1", "timeout_ms": 500},
        {"name": "fallback", "url": "http://fallback:9002", "api": "ollama", "key": "k2", "timeout_ms": 5000},
    ]
)

VEC_1024 = [0.01] * 1024


@pytest.fixture(autouse=True)
def _reset_chain(monkeypatch):
    embed_mod.reset_chain_cache()
    embed_mod._http = None
    yield
    embed_mod.reset_chain_cache()
    embed_mod._http = None


def _mock_client(handler) -> list[str]:
    calls: list[str] = []

    def tracking_handler(request):
        calls.append(f"{request.url.host}{request.url.path}")
        return handler(request)

    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(tracking_handler))
    return calls


# ---------------------------------------------------------------------------
# parse della catena
# ---------------------------------------------------------------------------


def test_parse_chain_valida(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()
    nodes = embed_mod.chain_nodes()
    assert [n.name for n in nodes] == ["primario", "fallback"]
    assert [n.api for n in nodes] == ["llamacpp", "ollama"]


def test_parse_legacy_quando_catena_vuota(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", "")
    monkeypatch.setattr(embed_mod, "EMBED_API", "ollama")
    monkeypatch.setattr(embed_mod, "EMBED_URL", "http://legacy:11434")
    monkeypatch.setattr(embed_mod, "EMBED_API_KEY", "lk")
    embed_mod.reset_chain_cache()
    nodes = embed_mod.chain_nodes()
    assert len(nodes) == 1
    assert nodes[0].name == "embed"
    assert nodes[0].api == "ollama"
    assert nodes[0].url == "http://legacy:11434"
    assert nodes[0].key == "lk"


def test_parse_chain_json_invalido_cade_su_legacy(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", "non-json")
    monkeypatch.setattr(embed_mod, "EMBED_URL", "http://legacy:11434")
    embed_mod.reset_chain_cache()
    assert [n.url for n in embed_mod.chain_nodes()] == ["http://legacy:11434"]


# ---------------------------------------------------------------------------
# fallback e cooldown
# ---------------------------------------------------------------------------


def test_fallback_primario_llamacpp_fallito(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()

    def handler(request):
        if request.url.host == "primario":
            return httpx.Response(500)
        # nodo ollama-style: risposta con campo "embeddings"
        return httpx.Response(200, json={"embeddings": [VEC_1024]})

    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    vector = asyncio.run(embed_mod.embed("test"))
    assert vector == VEC_1024


def test_cooldown_salta_primario(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()
    calls: list[str] = []

    def tracking(request):
        calls.append(request.url.host)
        if request.url.host == "primario":
            return httpx.Response(500)
        return httpx.Response(200, json={"embeddings": [VEC_1024]})

    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(tracking))
    asyncio.run(embed_mod.embed("test"))
    asyncio.run(embed_mod.embed("test"))
    # il primario fallito entra in cooldown: la seconda chiamata lo salta
    assert calls == ["primario", "fallback", "fallback"]


def test_dimensione_errata_salta_nodo(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()

    def handler(request):
        if request.url.host == "primario":
            return httpx.Response(200, json={"data": [{"embedding": [0.0] * 512}]})  # dim sbagliata
        return httpx.Response(200, json={"embeddings": [VEC_1024]})

    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    vector = asyncio.run(embed_mod.embed("test"))
    assert len(vector) == 1024


def test_tutti_nodi_falliti_rilancia(monkeypatch):
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()
    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(lambda r: httpx.Response(500)))
    with pytest.raises(RuntimeError, match="tutti i nodi embedding falliti"):
        asyncio.run(embed_mod.embed("test"))


def test_empty_text_comunque_chiamata(monkeypatch):
    # il modello pydantic valida già la query; qui verifichiamo il passthrough
    monkeypatch.setattr(embed_mod, "EMBED_CHAIN", CHAIN)
    embed_mod.reset_chain_cache()

    def handler(request):
        body = json.loads(request.content)
        return httpx.Response(200, json={"embeddings": [VEC_1024] if body["input"] else []})

    embed_mod._http = httpx.AsyncClient(transport=httpx.MockTransport(handler))
    assert asyncio.run(embed_mod.embed("ok")) == VEC_1024


# ---------------------------------------------------------------------------
# ResilientQdrant (retry transiente)
# ---------------------------------------------------------------------------


class FlakyQdrant:
    def __init__(self, failures: int, exc: Exception):
        self.calls = 0
        self.failures = failures
        self.exc = exc

    def upsert(self, **kw):
        self.calls += 1
        if self.calls <= self.failures:
            raise self.exc
        return "ok"


def test_retry_su_errore_transiente(monkeypatch):
    fake = FlakyQdrant(2, httpx.ConnectError("conn"))
    client = state.ResilientQdrant(fake, attempts=3, backoff_s=0.01)
    assert client.upsert(x=1) == "ok"
    assert fake.calls == 3


def test_niente_retry_su_errore_applicativo():
    fake = FlakyQdrant(2, ValueError("404 logico"))
    client = state.ResilientQdrant(fake, attempts=3, backoff_s=0.01)
    with pytest.raises(ValueError):
        client.upsert(x=1)
    assert fake.calls == 1


def test_retry_esaurito_rilancia():
    fake = FlakyQdrant(99, httpx.ReadTimeout("t"))
    client = state.ResilientQdrant(fake, attempts=2, backoff_s=0.01)
    with pytest.raises(httpx.ReadTimeout):
        client.upsert(x=1)
    assert fake.calls == 2