d1569a301f
- Updated README.md to clarify the conditions for `last_applied_treatment` and its integration into the advice JSON structure. - Refactored job.py to utilize a new `TreatmentHistory` class for managing treatment data, improving clarity and efficiency in treatment retrieval. - Enhanced advice_context.py to conditionally include `last_applied_treatment` based on the treatment history, ensuring no redundant queries are made. - Modified products.py to implement logic for excluding active ingredients based on the latest treatment's interval, refining the product prefiltering process. - Updated vector.py to ensure the SQL allowlist is applied correctly, improving the efficiency of product searches. - Added unit tests to validate the new treatment history functionality and product filtering logic.
127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
"""Unit tests for Weaviate product retrieval limits."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
from pipeline.stages.products import ProductCandidate
|
|
from pipeline.stages.vector import 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],
|
|
) -> SimpleNamespace:
|
|
self.calls.append({"query": query, "limit": limit, "filters": filters})
|
|
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, 7)]
|
|
hits = [_hit(f"p{i}", distance=0.1 * i) for i in range(1, 7)]
|
|
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", 4)],
|
|
)
|
|
self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
|
|
|
|
def test_two_queries_keep_two_each(self) -> None:
|
|
candidates = [_candidate(f"p{i}") for i in range(1, 7)]
|
|
client = FakeClient(
|
|
{
|
|
"q1": [_hit("p1"), _hit("p2"), _hit("p3")],
|
|
"q2": [_hit("p4"), _hit("p5"), _hit("p6")],
|
|
}
|
|
)
|
|
|
|
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", 2), ("q2", 4)],
|
|
)
|
|
self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
|
|
|
|
def test_overlapping_hits_keep_first_occurrence(self) -> None:
|
|
candidates = [_candidate(f"p{i}") for i in range(1, 6)]
|
|
client = FakeClient(
|
|
{
|
|
"q1": [_hit("p1"), _hit("p2"), _hit("p3")],
|
|
"q2": [_hit("p2"), _hit("p1"), _hit("p4"), _hit("p5")],
|
|
}
|
|
)
|
|
|
|
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", 2), ("q2", 4)],
|
|
)
|
|
self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
|
|
|
|
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, 10)]
|
|
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"], 4)
|
|
self.assertIsNotNone(client.query.calls[0]["filters"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|