Refactor treatment history handling and enhance product filtering logic

- 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.
This commit is contained in:
Arsham Mirehvandi 2026-09-20 23:38:25 +02:00
parent 884d42e125
commit d1569a301f
10 changed files with 643 additions and 197 deletions

View File

@ -363,7 +363,7 @@ Run once per `(field, disease)` job by `pipeline/job.py::run_job`:
3. **Observation** — `AI_agrosupport_agro_ril_pathogen` (`rilpato_layer`/`rilpato_date`/`rilpato_diffusion`); `observation` is true for `as_of` or any of the past 5 days if a row exists there with `rilpato_diffusion` other than `32` (including `NULL`) 3. **Observation** — `AI_agrosupport_agro_ril_pathogen` (`rilpato_layer`/`rilpato_date`/`rilpato_diffusion`); `observation` is true for `as_of` or any of the past 5 days if a row exists there with `rilpato_diffusion` other than `32` (including `NULL`)
4. **Weather** — `TDatiMeteo_D` for past days, `TDatiMeteo_D_FRC` for today/future 4. **Weather** — `TDatiMeteo_D` for past days, `TDatiMeteo_D_FRC` for today/future
5. **Phenology** — carry-forward of latest observation ≤ day; future days null; if today's phase is missing, use yesterday's 5. **Phenology** — carry-forward of latest observation ≤ day; future days null; if today's phase is missing, use yesterday's
6. **Applied treatments** — products sprayed over `as_of-5 .. as_of`, from `AI_agrosupport_agro_ril_operations*` with `rilop_operation = 10`. A separate lookback over `as_of-25 .. as_of-6` is used only for `last_applied_treatment` during context enrichment. 6. **Applied treatments** — one query loads crop-protection sprays over `as_of-25 .. as_of` from `AI_agrosupport_agro_ril_operations*` with `rilop_operation = 10`. Days in `as_of-5 .. as_of` fill meteorological `applied_treatment`. The latest spray in the full 25-day window drives the prefilter rotation gate. `last_applied_treatment` still uses only `as_of-25 .. as_of-6`, and only when the recent window is empty.
7. **LLM** — synthesise 1–2 anonymised search queries from the subset 7. **LLM** — synthesise 1–2 anonymised search queries from the subset
8. **Product prefilter** — organic / crop / disease gates, then one prefilter rule against the batch-wide, in-memory `products` index (`pipeline/stages/disease.py::select_prefilter_rule`, first match wins): 8. **Product prefilter** — organic / crop / disease gates, then one prefilter rule against the batch-wide, in-memory `products` index (`pipeline/stages/disease.py::select_prefilter_rule`, first match wins):
@ -377,8 +377,10 @@ Run once per `(field, disease)` job by `pipeline/job.py::run_job`:
| `FASE5` today or in the past only | `systemicity` in `{contact, mixed}` and `preventive_action = 1` and `curative_action = 0` and `eradicant_action = 0` | | `FASE5` today or in the past only | `systemicity` in `{contact, mixed}` and `preventive_action = 1` and `curative_action = 0` and `eradicant_action = 0` |
| none of the above | no FASE-specific cut; organic/crop/disease gates only | | none of the above | no FASE-specific cut; organic/crop/disease gates only |
9. **Vector search** — Weaviate `near_text` on `ProductProfile`, filtered server-side to the allowlisted `product_id`s. One synthesized query retrieves the top 4; two queries retrieve the top 2 unique each (the second search over-fetches so overlaps can still fill two products) After those gates, if the latest spray in `as_of-25 .. as_of` still has a non-NULL `product_uses.treatment_interval_min_days` strictly greater than days elapsed, candidates whose `active_ingredients` overlap that product's (tank-mix partners evaluated per product) are dropped. NULL or elapsed interval leaves the list unchanged.
9. **Vector search** — Weaviate `near_text` on `ProductProfile`, always filtered with `contains_any` on the SQL-allowlisted `product_id`s (`product_id` must use `tokenization: field`). One synthesized query retrieves the top 4; two queries retrieve the top 2 unique each (the second search over-fetches so overlaps can still fill two products)
10. **Product JSON** — label details per product from `products`, `product_uses`, `label_chunks`, cached per `(crop, disease)` 10. **Product JSON** — label details per product from `products`, `product_uses`, `label_chunks`, cached per `(crop, disease)`
11. **Context enrichment** — `allowed_products`, `applied_treatment` expansion, `last_advice` (with issuance `date`) from the `advice` table, and optional `last_applied_treatment` (spray date + products) when the recent window is empty but a spray exists within 25 days 11. **Context enrichment** — `allowed_products`, `applied_treatment` expansion, `last_advice` (with issuance `date`) from the `advice` table, and optional `last_applied_treatment` (spray date + products) when the recent window is empty but a spray exists in `as_of-25 .. as_of-6`. Enrichment reuses the job's 25-day treatment history; it does not query operations again.
12. **Advice generation** — second LLM profile returns the structured advisory 12. **Advice generation** — second LLM profile returns the structured advisory
13. **Persistence** — the advisory plus its collapsed input are written to `advice` inside one transaction 13. **Persistence** — the advisory plus its collapsed input are written to `advice` inside one transaction

View File

@ -42,8 +42,11 @@ from pipeline.stages.disease import resolve_disease_forecast_by_anmod, select_pr
from pipeline.stages.llm import synthesize_queries from pipeline.stages.llm import synthesize_queries
from pipeline.stages.observation import load_observation from pipeline.stages.observation import load_observation
from pipeline.stages.phenology import load_phenology from pipeline.stages.phenology import load_phenology
from pipeline.stages.products import prefilter_products from pipeline.stages.products import (
from pipeline.stages.treatments import load_applied_treatments excluded_active_ingredients_for_last_treatment,
prefilter_products,
)
from pipeline.stages.treatments import load_treatment_history
from pipeline.stages.vector import RecommendedProduct from pipeline.stages.vector import RecommendedProduct
from pipeline.stages.weather import load_weather from pipeline.stages.weather import load_weather
from pipeline.worklist import Job from pipeline.worklist import Job
@ -159,8 +162,19 @@ def _run_job_inner(resources: Resources, job: Job, as_of: date, dry_run: bool) -
weather = load_weather(conn, disease.station, as_of) weather = load_weather(conn, disease.station, as_of)
phenology = load_phenology(conn, job.field_id, as_of) phenology = load_phenology(conn, job.field_id, as_of)
applied_treatments = load_applied_treatments(conn, job.field_id, as_of) treatment_history = load_treatment_history(conn, job.field_id, as_of)
observation_by_day = load_observation(conn, job.field_id, as_of) observation_by_day = load_observation(conn, job.field_id, as_of)
product_index = resources.product_index
if product_index is None:
raise RuntimeError("load_product_index() must run before run_job().")
excluded_ais = excluded_active_ingredients_for_last_treatment(
conn,
product_index,
crop_english=job.crop_english,
disease_english=disease.disease_english,
latest_spray=treatment_history.latest(),
as_of=as_of,
)
# has_risk gates query synthesis and the product prefilter: it is true if # has_risk gates query synthesis and the product prefilter: it is true if
# any day is FASE2 / FASE5 / INCUBAZPRIMARIA (disease.has_risk, already # any day is FASE2 / FASE5 / INCUBAZPRIMARIA (disease.has_risk, already
@ -183,7 +197,7 @@ def _run_job_inner(resources: Resources, job: Job, as_of: date, dry_run: bool) -
weather, weather,
phenology, phenology,
disease.forecasts_by_day, disease.forecasts_by_day,
applied_treatments, treatment_history.recent(),
observation_by_day, observation_by_day,
) )
json_for_query_synthesis = build_json_for_query_synthesis( json_for_query_synthesis = build_json_for_query_synthesis(
@ -250,11 +264,12 @@ def _run_job_inner(resources: Resources, job: Job, as_of: date, dry_run: bool) -
prefilter_rule = select_prefilter_rule(disease.forecasts_by_day, observation_by_day, as_of) prefilter_rule = select_prefilter_rule(disease.forecasts_by_day, observation_by_day, as_of)
payload["run"]["prefilter_rule"] = prefilter_rule payload["run"]["prefilter_rule"] = prefilter_rule
candidates = prefilter_products( candidates = prefilter_products(
resources.product_index, product_index,
crop_english=job.crop_english, crop_english=job.crop_english,
disease_english=disease.disease_english, disease_english=disease.disease_english,
organic=job.organic, organic=job.organic,
rule=prefilter_rule, rule=prefilter_rule,
excluded_active_ingredients=excluded_ais,
) )
payload["candidate_count"] = len(candidates) payload["candidate_count"] = len(candidates)
@ -297,8 +312,7 @@ def _run_job_inner(resources: Resources, job: Job, as_of: date, dry_run: bool) -
last_advice_summary=last_summary, last_advice_summary=last_summary,
last_advice_treatments=last_treatments, last_advice_treatments=last_treatments,
last_advice_date=last_date, last_advice_date=last_date,
as_of=as_of, treatment_history=treatment_history,
field_id=job.field_id,
) )
payload["last_advice"] = json_for_advice_generation["last_advice"] payload["last_advice"] = json_for_advice_generation["last_advice"]

View File

@ -12,7 +12,7 @@ import pyodbc
from pipeline.db import fetch_one from pipeline.db import fetch_one
from pipeline.jsonutils import parse_json_list from pipeline.jsonutils import parse_json_list
from pipeline.stages.product_json import ProductJsonBuilder, normalize_name from pipeline.stages.product_json import ProductJsonBuilder, normalize_name
from pipeline.stages.treatments import load_last_treatment_before_recent_window from pipeline.stages.treatments import TreatmentHistory
from pipeline.stages.vector import RecommendedProduct from pipeline.stages.vector import RecommendedProduct
from pipeline.window import format_date from pipeline.window import format_date
@ -98,8 +98,7 @@ def enrich_advice_json(
recommended: Iterable[RecommendedProduct], recommended: Iterable[RecommendedProduct],
last_advice_summary: str, last_advice_summary: str,
last_advice_treatments: Iterable[str], last_advice_treatments: Iterable[str],
as_of: date, treatment_history: TreatmentHistory,
field_id: int,
last_advice_date: date | None = None, last_advice_date: date | None = None,
) -> None: ) -> None:
""" """
@ -115,6 +114,8 @@ def enrich_advice_json(
When allowed_products is non-empty, the meteorological window has no sprays, When allowed_products is non-empty, the meteorological window has no sprays,
and a spray exists in as_of-25 .. as_of-6, append last_applied_treatment and a spray exists in as_of-25 .. as_of-6, append last_applied_treatment
as the last key (date + products). Otherwise that key is omitted. as the last key (date + products). Otherwise that key is omitted.
`treatment_history` is the job's already-loaded 25-day operations window;
this function does not query treatments again.
""" """
recommended = list(recommended) recommended = list(recommended)
@ -146,8 +147,7 @@ def enrich_advice_json(
advice_json, advice_json,
conn=conn, conn=conn,
builder=builder, builder=builder,
as_of=as_of, treatment_history=treatment_history,
field_id=field_id,
) )
@ -156,15 +156,14 @@ def _maybe_add_last_applied_treatment(
*, *,
conn: pyodbc.Connection, conn: pyodbc.Connection,
builder: ProductJsonBuilder, builder: ProductJsonBuilder,
as_of: date, treatment_history: TreatmentHistory,
field_id: int,
) -> None: ) -> None:
if not advice_json.get("allowed_products"): if not advice_json.get("allowed_products"):
return return
if _has_recent_applied_treatment(advice_json): if _has_recent_applied_treatment(advice_json):
return return
last = load_last_treatment_before_recent_window(conn, field_id, as_of) last = treatment_history.latest_before_recent()
if last is None: if last is None:
return return

View File

@ -4,11 +4,14 @@ from __future__ import annotations
import logging import logging
from dataclasses import dataclass from dataclasses import dataclass
from datetime import date
from typing import Iterable
import pyodbc import pyodbc
from pipeline.db import fetch_all from pipeline.db import fetch_all
from pipeline.jsonutils import normalize_name, parse_json_list from pipeline.jsonutils import normalize_name, parse_json_list
from pipeline.stages.treatments import LAST_TREATMENT_MAX_AGE_DAYS
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -62,6 +65,14 @@ SELECT
FROM products FROM products
""" """
_USES_INTERVAL_SQL = """
SELECT product_id, treatment_interval_min_days
FROM product_uses
WHERE crop = ?
AND disease = ?
AND product_id IN ({placeholders})
"""
class ProductIndex: class ProductIndex:
""" """
@ -125,6 +136,97 @@ class ProductIndex:
return self._by_name.get(normalize_name(product_name)) return self._by_name.get(normalize_name(product_name))
def _normalize_ingredient(value: str) -> str:
return value.strip().casefold()
def excluded_active_ingredients_for_last_treatment(
conn: pyodbc.Connection,
index: ProductIndex,
*,
crop_english: str,
disease_english: str,
latest_spray: tuple[date, list[str]] | None,
as_of: date,
) -> frozenset[str]:
"""
Active ingredients still locked by the latest spray's treatment interval.
For each product in `latest_spray`, look up `product_uses.treatment_interval_min_days`
for this crop/disease. A product is locked if any matching use row has a
non-NULL min interval strictly greater than days elapsed since the spray.
Only locked products contribute ingredients. Tank-mix partners are
evaluated independently. Names missing from `products` are skipped.
"""
if latest_spray is None:
return frozenset()
spray_date, product_names = latest_spray
days_elapsed = (as_of - spray_date).days
if days_elapsed < 0 or days_elapsed > LAST_TREATMENT_MAX_AGE_DAYS:
return frozenset()
resolved: list[ProductRow] = []
seen_ids: set[str] = set()
for name in product_names:
row = index.by_name(name)
if row is None:
logger.warning(
"Last applied product '%s' has no row in products; "
"it cannot contribute to the rotation gate.",
name,
)
continue
if row.product_id in seen_ids:
continue
seen_ids.add(row.product_id)
resolved.append(row)
if not resolved:
return frozenset()
product_ids = [row.product_id for row in resolved]
placeholders = ",".join("?" * len(product_ids))
use_rows = fetch_all(
conn,
_USES_INTERVAL_SQL.format(placeholders=placeholders),
(crop_english, disease_english, *product_ids),
)
locked_ids: set[str] = set()
for use_row in use_rows:
interval = use_row.treatment_interval_min_days
if interval is None:
continue
try:
min_days = int(interval)
except (TypeError, ValueError):
continue
if min_days > days_elapsed:
locked_ids.add(str(use_row.product_id))
excluded: set[str] = set()
locked_names: list[str] = []
for row in resolved:
if row.product_id not in locked_ids:
continue
locked_names.append(row.product_name)
for ingredient in row.active_ingredients:
key = _normalize_ingredient(str(ingredient))
if key:
excluded.add(key)
if locked_names:
logger.info(
"Rotation gate: spray_date=%s days_elapsed=%s locked=%s excluded_ingredients=%d",
spray_date.isoformat(),
days_elapsed,
locked_names,
len(excluded),
)
return frozenset(excluded)
def prefilter_products( def prefilter_products(
index: ProductIndex, index: ProductIndex,
*, *,
@ -132,6 +234,7 @@ def prefilter_products(
disease_english: str, disease_english: str,
organic: bool, organic: bool,
rule: str | None, rule: str | None,
excluded_active_ingredients: Iterable[str] | None = None,
) -> list[ProductCandidate]: ) -> list[ProductCandidate]:
""" """
Filter the (crop, disease) bucket by organic flag and the selected Filter the (crop, disease) bucket by organic flag and the selected
@ -148,7 +251,15 @@ def prefilter_products(
fase5_past -> systemicity in {contact, mixed} AND preventive_action = 1 fase5_past -> systemicity in {contact, mixed} AND preventive_action = 1
AND curative_action = 0 AND eradicant_action = 0 AND curative_action = 0 AND eradicant_action = 0
None -> no FASE-specific cut; crop/disease/organic gates only None -> no FASE-specific cut; crop/disease/organic gates only
When `excluded_active_ingredients` is non-empty, a candidate is also
dropped if any of its `active_ingredients` matches (case-insensitive).
""" """
blocked = {
_normalize_ingredient(str(item))
for item in (excluded_active_ingredients or [])
if item is not None and str(item).strip()
}
candidates: list[ProductCandidate] = [] candidates: list[ProductCandidate] = []
for row in index.candidates_for(crop_english, disease_english): for row in index.candidates_for(crop_english, disease_english):
@ -198,14 +309,21 @@ def prefilter_products(
continue continue
# rule is None: no FASE-specific cut, only the gates above apply. # rule is None: no FASE-specific cut, only the gates above apply.
if blocked and any(
_normalize_ingredient(str(ingredient)) in blocked
for ingredient in row.active_ingredients
):
continue
candidates.append(row.as_candidate()) candidates.append(row.as_candidate())
logger.info( logger.info(
"Product prefilter: crop=%s disease=%s organic=%s rule=%s -> %d candidates", "Product prefilter: crop=%s disease=%s organic=%s rule=%s excluded_ais=%d -> %d candidates",
crop_english, crop_english,
disease_english, disease_english,
organic, organic,
rule, rule,
len(blocked),
len(candidates), len(candidates),
) )
return candidates return candidates

View File

@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
import logging import logging
from dataclasses import dataclass, field
from datetime import date, timedelta from datetime import date, timedelta
from typing import Any, Iterable from typing import Any, Iterable
@ -16,7 +17,7 @@ logger = logging.getLogger(__name__)
_TREATMENT_OPERATION = 10 _TREATMENT_OPERATION = 10
# Meteorological applied_treatment covers as_of-RECENT .. as_of. # Meteorological applied_treatment covers as_of-RECENT .. as_of.
RECENT_TREATMENT_DAYS = 5 RECENT_TREATMENT_DAYS = 5
# last_applied_treatment looks at sprays strictly before that window, back to as_of-25. # Rotation gate and last_applied_treatment look back to as_of-25.
LAST_TREATMENT_MAX_AGE_DAYS = 25 LAST_TREATMENT_MAX_AGE_DAYS = 25
_APPLIED_TREATMENTS_SQL = """ _APPLIED_TREATMENTS_SQL = """
@ -59,79 +60,92 @@ def _group_treatments_by_day(rows: Iterable[Any]) -> dict[date, list[str]]:
return by_day return by_day
def load_applied_treatments( @dataclass
class TreatmentHistory:
"""Crop-protection sprays on one field over as_of-25 .. as_of, loaded once."""
as_of: date
by_day: dict[date, list[str]] = field(default_factory=dict)
def _window(self, start: date, end: date) -> dict[date, list[str]]:
return {day: names for day, names in self.by_day.items() if start <= day <= end}
def recent(self) -> dict[date, list[str]]:
"""Product names per day over as_of-5 .. as_of (meteorological applied_treatment)."""
start = self.as_of - timedelta(days=RECENT_TREATMENT_DAYS)
return self._window(start, self.as_of)
def latest(self) -> tuple[date, list[str]] | None:
"""Most recent spray in as_of-25 .. as_of, or None. Used by the rotation gate."""
start = self.as_of - timedelta(days=LAST_TREATMENT_MAX_AGE_DAYS)
window = self._window(start, self.as_of)
if not window:
return None
day = max(window)
return day, list(window[day])
def latest_before_recent(self) -> tuple[date, list[str]] | None:
"""Most recent spray in as_of-25 .. as_of-6, or None (last_applied_treatment)."""
start = self.as_of - timedelta(days=LAST_TREATMENT_MAX_AGE_DAYS)
end = self.as_of - timedelta(days=RECENT_TREATMENT_DAYS + 1)
window = self._window(start, end)
if not window:
return None
day = max(window)
return day, list(window[day])
def load_treatment_history(
conn: pyodbc.Connection, conn: pyodbc.Connection,
field_id: int, field_id: int,
as_of: date, as_of: date,
) -> dict[date, list[str]]: ) -> TreatmentHistory:
""" """
Return the product names applied on the field per day over as_of-5 .. as_of. Load every crop-protection spray on the field over as_of-25 .. as_of.
Days without a treatment are absent from the mapping. An empty result means Callers slice the result: recent() for meteorological applied_treatment,
no operation matched the field, the date window, or the treatment operation latest() for the rotation gate, latest_before_recent() for
type, in which case every applied_treatment stays empty downstream. last_applied_treatment. Days without a treatment are absent from by_day.
""" """
start = as_of - timedelta(days=RECENT_TREATMENT_DAYS) start = as_of - timedelta(days=LAST_TREATMENT_MAX_AGE_DAYS)
rows = fetch_all( rows = fetch_all(
conn, conn,
_APPLIED_TREATMENTS_SQL, _APPLIED_TREATMENTS_SQL,
(field_id, start, as_of, _TREATMENT_OPERATION), (field_id, start, as_of, _TREATMENT_OPERATION),
) )
by_day = _group_treatments_by_day(rows)
history = TreatmentHistory(as_of=as_of, by_day=by_day)
if not rows: recent = history.recent()
if not recent:
logger.warning( logger.warning(
"No crop protection operations for field %s between %s and %s; " "No crop protection operations for field %s between %s and %s; "
"applied_treatment stays empty.", "applied_treatment stays empty.",
field_id, field_id,
start.isoformat(), (as_of - timedelta(days=RECENT_TREATMENT_DAYS)).isoformat(),
as_of.isoformat(), as_of.isoformat(),
) )
return {} else:
by_day = _group_treatments_by_day(rows)
logger.info( logger.info(
"Applied treatments for field %s: %d day(s), %d product mention(s)", "Applied treatments for field %s: %d day(s), %d product mention(s)",
field_id, field_id,
len(by_day), len(recent),
sum(len(v) for v in by_day.values()), sum(len(names) for names in recent.values()),
) )
return by_day
latest = history.latest()
def load_last_treatment_before_recent_window( if latest is None:
conn: pyodbc.Connection,
field_id: int,
as_of: date,
) -> tuple[date, list[str]] | None:
"""
Return the most recent spray in as_of-25 .. as_of-6, or None.
The meteorological window already covers as_of-5 .. as_of; this lookback
is only used for last_applied_treatment when that recent window is empty.
"""
start = as_of - timedelta(days=LAST_TREATMENT_MAX_AGE_DAYS)
end = as_of - timedelta(days=RECENT_TREATMENT_DAYS + 1)
rows = fetch_all(
conn,
_APPLIED_TREATMENTS_SQL,
(field_id, start, end, _TREATMENT_OPERATION),
)
by_day = _group_treatments_by_day(rows)
if not by_day:
logger.info( logger.info(
"No crop protection operations for field %s between %s and %s.", "No crop protection operations for field %s between %s and %s.",
field_id, field_id,
start.isoformat(), start.isoformat(),
end.isoformat(), as_of.isoformat(),
) )
return None else:
latest = max(by_day)
names = by_day[latest]
logger.info( logger.info(
"Last applied treatment for field %s on %s: %s", "Latest treatment for field %s on %s: %s",
field_id, field_id,
latest.isoformat(), latest[0].isoformat(),
names, latest[1],
) )
return latest, names return history

View File

@ -14,13 +14,6 @@ logger = logging.getLogger(__name__)
COLLECTION = "ProductProfile" COLLECTION = "ProductProfile"
# Above this many allowlisted product_ids, encoding every one of them as an
# OR filter term gets expensive to build and transmit. Past that size the
# allowlist is a large fraction of the whole catalog anyway, so falling back
# to a global over-fetch + Python filter (the original approach) is cheaper
# and just as correct.
_MAX_FILTERED_IDS = 200
@dataclass(frozen=True) @dataclass(frozen=True)
class RecommendedProduct: class RecommendedProduct:
@ -78,15 +71,10 @@ def search_products(
`pipeline.resources.Resources`); this function does not open or close a `pipeline.resources.Resources`); this function does not open or close a
connection itself, so it is safe to call from many jobs in a row. connection itself, so it is safe to call from many jobs in a row.
The allowlist is applied as a `product_id` OR filter before ranking, The SQL allowlist is always applied as a `product_id` `contains_any`
rather than over-fetching globally and intersecting in Python afterwards. filter before ranking, so `limit` is the nearest neighbours among
The original code fetched a flat `limit=30` and only worked because the allowed products even when the shortlist is large. This assumes
whole `ProductProfile` collection had ~30 objects, so the top 30 by raw `product_id` is filterable with `tokenization: field` (exact ID match).
similarity were, by construction, the entire catalog. Once the catalog
grows past that, a selective allowlist (organic + crop + disease + FASE
can easily narrow it to a handful of products) would often be missing
entirely from an unfiltered top-30, silently returning zero or too few
recommendations instead of raising an error.
""" """
if not candidates: if not candidates:
logger.warning("No product candidates after prefilter; skipping vector search.") logger.warning("No product candidates after prefilter; skipping vector search.")
@ -98,22 +86,7 @@ def search_products(
allowlist = {c.product_id: c for c in candidates} allowlist = {c.product_id: c for c in candidates}
per_query_limit = 4 if len(queries) == 1 else 2 per_query_limit = 4 if len(queries) == 1 else 2
query_filter = Filter.by_property("product_id").contains_any(list(allowlist))
query_filter = None
unfiltered_overfetch = False
if len(allowlist) <= _MAX_FILTERED_IDS:
query_filter = Filter.any_of(
[Filter.by_property("product_id").equal(pid) for pid in allowlist]
)
else:
unfiltered_overfetch = True
logger.warning(
"Allowlist has %d product(s), above the %d-id server-side filter "
"threshold; falling back to a global over-fetch + Python filter "
"for this query.",
len(allowlist),
_MAX_FILTERED_IDS,
)
collection = client.collections.get(COLLECTION) collection = client.collections.get(COLLECTION)
results: list[RecommendedProduct] = [] results: list[RecommendedProduct] = []
@ -121,8 +94,6 @@ def search_products(
for query_index, query in enumerate(queries): for query_index, query in enumerate(queries):
fetch_limit = per_query_limit * (query_index + 1) fetch_limit = per_query_limit * (query_index + 1)
if unfiltered_overfetch:
fetch_limit = min(200, max(fetch_limit, per_query_limit * 5, 15))
response = collection.query.near_text( response = collection.query.near_text(
query=query, query=query,
limit=fetch_limit, limit=fetch_limit,

View File

@ -4,15 +4,13 @@ from __future__ import annotations
import unittest import unittest
from datetime import date from datetime import date
from types import SimpleNamespace from unittest.mock import MagicMock, patch
from unittest.mock import ANY, MagicMock, patch
from pipeline.stages.advice_context import collapse_product_json, enrich_advice_json from pipeline.stages.advice_context import collapse_product_json, enrich_advice_json
from pipeline.stages.treatments import load_last_treatment_before_recent_window from pipeline.stages.treatments import TreatmentHistory
from pipeline.stages.vector import RecommendedProduct from pipeline.stages.vector import RecommendedProduct
_AS_OF = date(2026, 4, 16) _AS_OF = date(2026, 4, 16)
_FIELD_ID = 7216
_SPRAY_DATE = date(2026, 4, 10) _SPRAY_DATE = date(2026, 4, 10)
@ -41,12 +39,18 @@ def _payload(*, applied: list[str] | None = None) -> dict:
} }
def _history(*product_names: str, spray_date: date = _SPRAY_DATE) -> TreatmentHistory:
by_day = {spray_date: list(product_names)} if product_names else {}
return TreatmentHistory(as_of=_AS_OF, by_day=by_day)
def _enrich( def _enrich(
advice_json: dict, advice_json: dict,
*, *,
recommended: list[RecommendedProduct] | None = None, recommended: list[RecommendedProduct] | None = None,
last_advice_treatments: list[str] | None = None, last_advice_treatments: list[str] | None = None,
last_advice_date: date | None = None, last_advice_date: date | None = None,
treatment_history: TreatmentHistory | None = None,
) -> None: ) -> None:
enrich_advice_json( enrich_advice_json(
advice_json, advice_json,
@ -55,49 +59,53 @@ def _enrich(
recommended=recommended or [], recommended=recommended or [],
last_advice_summary="", last_advice_summary="",
last_advice_treatments=last_advice_treatments or [], last_advice_treatments=last_advice_treatments or [],
as_of=_AS_OF,
field_id=_FIELD_ID,
last_advice_date=last_advice_date, last_advice_date=last_advice_date,
treatment_history=treatment_history if treatment_history is not None else _history(),
) )
class LastAppliedTreatmentTests(unittest.TestCase): class LastAppliedTreatmentTests(unittest.TestCase):
def test_omits_when_allowed_products_empty_and_does_not_query(self) -> None: def test_omits_when_allowed_products_empty_and_does_not_query(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( with patch("pipeline.stages.treatments.fetch_all") as fetch:
"pipeline.stages.advice_context.load_last_treatment_before_recent_window" _enrich(
) as loader: advice_json,
_enrich(advice_json, recommended=[]) recommended=[],
loader.assert_not_called() treatment_history=_history("LIETO"),
)
fetch.assert_not_called()
self.assertNotIn("last_applied_treatment", advice_json) self.assertNotIn("last_applied_treatment", advice_json)
self.assertEqual(advice_json["allowed_products"], []) self.assertEqual(advice_json["allowed_products"], [])
def test_omits_when_recent_applied_treatment_present(self) -> None: def test_omits_when_recent_applied_treatment_present(self) -> None:
advice_json = _payload(applied=["LIETO"]) advice_json = _payload(applied=["LIETO"])
with patch( with patch("pipeline.stages.treatments.fetch_all") as fetch:
"pipeline.stages.advice_context.load_last_treatment_before_recent_window" _enrich(
) as loader: advice_json,
_enrich(advice_json, recommended=[_recommended("LIETO")]) recommended=[_recommended("LIETO")],
loader.assert_not_called() treatment_history=_history("LIETO"),
)
fetch.assert_not_called()
self.assertNotIn("last_applied_treatment", advice_json) self.assertNotIn("last_applied_treatment", advice_json)
def test_omits_when_loader_returns_none(self) -> None: def test_omits_when_history_has_no_older_spray(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( with patch("pipeline.stages.treatments.fetch_all") as fetch:
"pipeline.stages.advice_context.load_last_treatment_before_recent_window", _enrich(
return_value=None, advice_json,
) as loader: recommended=[_recommended("LIETO")],
_enrich(advice_json, recommended=[_recommended("LIETO")]) treatment_history=TreatmentHistory(as_of=_AS_OF, by_day={}),
loader.assert_called_once_with(ANY, _FIELD_ID, _AS_OF) )
fetch.assert_not_called()
self.assertNotIn("last_applied_treatment", advice_json) self.assertNotIn("last_applied_treatment", advice_json)
def test_includes_at_end_when_all_gates_pass(self) -> None: def test_includes_at_end_when_all_gates_pass(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( _enrich(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window", advice_json,
return_value=(_SPRAY_DATE, ["LIETO"]), recommended=[_recommended("LIETO")],
): treatment_history=_history("LIETO"),
_enrich(advice_json, recommended=[_recommended("LIETO")]) )
self.assertEqual(list(advice_json)[-1], "last_applied_treatment") self.assertEqual(list(advice_json)[-1], "last_applied_treatment")
self.assertEqual( self.assertEqual(
advice_json["last_applied_treatment"], advice_json["last_applied_treatment"],
@ -106,37 +114,31 @@ class LastAppliedTreatmentTests(unittest.TestCase):
def test_name_only_when_product_in_allowed_products(self) -> None: def test_name_only_when_product_in_allowed_products(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( _enrich(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window", advice_json,
return_value=(_SPRAY_DATE, ["LIETO"]), recommended=[_recommended("LIETO")],
): treatment_history=_history("LIETO"),
_enrich(advice_json, recommended=[_recommended("LIETO")]) )
self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"]) self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"])
def test_name_only_when_product_in_suggested_products_as_string(self) -> None: def test_name_only_when_product_in_suggested_products_as_string(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window",
return_value=(_SPRAY_DATE, ["LIETO"]),
):
_enrich( _enrich(
advice_json, advice_json,
recommended=[_recommended("LIETO")], recommended=[_recommended("LIETO")],
last_advice_treatments=["LIETO"], last_advice_treatments=["LIETO"],
treatment_history=_history("LIETO"),
) )
self.assertEqual(advice_json["last_advice"]["suggested_products"], ["LIETO"]) self.assertEqual(advice_json["last_advice"]["suggested_products"], ["LIETO"])
self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"]) self.assertEqual(advice_json["last_applied_treatment"]["products"], ["LIETO"])
def test_name_only_when_product_in_suggested_products_as_object(self) -> None: def test_name_only_when_product_in_suggested_products_as_object(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window",
return_value=(_SPRAY_DATE, ["TWINGO"]),
):
_enrich( _enrich(
advice_json, advice_json,
recommended=[_recommended("LIETO")], recommended=[_recommended("LIETO")],
last_advice_treatments=["TWINGO"], last_advice_treatments=["TWINGO"],
treatment_history=_history("TWINGO"),
) )
self.assertEqual( self.assertEqual(
advice_json["last_advice"]["suggested_products"], advice_json["last_advice"]["suggested_products"],
@ -146,11 +148,11 @@ class LastAppliedTreatmentTests(unittest.TestCase):
def test_full_json_when_product_in_neither_list(self) -> None: def test_full_json_when_product_in_neither_list(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( _enrich(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window", advice_json,
return_value=(_SPRAY_DATE, ["UNKNOWN"]), recommended=[_recommended("LIETO")],
): treatment_history=_history("UNKNOWN"),
_enrich(advice_json, recommended=[_recommended("LIETO")]) )
self.assertEqual( self.assertEqual(
advice_json["last_applied_treatment"]["products"], advice_json["last_applied_treatment"]["products"],
[{"UNKNOWN": {"uses": []}}], [{"UNKNOWN": {"uses": []}}],
@ -158,16 +160,27 @@ class LastAppliedTreatmentTests(unittest.TestCase):
def test_tank_mix_mixes_name_and_full_json(self) -> None: def test_tank_mix_mixes_name_and_full_json(self) -> None:
advice_json = _payload() advice_json = _payload()
with patch( _enrich(
"pipeline.stages.advice_context.load_last_treatment_before_recent_window", advice_json,
return_value=(_SPRAY_DATE, ["LIETO", "UNKNOWN"]), recommended=[_recommended("LIETO")],
): treatment_history=_history("LIETO", "UNKNOWN"),
_enrich(advice_json, recommended=[_recommended("LIETO")]) )
self.assertEqual( self.assertEqual(
advice_json["last_applied_treatment"]["products"], advice_json["last_applied_treatment"]["products"],
["LIETO", {"UNKNOWN": {"uses": []}}], ["LIETO", {"UNKNOWN": {"uses": []}}],
) )
def test_does_not_query_operations_when_including_last_applied_treatment(self) -> None:
advice_json = _payload()
with patch("pipeline.stages.treatments.fetch_all") as fetch:
_enrich(
advice_json,
recommended=[_recommended("LIETO")],
treatment_history=_history("LIETO"),
)
fetch.assert_not_called()
self.assertIn("last_applied_treatment", advice_json)
class CollapseProductJsonTests(unittest.TestCase): class CollapseProductJsonTests(unittest.TestCase):
def test_collapses_last_applied_treatment_products(self) -> None: def test_collapses_last_applied_treatment_products(self) -> None:
@ -193,29 +206,5 @@ class CollapseProductJsonTests(unittest.TestCase):
) )
class LoadLastTreatmentBeforeRecentWindowTests(unittest.TestCase):
def test_returns_latest_day_products(self) -> None:
rows = [
SimpleNamespace(rilop_date=date(2026, 3, 22), prodgov_name="OLD"),
SimpleNamespace(rilop_date=date(2026, 4, 10), prodgov_name="LIETO"),
SimpleNamespace(rilop_date=date(2026, 4, 10), prodgov_name="TWINGO"),
]
with patch("pipeline.stages.treatments.fetch_all", return_value=rows) as fetch:
result = load_last_treatment_before_recent_window(
MagicMock(), _FIELD_ID, _AS_OF
)
params = fetch.call_args[0][2]
self.assertEqual(params[1], date(2026, 3, 22)) # as_of-25
self.assertEqual(params[2], date(2026, 4, 10)) # as_of-6
self.assertEqual(result, (_SPRAY_DATE, ["LIETO", "TWINGO"]))
def test_returns_none_when_no_rows(self) -> None:
with patch("pipeline.stages.treatments.fetch_all", return_value=[]):
result = load_last_treatment_before_recent_window(
MagicMock(), _FIELD_ID, _AS_OF
)
self.assertIsNone(result)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

225
tests/test_products.py Normal file
View File

@ -0,0 +1,225 @@
"""Unit tests for the treatment-interval active-ingredient rotation gate."""
from __future__ import annotations
import unittest
from datetime import date
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from pipeline.stages.products import (
ProductIndex,
ProductRow,
excluded_active_ingredients_for_last_treatment,
prefilter_products,
)
_AS_OF = date(2026, 4, 16)
_SPRAY_DATE = date(2026, 4, 6) # 10 days before as_of
_CROP = "grapevine"
_DISEASE = "downy mildew"
def _product(
*,
product_id: str = "p1",
product_name: str = "LIETO",
active_ingredients: list[str] | None = None,
systemicity: str = "contact",
preventive_action: bool = True,
curative_action: bool = False,
eradicant_action: bool = False,
organic_certified: bool = False,
) -> ProductRow:
return ProductRow(
product_id=product_id,
product_name=product_name,
registration_number=None,
organic_certified=organic_certified,
systemicity=systemicity,
preventive_action=preventive_action,
curative_action=curative_action,
eradicant_action=eradicant_action,
target_crops=[_CROP],
target_diseases=[_DISEASE],
active_ingredients=active_ingredients or ["mancozeb"],
frac_groups=[],
)
def _excluded(
index: ProductIndex,
*,
names: list[str],
uses: list[SimpleNamespace],
spray_date: date = _SPRAY_DATE,
as_of: date = _AS_OF,
) -> frozenset[str]:
with patch("pipeline.stages.products.fetch_all", return_value=uses) as fetch:
result = excluded_active_ingredients_for_last_treatment(
MagicMock(),
index,
crop_english=_CROP,
disease_english=_DISEASE,
latest_spray=(spray_date, names),
as_of=as_of,
)
fetch.assert_called_once()
return result
class ExcludedActiveIngredientsTests(unittest.TestCase):
def test_none_spray_returns_empty_and_does_not_query(self) -> None:
index = ProductIndex([_product()])
with patch("pipeline.stages.products.fetch_all") as fetch:
result = excluded_active_ingredients_for_last_treatment(
MagicMock(),
index,
crop_english=_CROP,
disease_english=_DISEASE,
latest_spray=None,
as_of=_AS_OF,
)
fetch.assert_not_called()
self.assertEqual(result, frozenset())
def test_null_interval_returns_empty(self) -> None:
index = ProductIndex([_product()])
result = _excluded(
index,
names=["LIETO"],
uses=[SimpleNamespace(product_id="p1", treatment_interval_min_days=None)],
)
self.assertEqual(result, frozenset())
def test_interval_less_or_equal_days_elapsed_does_not_lock(self) -> None:
index = ProductIndex([_product()])
equal = _excluded(
index,
names=["LIETO"],
uses=[SimpleNamespace(product_id="p1", treatment_interval_min_days=10)],
)
shorter = _excluded(
index,
names=["LIETO"],
uses=[SimpleNamespace(product_id="p1", treatment_interval_min_days=7)],
)
self.assertEqual(equal, frozenset())
self.assertEqual(shorter, frozenset())
def test_interval_greater_than_days_elapsed_locks_ingredients(self) -> None:
index = ProductIndex([_product(active_ingredients=["Mancozeb", "copper"])])
result = _excluded(
index,
names=["LIETO"],
uses=[SimpleNamespace(product_id="p1", treatment_interval_min_days=14)],
)
self.assertEqual(result, frozenset({"mancozeb", "copper"}))
def test_any_qualifying_use_row_locks_the_product(self) -> None:
index = ProductIndex([_product()])
result = _excluded(
index,
names=["LIETO"],
uses=[
SimpleNamespace(product_id="p1", treatment_interval_min_days=None),
SimpleNamespace(product_id="p1", treatment_interval_min_days=21),
],
)
self.assertEqual(result, frozenset({"mancozeb"}))
def test_tank_mix_only_locked_partner_contributes_ingredients(self) -> None:
index = ProductIndex(
[
_product(
product_id="p1",
product_name="LIETO",
active_ingredients=["mancozeb"],
),
_product(
product_id="p2",
product_name="TWINGO",
active_ingredients=["metalaxyl"],
),
]
)
result = _excluded(
index,
names=["LIETO", "TWINGO"],
uses=[
SimpleNamespace(product_id="p1", treatment_interval_min_days=14),
SimpleNamespace(product_id="p2", treatment_interval_min_days=None),
],
)
self.assertEqual(result, frozenset({"mancozeb"}))
def test_catalog_miss_is_ignored_other_partner_still_locks(self) -> None:
index = ProductIndex([_product()])
result = _excluded(
index,
names=["UNKNOWN", "LIETO"],
uses=[SimpleNamespace(product_id="p1", treatment_interval_min_days=14)],
)
self.assertEqual(result, frozenset({"mancozeb"}))
def test_catalog_miss_only_does_not_query_uses(self) -> None:
index = ProductIndex([_product()])
with patch("pipeline.stages.products.fetch_all") as fetch:
result = excluded_active_ingredients_for_last_treatment(
MagicMock(),
index,
crop_english=_CROP,
disease_english=_DISEASE,
latest_spray=(_SPRAY_DATE, ["UNKNOWN"]),
as_of=_AS_OF,
)
fetch.assert_not_called()
self.assertEqual(result, frozenset())
class PrefilterRotationGateTests(unittest.TestCase):
def test_drops_candidates_sharing_an_excluded_ingredient(self) -> None:
index = ProductIndex(
[
_product(
product_id="a",
product_name="A",
active_ingredients=["mancozeb"],
),
_product(
product_id="b",
product_name="B",
active_ingredients=["copper hydroxide"],
),
]
)
candidates = prefilter_products(
index,
crop_english=_CROP,
disease_english=_DISEASE,
organic=False,
rule=None,
excluded_active_ingredients={"MANCOZEB"},
)
self.assertEqual([item.product_name for item in candidates], ["B"])
def test_empty_exclusion_set_leaves_candidates_unchanged(self) -> None:
index = ProductIndex(
[
_product(product_id="a", product_name="A"),
_product(product_id="b", product_name="B", active_ingredients=["copper"]),
]
)
candidates = prefilter_products(
index,
crop_english=_CROP,
disease_english=_DISEASE,
organic=False,
rule=None,
excluded_active_ingredients=frozenset(),
)
self.assertEqual([item.product_name for item in candidates], ["A", "B"])
if __name__ == "__main__":
unittest.main()

96
tests/test_treatments.py Normal file
View File

@ -0,0 +1,96 @@
"""Unit tests for TreatmentHistory window slicing and the single 25-day load."""
from __future__ import annotations
import unittest
from datetime import date
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
from pipeline.stages.treatments import TreatmentHistory, load_treatment_history
_AS_OF = date(2026, 4, 16)
_FIELD_ID = 7216
_SPRAY_DATE = date(2026, 4, 10) # as_of-6
class TreatmentHistorySliceTests(unittest.TestCase):
def test_latest_picks_newest_day_and_keeps_tank_mix_order(self) -> None:
history = TreatmentHistory(
as_of=_AS_OF,
by_day={
date(2026, 3, 22): ["OLD"],
_SPRAY_DATE: ["LIETO", "TWINGO"],
},
)
self.assertEqual(history.latest(), (_SPRAY_DATE, ["LIETO", "TWINGO"]))
def test_latest_prefers_a_spray_inside_the_recent_window(self) -> None:
history = TreatmentHistory(
as_of=_AS_OF,
by_day={
_SPRAY_DATE: ["OLD"],
_AS_OF: ["TODAY"],
},
)
self.assertEqual(history.latest(), (_AS_OF, ["TODAY"]))
self.assertEqual(history.latest_before_recent(), (_SPRAY_DATE, ["OLD"]))
def test_recent_drops_days_older_than_five(self) -> None:
history = TreatmentHistory(
as_of=_AS_OF,
by_day={
_SPRAY_DATE: ["OLD"],
date(2026, 4, 11): ["RECENT"],
_AS_OF: ["TODAY"],
},
)
self.assertEqual(
history.recent(),
{date(2026, 4, 11): ["RECENT"], _AS_OF: ["TODAY"]},
)
def test_latest_before_recent_matches_as_of_minus_25_to_minus_6(self) -> None:
history = TreatmentHistory(
as_of=_AS_OF,
by_day={
date(2026, 3, 22): ["OLD"],
_SPRAY_DATE: ["LIETO"],
date(2026, 4, 11): ["RECENT"],
},
)
self.assertEqual(history.latest_before_recent(), (_SPRAY_DATE, ["LIETO"]))
def test_latest_and_latest_before_recent_return_none_when_empty(self) -> None:
history = TreatmentHistory(as_of=_AS_OF, by_day={})
self.assertIsNone(history.latest())
self.assertIsNone(history.latest_before_recent())
self.assertEqual(history.recent(), {})
class LoadTreatmentHistoryTests(unittest.TestCase):
def test_queries_as_of_minus_25_through_as_of(self) -> None:
rows = [
SimpleNamespace(rilop_date=date(2026, 3, 22), prodgov_name="OLD"),
SimpleNamespace(rilop_date=_SPRAY_DATE, prodgov_name="LIETO"),
SimpleNamespace(rilop_date=_SPRAY_DATE, prodgov_name="TWINGO"),
SimpleNamespace(rilop_date=_SPRAY_DATE, prodgov_name="LIETO"),
]
with patch("pipeline.stages.treatments.fetch_all", return_value=rows) as fetch:
history = load_treatment_history(MagicMock(), _FIELD_ID, _AS_OF)
params = fetch.call_args[0][2]
self.assertEqual(params[1], date(2026, 3, 22)) # as_of-25
self.assertEqual(params[2], _AS_OF)
self.assertEqual(history.latest(), (_SPRAY_DATE, ["LIETO", "TWINGO"]))
self.assertEqual(history.latest_before_recent(), (_SPRAY_DATE, ["LIETO", "TWINGO"]))
self.assertEqual(history.recent(), {})
def test_returns_empty_history_when_no_rows(self) -> None:
with patch("pipeline.stages.treatments.fetch_all", return_value=[]):
history = load_treatment_history(MagicMock(), _FIELD_ID, _AS_OF)
self.assertEqual(history.by_day, {})
self.assertIsNone(history.latest())
if __name__ == "__main__":
unittest.main()

View File

@ -38,7 +38,7 @@ class FakeNearTextQuery:
return_metadata: Any, return_metadata: Any,
return_properties: list[str], return_properties: list[str],
) -> SimpleNamespace: ) -> SimpleNamespace:
self.calls.append({"query": query, "limit": limit}) self.calls.append({"query": query, "limit": limit, "filters": filters})
hits = list(self.hits_by_query.get(query, [])) hits = list(self.hits_by_query.get(query, []))
return SimpleNamespace(objects=hits[:limit]) return SimpleNamespace(objects=hits[:limit])
@ -61,7 +61,11 @@ class SearchProductsLimitTests(unittest.TestCase):
self.assertEqual([item.product_id for item in results], ["p1", "p2", "p3", "p4"]) 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([item.query_index for item in results], [0, 0, 0, 0])
self.assertEqual(client.query.calls, [{"query": "q1", "limit": 4}]) 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: def test_two_queries_keep_two_each(self) -> None:
candidates = [_candidate(f"p{i}") for i in range(1, 7)] candidates = [_candidate(f"p{i}") for i in range(1, 7)]
@ -79,9 +83,10 @@ class SearchProductsLimitTests(unittest.TestCase):
[("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)], [("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)],
) )
self.assertEqual( self.assertEqual(
client.query.calls, [(c["query"], c["limit"]) for c in client.query.calls],
[{"query": "q1", "limit": 2}, {"query": "q2", "limit": 4}], [("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: def test_overlapping_hits_keep_first_occurrence(self) -> None:
candidates = [_candidate(f"p{i}") for i in range(1, 6)] candidates = [_candidate(f"p{i}") for i in range(1, 6)]
@ -99,9 +104,22 @@ class SearchProductsLimitTests(unittest.TestCase):
[("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)], [("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)],
) )
self.assertEqual( self.assertEqual(
client.query.calls, [(c["query"], c["limit"]) for c in client.query.calls],
[{"query": "q1", "limit": 2}, {"query": "q2", "limit": 4}], [("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__": if __name__ == "__main__":