AI_Agro_Support/tests/test_vector.py

"""Unit tests for Weaviate product retrieval limits."""

from __future__ import annotations

import unittest
from types import SimpleNamespace
from typing import Any

from pipeline.errors import PipelineError
from pipeline.stages.products import ProductCandidate
from pipeline.stages.vector import _ensure_transformers_reranker, search_products


def _candidate(product_id: str) -> ProductCandidate:
    return ProductCandidate(
        product_id=product_id,
        product_name=product_id.upper(),
        registration_number=f"reg-{product_id}",
    )


def _hit(product_id: str, distance: float = 0.1) -> SimpleNamespace:
    return SimpleNamespace(
        properties={"product_id": product_id, "product_name": product_id.upper()},
        metadata=SimpleNamespace(distance=distance),
    )


class FakeNearTextQuery:
    def __init__(self, hits_by_query: dict[str, list[Any]]) -> None:
        self.hits_by_query = hits_by_query
        self.calls: list[dict[str, Any]] = []

    def near_text(
        self,
        query: str,
        limit: int,
        filters: Any,
        return_metadata: Any,
        return_properties: list[str],
        rerank: Any = None,
    ) -> SimpleNamespace:
        self.calls.append(
            {"query": query, "limit": limit, "filters": filters, "rerank": rerank}
        )
        hits = list(self.hits_by_query.get(query, []))
        return SimpleNamespace(objects=hits[:limit])


class FakeClient:
    def __init__(self, hits_by_query: dict[str, list[Any]]) -> None:
        self.query = FakeNearTextQuery(hits_by_query)
        self.collections = SimpleNamespace(
            get=lambda _name: SimpleNamespace(query=self.query)
        )


class SearchProductsLimitTests(unittest.TestCase):
    def test_one_query_keeps_top_four(self) -> None:
        candidates = [_candidate(f"p{i}") for i in range(1, 13)]
        hits = [_hit(f"p{i}", distance=0.1 * i) for i in range(1, 13)]
        client = FakeClient({"q1": hits})

        results = search_products(client, ["q1"], candidates)

        self.assertEqual([item.product_id for item in results], ["p1", "p2", "p3", "p4"])
        self.assertEqual([item.query_index for item in results], [0, 0, 0, 0])
        self.assertEqual(
            [(c["query"], c["limit"]) for c in client.query.calls],
            [("q1", 10)],
        )
        self._assert_rerank(client, ["q1"])

    def test_two_queries_keep_two_each(self) -> None:
        candidates = [_candidate(f"p{i}") for i in range(1, 13)]
        client = FakeClient(
            {
                "q1": [_hit(f"p{i}") for i in range(1, 7)],
                "q2": [_hit(f"p{i}") for i in range(7, 13)],
            }
        )

        results = search_products(client, ["q1", "q2"], candidates)

        self.assertEqual(
            [(item.product_id, item.query_index) for item in results],
            [("p1", 0), ("p2", 0), ("p7", 1), ("p8", 1)],
        )
        self.assertEqual(
            [(c["query"], c["limit"]) for c in client.query.calls],
            [("q1", 6), ("q2", 6)],
        )
        self._assert_rerank(client, ["q1", "q2"])

    def test_overlapping_hits_keep_first_occurrence(self) -> None:
        candidates = [_candidate(f"p{i}") for i in range(1, 9)]
        client = FakeClient(
            {
                "q1": [_hit(f"p{i}") for i in range(1, 7)],
                "q2": [
                    _hit("p2"),
                    _hit("p1"),
                    _hit("p4"),
                    _hit("p5"),
                    _hit("p7"),
                    _hit("p8"),
                ],
            }
        )

        results = search_products(client, ["q1", "q2"], candidates)

        self.assertEqual(
            [(item.product_id, item.query_index) for item in results],
            [("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)],
        )
        self.assertEqual(
            [(c["query"], c["limit"]) for c in client.query.calls],
            [("q1", 6), ("q2", 6)],
        )
        self._assert_rerank(client, ["q1", "q2"])

    def test_large_allowlist_still_filters_and_keeps_top_four(self) -> None:
        candidates = [_candidate(f"p{i}") for i in range(1, 202)]
        hits = [_hit(f"p{i}") for i in range(1, 15)]
        client = FakeClient({"q1": hits})

        results = search_products(client, ["q1"], candidates)

        self.assertEqual([item.product_id for item in results], ["p1", "p2", "p3", "p4"])
        self.assertEqual(len(client.query.calls), 1)
        self.assertEqual(client.query.calls[0]["limit"], 10)
        self._assert_rerank(client, ["q1"])

    def _assert_rerank(self, client: FakeClient, queries: list[str]) -> None:
        self.assertEqual(len(client.query.calls), len(queries))
        for call, query in zip(client.query.calls, queries):
            self.assertIsNotNone(call["filters"])
            self.assertEqual(call["rerank"].prop, "situation_fit_summary")
            self.assertEqual(call["rerank"].query, query)


class EnsureRerankerTests(unittest.TestCase):
    @staticmethod
    def _client(modules: dict[str, Any], reranker: Any) -> SimpleNamespace:
        updates: list[Any] = []
        config = SimpleNamespace(
            get=lambda: SimpleNamespace(reranker_config=reranker),
            update=lambda **kwargs: updates.append(kwargs),
        )
        client = SimpleNamespace(
            get_meta=lambda: {"modules": modules},
            collections=SimpleNamespace(get=lambda _name: SimpleNamespace(config=config)),
        )
        client.updates = updates
        return client

    def test_missing_module_raises_clear_error(self) -> None:
        client = self._client({"text2vec-google": {}}, None)

        with self.assertRaises(PipelineError) as ctx:
            _ensure_transformers_reranker(client)

        self.assertIn("docker compose up -d", str(ctx.exception))
        self.assertEqual(client.updates, [])

    def test_sets_reranker_when_missing(self) -> None:
        client = self._client({"reranker-transformers": {}}, None)

        _ensure_transformers_reranker(client)

        self.assertEqual(len(client.updates), 1)
        self.assertIn("reranker_config", client.updates[0])

    def test_existing_reranker_is_left_alone(self) -> None:
        current = SimpleNamespace(reranker="reranker-transformers")
        client = self._client({"reranker-transformers": {}}, current)

        _ensure_transformers_reranker(client)

        self.assertEqual(client.updates, [])


if __name__ == "__main__":
    unittest.main()