From 46cc17e5d5a00641a7d6826030ae5f49030a80e5 Mon Sep 17 00:00:00 2001 From: Arsham Mirehvandi <130353205+arshammi@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:25:54 +0200 Subject: [PATCH] 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 | 17 ++++--- config.yaml | 14 +++--- docker-compose.yml | 16 +++++-- pipeline/config.py | 24 +++++++--- pipeline/stages/llm.py | 12 ++--- pipeline/stages/vector.py | 55 ++++++++++++++++++---- tests/test_vector.py | 99 +++++++++++++++++++++++++++++++-------- 7 files changed, 176 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index 75117d1..8f54725 100644 --- a/README.md +++ b/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 - `advice_generation` span (gemini-2.5-pro), each wrapping the provider's own +- a `query_synthesis` span (gemini-3.8-flash by default) and an + `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 -already includes `gemini-2.5-flash` and `gemini-2.5-pro` — no custom entry -under **Settings → Models** is needed for either model as configured today. -If a model is ever renamed or swapped to one Phoenix doesn't recognise, add -it there (regex `Name Pattern`, provider `google`, per-1M-token prices). +Phoenix computes cost per span from its built-in model pricing table. If a +configured model (`gemini-3.8-flash` or `gemini-3.1-pro-preview`) is not +recognised, add it under **Settings → Models** (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 diff --git a/config.yaml b/config.yaml index d0f534e..054a9d1 100644 --- a/config.yaml +++ b/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, -# Gemini's max_output_tokens is max_tokens + thinking_budget, so max_tokens -# is the visible-output budget. Omit it to keep max_tokens as the shared cap. +# thinking_level (optional, Gemini 3 only): low | medium | high. +# Gemini counts thinking tokens inside max_tokens, so max_tokens is the +# 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 - max_tokens: 5120 - thinking_budget: 8192 + gemini_model: gemini-3.1-pro-preview + max_tokens: 13312 + thinking_level: medium temperature: 0.0 # ── Crop-first worklist ──────────────────────────────────────────────────────── diff --git a/docker-compose.yml b/docker-compose.yml index 3a015c2..3dfb817 100644 --- a/docker-compose.yml +++ b/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: diff --git a/pipeline/config.py b/pipeline/config.py index 217ddc7..5ac57a6 100644 --- a/pipeline/config.py +++ b/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( diff --git a/pipeline/stages/llm.py b/pipeline/stages/llm.py index 8adbc1f..6efd653 100644 --- a/pipeline/stages/llm.py +++ b/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: - max_output_tokens = llm_cfg.max_tokens + llm_cfg.thinking_budget - thinking_config = types.ThinkingConfig(thinking_budget=llm_cfg.thinking_budget) + if llm_cfg.thinking_level is not None: + thinking_config = types.ThinkingConfig(thinking_level=llm_cfg.thinking_level) 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, ), diff --git a/pipeline/stages/vector.py b/pipeline/stages/vector.py index 3f66fa4..0e94c6e 100644 --- a/pipeline/stages/vector.py +++ b/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. - Two queries → top 2 unique per query (deduplicated preferring first - occurrence). Later queries over-fetch so duplicates of earlier hits - can be replaced from the same response without a second search. + One query → retrieve 10, keep the top 4 in rerank order. + Two queries → retrieve 6 per query, keep the top 2 unique per query + (deduplicated preferring first occurrence). A product already kept for + 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"], ) diff --git a/tests/test_vector.py b/tests/test_vector.py index 337fb12..80f4582 100644 --- a/tests/test_vector.py +++ b/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)] - hits = [_hit(f"p{i}", distance=0.1 * 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, 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")], - "q2": [_hit("p4"), _hit("p5"), _hit("p6")], + "q1": [_hit(f"p{i}") for i in range(1, 7)], + "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")], - "q2": [_hit("p2"), _hit("p1"), _hit("p4"), _hit("p5")], + "q1": [_hit(f"p{i}") for i in range(1, 7)], + "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.assertIsNotNone(client.query.calls[0]["filters"]) + 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__":