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.
This commit is contained in:
Arsham Mirehvandi
parent
e45376fd4b
commit
46cc17e5d5
17
README.md
17
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
|
||||
|
||||
14
config.yaml
14
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 ────────────────────────────────────────────────────────
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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,
|
||||
),
|
||||
|
||||
@ -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"],
|
||||
)
|
||||
|
||||
@ -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__":
|
||||
|
||||
Loading…
Reference in New Issue
Block a user