Update LLM configuration and enhance vector search capabilities

- Changed `thinking_budget` to `thinking_level` in `config.yaml` for Gemini models, updating the logic to reflect the new parameter.
- Updated LLM models in `config.yaml` to `gemini-3.8-flash` and `gemini-3.1-pro-preview`, adjusting their respective `max_tokens` and `thinking_level`.
- Enhanced `docker-compose.yml` to include a new `reranker-transformers` service for improved product ranking.
- Modified vector search logic in `vector.py` to incorporate reranking with the local transformers model, improving the accuracy of product recommendations.
- Updated unit tests to validate the new reranking functionality and ensure correct behavior in product searches.

README.md

@@ -310,21 +310,20 @@ collector cannot block a worker or push the run past `schedule.deadline`.
 Every `(field, crop, disease, as_of)` job produces one trace, rooted at a
 `run_job` span, with:
 
-- a `query_synthesis` span (gemini-2.5-flash by default) and an
+- a `query_synthesis` span (gemini-3.8-flash by default) and an
-  `advice_generation` span (gemini-2.5-pro), each wrapping the provider's own
+  `advice_generation` span (gemini-3.1-pro-preview), each wrapping the provider's own
   auto-instrumented LLM span with `llm.model_name`, `llm.provider`, and
   token counts (`llm.token_count.prompt` / `completion` / `total`, plus
-  `completion_details.reasoning` for Gemini 2.5 thinking tokens);
+  `completion_details.reasoning` for Gemini thinking tokens);
 - a `search_products` retriever span with the synthesized queries and the
   retrieved product IDs/distances (never product names or label content);
 - `dry_run` jobs get only the `run_job` span — no LLM or retrieval children,
   since `--dry-run` skips those calls entirely.
 
-Phoenix computes cost per span from its built-in model pricing table, which
+Phoenix computes cost per span from its built-in model pricing table. If a
-already includes `gemini-2.5-flash` and `gemini-2.5-pro` — no custom entry
+configured model (`gemini-3.8-flash` or `gemini-3.1-pro-preview`) is not
-under **Settings → Models** is needed for either model as configured today.
+recognised, add it under **Settings → Models** (regex `Name Pattern`,
-If a model is ever renamed or swapped to one Phoenix doesn't recognise, add
+provider `google`, per-1M-token prices).
-it there (regex `Name Pattern`, provider `google`, per-1M-token prices).
 
 Trace and span inputs are hidden by default (`observability.hide_prompts`):
 the full `json_for_advice_generation` payload (weather, phenology, product
@@ -383,7 +382,7 @@ Run once per `(field, disease)` job by `pipeline/job.py::run_job`:
 
    This stage is skipped (along with query synthesis and vector search) when the latest spray's min treatment interval is still open; see stage 6.
 
-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)
+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`), then reranked locally with `reranker-transformers` (`ms-marco-MiniLM-L-6-v2`) on `situation_fit_summary`. One synthesized query retrieves 10 and keeps the top 4 in rerank order; two queries retrieve 6 each and keep the top 2 unique each (a product already kept for the first query is skipped so the second can fill its two slots from the rest of its reranked six)
 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 `advisories` 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`, including when `allowed_products` is empty. 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

config.yaml

@@ -2,9 +2,9 @@
 # Two independent profiles, one per pipeline stage. Each accepts the same keys:
 #   provider: openai | anthropic | gemini
 #   openai_model / anthropic_model / gemini_model / max_tokens / temperature
-#   thinking_budget (optional, Gemini only): reserved thinking tokens. When set,
+#   thinking_level (optional, Gemini 3 only): low | medium | high.
-#   Gemini's max_output_tokens is max_tokens + thinking_budget, so max_tokens
+#   Gemini counts thinking tokens inside max_tokens, so max_tokens is the
-#   is the visible-output budget. Omit it to keep max_tokens as the shared cap.
+#   combined thinking-plus-output cap. Omit thinking_level to use the model's default.
 llm:
   # Part_One: vector-search query synthesis (small structured JSON in,
   # 1-2 short queries out). A lighter/cheaper model is sufficient here.
@@ -12,7 +12,7 @@ llm:
     provider: gemini
     openai_model: gpt-4o-mini
     anthropic_model: claude-haiku-4-5
-    gemini_model: gemini-2.5-flash
+    gemini_model: gemini-3.8-flash
     max_tokens: 2048
     temperature: 0.0
 
@@ -24,9 +24,9 @@ llm:
     provider: gemini
     openai_model: gpt-4o
     anthropic_model: claude-opus-4-5
-    gemini_model: gemini-2.5-pro
+    gemini_model: gemini-3.1-pro-preview
-    max_tokens: 5120
+    max_tokens: 13312
-    thinking_budget: 8192
+    thinking_level: medium
     temperature: 0.0
 
 # ── Crop-first worklist ────────────────────────────────────────────────────────

docker-compose.yml

@@ -12,15 +12,24 @@ services:
       - 8080:8080
       - 50051:50051
     volumes:
-      - weaviate_data:/var/lib/weaviate
+      - product_profile_ingestion_weaviate_data:/var/lib/weaviate
     restart: on-failure:0
+    depends_on:
+      - reranker-transformers
     environment:
       QUERY_DEFAULTS_LIMIT: 25
       PERSISTENCE_DATA_PATH: "/var/lib/weaviate"
       CLUSTER_HOSTNAME: "node1"
       # Swap OpenAI to Google modules
-      ENABLE_MODULES: "text2vec-google"
+      ENABLE_MODULES: "text2vec-google,reranker-transformers"
       DEFAULT_VECTORIZER_MODULE: "text2vec-google"
+      RERANKER_INFERENCE_API: "http://reranker-transformers:8080"
+
+  # Local cross-encoder. Reranks the near_text candidate pool; no external API key.
+  reranker-transformers:
+    image: cr.weaviate.io/semitechnologies/reranker-transformers:cross-encoder-ms-marco-MiniLM-L-6-v2
+    environment:
+      ENABLE_CUDA: 0
 
   # Self-hosted tracing/eval UI (Arize Phoenix). Opt-in from the pipeline's
   # side (config.yaml observability.enabled / PHOENIX_TRACING_ENABLED); this
@@ -38,5 +47,6 @@ services:
     restart: unless-stopped
 
 volumes:
-  weaviate_data:
+  product_profile_ingestion_weaviate_data:
+    external: true
   phoenix_data:

pipeline/config.py

@@ -20,7 +20,7 @@ class LlmSettings:
     gemini_model: str
     max_tokens: int
     temperature: float
-    thinking_budget: int | None
+    thinking_level: str | None
 
 
 @dataclass(frozen=True)
@@ -118,9 +118,19 @@ class Settings:
     output_dir: Path
 
 
+_THINKING_LEVELS = frozenset({"low", "medium", "high"})
+
+
 def _parse_llm_settings(raw: dict, defaults: dict) -> LlmSettings:
     """Build an LlmSettings profile from a config sub-block, falling back to defaults."""
-    thinking_raw = raw["thinking_budget"] if "thinking_budget" in raw else defaults.get("thinking_budget")
+    thinking_raw = raw["thinking_level"] if "thinking_level" in raw else defaults.get("thinking_level")
+    if thinking_raw is None or str(thinking_raw).strip() == "":
+        thinking_level = None
+    else:
+        thinking_level = str(thinking_raw).strip().lower()
+        if thinking_level not in _THINKING_LEVELS:
+            allowed = ", ".join(sorted(_THINKING_LEVELS))
+            raise ValueError(f"thinking_level must be one of {allowed}, got {thinking_raw!r}")
     return LlmSettings(
         provider=str(raw.get("provider", defaults["provider"])).strip().lower(),
         openai_model=str(raw.get("openai_model", defaults["openai_model"])),
@@ -128,7 +138,7 @@ def _parse_llm_settings(raw: dict, defaults: dict) -> LlmSettings:
         gemini_model=str(raw.get("gemini_model", defaults["gemini_model"])),
         max_tokens=int(raw.get("max_tokens", defaults["max_tokens"])),
         temperature=float(raw.get("temperature", defaults["temperature"])),
-        thinking_budget=None if thinking_raw is None else int(thinking_raw),
+        thinking_level=thinking_level,
     )
 
 
@@ -235,19 +245,19 @@ def load_settings(
         "provider": "gemini",
         "openai_model": "gpt-4o-mini",
         "anthropic_model": "claude-haiku-4-5",
-        "gemini_model": "gemini-2.5-flash",
+        "gemini_model": "gemini-3.8-flash",
         "max_tokens": 8192,
         "temperature": 0.0,
-        "thinking_budget": None,
+        "thinking_level": None,
     }
     advice_generation_defaults = {
         "provider": "gemini",
         "openai_model": "gpt-4o",
         "anthropic_model": "claude-opus-4-5",
-        "gemini_model": "gemini-2.5-pro",
+        "gemini_model": "gemini-3.1-pro-preview",
         "max_tokens": 8192,
         "temperature": 0.0,
-        "thinking_budget": 8192,
+        "thinking_level": "medium",
     }
 
     return Settings(

pipeline/stages/llm.py

@@ -102,10 +102,10 @@ def call_llm(
     provider = llm_cfg.provider
 
     logger.info(
-        "Calling LLM provider=%s model=%s thinking_budget=%s",
+        "Calling LLM provider=%s model=%s thinking_level=%s",
         provider,
         _model_name(llm_cfg),
-        llm_cfg.thinking_budget,
+        llm_cfg.thinking_level,
     )
     logger.debug(
         "LLM REQUEST (system_prompt + JSON user content):\n--- SYSTEM PROMPT ---\n%s\n--- USER JSON ---\n%s",
@@ -145,11 +145,9 @@ def _call_gemini(
     from google import genai
     from google.genai import types
 
-    max_output_tokens = llm_cfg.max_tokens
     thinking_config = None
-    if llm_cfg.thinking_budget is not None:
+    if llm_cfg.thinking_level is not None:
-        max_output_tokens = llm_cfg.max_tokens + llm_cfg.thinking_budget
+        thinking_config = types.ThinkingConfig(thinking_level=llm_cfg.thinking_level)
-        thinking_config = types.ThinkingConfig(thinking_budget=llm_cfg.thinking_budget)
 
     client = genai.Client(api_key=settings.gemini_api_key)
     response = client.models.generate_content(
@@ -158,7 +156,7 @@ def _call_gemini(
         config=types.GenerateContentConfig(
             system_instruction=system_prompt,
             temperature=llm_cfg.temperature,
-            max_output_tokens=max_output_tokens,
+            max_output_tokens=llm_cfg.max_tokens,
             thinking_config=thinking_config,
             response_mime_type="application/json" if json_output else None,
         ),

pipeline/stages/vector.py

@@ -13,6 +13,8 @@ from pipeline.stages.products import ProductCandidate
 logger = logging.getLogger(__name__)
 
 COLLECTION = "ProductProfile"
+RERANK_PROPERTY = "situation_fit_summary"
+_RERANKER_MODULE = "reranker-transformers"
 
 
 @dataclass(frozen=True)
@@ -45,12 +47,46 @@ def connect_client(settings: Settings) -> Any:
         headers["X-Goog-Studio-Api-Key"] = settings.gemini_api_key
         headers["X-Palm-Api-Key"] = settings.gemini_api_key
 
-    return weaviate.connect_to_local(
+    client = weaviate.connect_to_local(
         host=settings.weaviate.host,
         port=settings.weaviate.http_port,
         grpc_port=settings.weaviate.grpc_port,
         headers=headers or None,
     )
+    _ensure_transformers_reranker(client)
+    return client
+
+
+def _ensure_transformers_reranker(client: Any) -> None:
+    """
+    Attach the local transformers reranker to ProductProfile when none is set.
+
+    The collection setting is mutable on Weaviate 1.32 and does not re-embed
+    existing objects. The model itself is the reranker-transformers container
+    in docker-compose.yml (ms-marco-MiniLM-L-6-v2).
+    """
+    from weaviate.classes.config import Reconfigure
+
+    enabled_modules = (client.get_meta() or {}).get("modules") or {}
+    if _RERANKER_MODULE not in enabled_modules:
+        raise PipelineError(
+            f"Weaviate does not have the '{_RERANKER_MODULE}' module enabled. "
+            "Recreate the containers with `docker compose up -d` so the "
+            "reranker-transformers service from docker-compose.yml is running."
+        )
+
+    collection = client.collections.get(COLLECTION)
+    current = collection.config.get().reranker_config
+    if current is not None:
+        reranker = current.reranker
+        module = str(getattr(reranker, "value", reranker))
+        if module == _RERANKER_MODULE:
+            return
+        logger.warning("ProductProfile reranker is %s; leaving it unchanged.", module)
+        return
+
+    collection.config.update(reranker_config=Reconfigure.Reranker.transformers())
+    logger.info("Enabled %s on %s.", _RERANKER_MODULE, COLLECTION)
 
 
 def search_products(
@@ -60,12 +96,14 @@ def search_products(
 ) -> list[RecommendedProduct]:
     """
     Embed each query and retrieve the nearest ProductProfile objects, scoped
-    server-side to the SQL-prefiltered allowlist.
+    server-side to the SQL-prefiltered allowlist, then rerank that pool with
+    the local transformers cross-encoder on `situation_fit_summary`.
 
-    One query  → top 4 overall.
+    One query  → retrieve 10, keep the top 4 in rerank order.
-    Two queries → top 2 unique per query (deduplicated preferring first
+    Two queries → retrieve 6 per query, keep the top 2 unique per query
-    occurrence). Later queries over-fetch so duplicates of earlier hits
+    (deduplicated preferring first occurrence). A product already kept for
-    can be replaced from the same response without a second search.
+    an earlier query is skipped so the later query can fill its two slots
+    from the rest of its reranked six.
 
     `client` is a shared, already-connected Weaviate client (see
     `pipeline.resources.Resources`); this function does not open or close a
@@ -82,10 +120,11 @@ def search_products(
     if not queries:
         raise PipelineError("Cannot run vector search without at least one query.")
 
-    from weaviate.classes.query import Filter, MetadataQuery
+    from weaviate.classes.query import Filter, MetadataQuery, Rerank
 
     allowlist = {c.product_id: c for c in candidates}
     per_query_limit = 4 if len(queries) == 1 else 2
+    fetch_limit = 10 if len(queries) == 1 else 6
     query_filter = Filter.by_property("product_id").contains_any(list(allowlist))
 
     collection = client.collections.get(COLLECTION)
@@ -93,11 +132,11 @@ def search_products(
     seen_ids: set[str] = set()
 
     for query_index, query in enumerate(queries):
-        fetch_limit = per_query_limit * (query_index + 1)
         response = collection.query.near_text(
             query=query,
             limit=fetch_limit,
             filters=query_filter,
+            rerank=Rerank(prop=RERANK_PROPERTY, query=query),
             return_metadata=MetadataQuery(distance=True),
             return_properties=["product_id", "product_name"],
         )

tests/test_vector.py

@@ -6,8 +6,9 @@ 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 search_products
+from pipeline.stages.vector import _ensure_transformers_reranker, search_products
 
 
 def _candidate(product_id: str) -> ProductCandidate:
@@ -37,8 +38,11 @@ class FakeNearTextQuery:
         filters: Any,
         return_metadata: Any,
         return_properties: list[str],
+        rerank: Any = None,
     ) -> SimpleNamespace:
-        self.calls.append({"query": query, "limit": limit, "filters": filters})
+        self.calls.append(
+            {"query": query, "limit": limit, "filters": filters, "rerank": rerank}
+        )
         hits = list(self.hits_by_query.get(query, []))
         return SimpleNamespace(objects=hits[:limit])
 
@@ -53,8 +57,8 @@ class FakeClient:
 
 class SearchProductsLimitTests(unittest.TestCase):
     def test_one_query_keeps_top_four(self) -> None:
-        candidates = [_candidate(f"p{i}") for i in range(1, 7)]
+        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, 7)]
+        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)
@@ -63,16 +67,16 @@ class SearchProductsLimitTests(unittest.TestCase):
         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)],
+            [("q1", 10)],
         )
-        self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
+        self._assert_rerank(client, ["q1"])
 
     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, 13)]
         client = FakeClient(
             {
-                "q1": [_hit("p1"), _hit("p2"), _hit("p3")],
+                "q1": [_hit(f"p{i}") for i in range(1, 7)],
-                "q2": [_hit("p4"), _hit("p5"), _hit("p6")],
+                "q2": [_hit(f"p{i}") for i in range(7, 13)],
             }
         )
 
@@ -80,20 +84,27 @@ class SearchProductsLimitTests(unittest.TestCase):
 
         self.assertEqual(
             [(item.product_id, item.query_index) for item in results],
-            [("p1", 0), ("p2", 0), ("p4", 1), ("p5", 1)],
+            [("p1", 0), ("p2", 0), ("p7", 1), ("p8", 1)],
         )
         self.assertEqual(
             [(c["query"], c["limit"]) for c in client.query.calls],
-            [("q1", 2), ("q2", 4)],
+            [("q1", 6), ("q2", 6)],
         )
-        self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
+        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, 6)]
+        candidates = [_candidate(f"p{i}") for i in range(1, 9)]
         client = FakeClient(
             {
-                "q1": [_hit("p1"), _hit("p2"), _hit("p3")],
+                "q1": [_hit(f"p{i}") for i in range(1, 7)],
-                "q2": [_hit("p2"), _hit("p1"), _hit("p4"), _hit("p5")],
+                "q2": [
+                    _hit("p2"),
+                    _hit("p1"),
+                    _hit("p4"),
+                    _hit("p5"),
+                    _hit("p7"),
+                    _hit("p8"),
+                ],
             }
         )
 
@@ -105,21 +116,69 @@ class SearchProductsLimitTests(unittest.TestCase):
         )
         self.assertEqual(
             [(c["query"], c["limit"]) for c in client.query.calls],
-            [("q1", 2), ("q2", 4)],
+            [("q1", 6), ("q2", 6)],
         )
-        self.assertTrue(all(c["filters"] is not None for c in client.query.calls))
+        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, 10)]
+        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"], 4)
+        self.assertEqual(client.query.calls[0]["limit"], 10)
-        self.assertIsNotNone(client.query.calls[0]["filters"])
+        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__":