"""Resources shared across every job in a batch run. A single-field CLI run opened one SQL connection, one Weaviate client, and made its two LLM calls without needing to coordinate with anyone else. The batch runs many jobs concurrently across a thread pool, so several things that used to be "create it, use it, close it" per field now need to be long-lived, shared, and safe to touch from multiple threads at once: - pyodbc connections must not be shared across threads, so each worker thread gets its own, created lazily and reused for every job it picks up. - The Weaviate client is documented as safe for concurrent query use, so one instance is opened for the whole batch instead of one per field. - `near_text` search embeds through the same Gemini API key as the two LLM calls (`ProductProfile` is vectorised with text2vec-palm / gemini-embedding-001), so a single semaphore + rate limiter covers all three Gemini-backed calls per job instead of treating "LLM" and "Weaviate" as independent budgets. - The `products` table and the `ProductJsonBuilder` per-product cache are loaded/shared once instead of once per field. """ from __future__ import annotations import logging import threading import time from typing import Any from pipeline.config import LlmSettings, Settings from pipeline.db import connect_raw from pipeline.prompts import PromptRegistry from pipeline.retry import call_with_llm_retry, call_with_weaviate_retry from pipeline.stages.llm import call_llm as _raw_call_llm from pipeline.stages.product_json import ProductJsonBuilder from pipeline.stages.products import ProductCandidate, ProductIndex from pipeline.stages.vector import RecommendedProduct, connect_client from pipeline.stages.vector import search_products as _raw_search_products logger = logging.getLogger(__name__) class RateLimiter: """ Thread-safe pacing limiter: spreads calls out to at most N per minute. This paces call *starts* evenly (a leaky bucket, not a bursty token bucket with saved-up credit) — the simplest thing that reliably keeps a pool of worker threads under an external rate limit, which matters more here than allowing occasional bursts. """ def __init__(self, requests_per_minute: int) -> None: self._interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0.0 self._lock = threading.Lock() self._next_allowed = 0.0 def acquire(self) -> None: if self._interval <= 0: return with self._lock: now = time.monotonic() wait = self._next_allowed - now if wait > 0: time.sleep(wait) now = time.monotonic() self._next_allowed = max(now, self._next_allowed) + self._interval class Resources: """Shared, thread-safe state for one batch run. Create once, close once.""" def __init__(self, settings: Settings, prompts: PromptRegistry) -> None: self.settings = settings self.prompts = prompts self._local = threading.local() self._connections: list[Any] = [] self._connections_lock = threading.Lock() self.sql_sem = threading.Semaphore(settings.concurrency.sql_limit) self.gemini_sem = threading.Semaphore(settings.concurrency.gemini_limit) self.gemini_rate_limiter = RateLimiter(settings.concurrency.gemini_requests_per_minute) self._weaviate_client: Any = None self._weaviate_lock = threading.Lock() self.product_index: ProductIndex | None = None self._builder_cache: dict[tuple[str, str], ProductJsonBuilder] = {} self._builder_lock = threading.Lock() # -- SQL ----------------------------------------------------------------- def sql_connection(self) -> Any: """ Return this thread's SQL connection, opening one on first use. Not gated by `sql_sem`: with one connection per worker thread kept open for the batch's whole lifetime, gating *creation* would only throttle how fast the pool warms up, and gating it for the connection's entire lifetime would deadlock as soon as `sql_limit` is set below `workers` (the remaining threads would block forever on a permit nothing ever releases). `sql_sem` instead bounds how many jobs may run their SQL-heavy phase at the same instant — see `pipeline.job.run_job`, which acquires it around each phase, not around the connection. """ conn = getattr(self._local, "conn", None) if conn is None: conn = connect_raw(self.settings.sql) self._local.conn = conn with self._connections_lock: self._connections.append(conn) return conn def load_product_index(self) -> ProductIndex: """Load the whole `products` table once; call before starting the pool.""" self.product_index = ProductIndex.load(self.sql_connection()) return self.product_index def product_json_builder(self, crop_english: str, disease_english: str) -> ProductJsonBuilder: """ One `ProductJsonBuilder` per (crop, disease), shared by every field growing that crop with that disease across all worker threads. """ if self.product_index is None: raise RuntimeError("load_product_index() must run before product_json_builder().") key = (crop_english, disease_english) with self._builder_lock: builder = self._builder_cache.get(key) if builder is None: builder = ProductJsonBuilder(self.product_index, crop_english, disease_english) self._builder_cache[key] = builder return builder # -- Weaviate / Gemini ----------------------------------------------------- def _weaviate(self) -> Any: if self._weaviate_client is None: with self._weaviate_lock: if self._weaviate_client is None: self._weaviate_client = connect_client(self.settings) return self._weaviate_client def call_llm( self, settings: Settings, llm_cfg: LlmSettings, system_prompt: str, user_content: str, json_output: bool = False, ) -> str: """ Rate-limited, retried, concurrency-bounded drop-in for `pipeline.stages.llm.call_llm`. Signature-compatible so it can be passed as `synthesize_queries(..., llm_call=resources.call_llm)` / `generate_advice(..., llm_call=resources.call_llm)`. """ def _attempt() -> str: with self.gemini_sem: self.gemini_rate_limiter.acquire() return _raw_call_llm(settings, llm_cfg, system_prompt, user_content, json_output) return call_with_llm_retry(self.settings.retry, _attempt) def search_products( self, queries: list[str], candidates: list[ProductCandidate], ) -> list[RecommendedProduct]: """Rate-limited, retried wrapper around `pipeline.stages.vector.search_products`.""" def _attempt() -> list[RecommendedProduct]: with self.gemini_sem: self.gemini_rate_limiter.acquire() return _raw_search_products(self._weaviate(), queries, candidates) return call_with_weaviate_retry(self.settings.retry, _attempt) # -- Lifecycle ------------------------------------------------------------- def close(self) -> None: """Close every pooled SQL connection and the shared Weaviate client.""" with self._connections_lock: for conn in self._connections: try: conn.close() except Exception: logger.warning("Error closing a pooled SQL connection.", exc_info=True) self._connections.clear() if self._weaviate_client is not None: try: self._weaviate_client.close() except Exception: logger.warning("Error closing the Weaviate client.", exc_info=True) self._weaviate_client = None