diff --git a/workers/executor/executors/lookup_enrichment.py b/workers/executor/executors/lookup_enrichment.py index 0b224badb5..3cfbe6b767 100644 --- a/workers/executor/executors/lookup_enrichment.py +++ b/workers/executor/executors/lookup_enrichment.py @@ -32,6 +32,30 @@ def is_blank(value: Any) -> bool: return False +def preload_reference_texts(lookup_config: dict[str, Any] | None) -> Any: + """Read a lookup's reference files once via the enrichment plugin. + + For callers that enrich many rows against one ``lookup_config`` (per-row + enrichment), pass the result as ``reference_texts`` to + ``run_lookup_enrichment`` so the remote read isn't repeated per row. + Returns ``None`` (fall back to per-call loading) when the plugin/config is + unavailable or the read fails. + """ + lookup_cls = ExecutorPluginLoader.get("lookup-enrichment") + if not (lookup_config and lookup_cls): + return None + preload = getattr(lookup_cls, "preload_reference_texts", None) + if not callable(preload): + return None + try: + return preload(lookup_config) + except Exception: + # Honour the None-on-failure contract even if the plugin's preload + # raises — the caller falls back to per-call reference loading. + logger.exception("Failed to preload lookup reference texts") + return None + + def run_lookup_enrichment( output: dict[str, Any], structured_output: dict[str, Any], @@ -40,11 +64,15 @@ def run_lookup_enrichment( shim: Any, llm_cls: Any, usage_kwargs: dict[str, Any] | None = None, + reference_texts: dict[str, str] | None = None, ) -> list[dict[str, Any]]: """Run lookup enrichment plugin if enabled and available. Returns any usage records the plugin emitted (recovered even on plugin failure) so the caller can extend its billing batch. + + ``reference_texts`` (from :func:`preload_reference_texts`) skips the + plugin's per-call remote read of the reference files. """ prompt_name = output[PSKeys.NAME] current_value = structured_output.get(prompt_name) @@ -73,6 +101,7 @@ def run_lookup_enrichment( prompt_name=prompt_name, shim=shim, usage_kwargs=usage_kwargs, + reference_texts=reference_texts, ) metrics.setdefault(prompt_name, {})[lookup_cls.METRICS_KEY] = outcome.llm_metrics except Exception: diff --git a/workers/file_processing/structure_tool_task.py b/workers/file_processing/structure_tool_task.py index 4a3c2bb4ce..971a783980 100644 --- a/workers/file_processing/structure_tool_task.py +++ b/workers/file_processing/structure_tool_task.py @@ -479,6 +479,9 @@ def _execute_structure_tool_impl(params: dict) -> dict: "execution_id": execution_id, "PLATFORM_SERVICE_API_KEY": platform_service_api_key, "group_key": at_settings.get("group_key", ""), + # Set on the output at export by prompt_studio_registry_helper.py + # when a lookup is assigned; None otherwise. + "lookup_config": at_output.get("lookup_config"), } at_ctx = ExecutionContext( executor_name="agentic_table", diff --git a/workers/shared/parallel_map.py b/workers/shared/parallel_map.py new file mode 100644 index 0000000000..6467034939 --- /dev/null +++ b/workers/shared/parallel_map.py @@ -0,0 +1,122 @@ +"""Generic bounded-concurrency map for LLM (and other I/O-bound) call sites. + +Runs ``worker(item)`` across a thread pool capped at ``max_workers`` and returns +results in input order. It is **length-preserving**: the output list is always +the same length as the input and failed items keep their slot, so callers can +realign results to inputs by index (row *i* in maps to row *i* out). + +``max_workers=1`` runs effectively sequentially, so a single knob covers both the +sequential and bounded-parallel strategies with no code change. + +Rate limiting / retries are intentionally NOT handled here. Bounding +``max_workers`` caps concurrent calls; anything finer (provider rate limits, +``Retry-After``, backoff on 429/5xx/overloaded) is the LLM client's job, not +this generic utility's. Size ``max_workers`` conservatively for the provider. + +NOTE: there is no early-abort. All submitted work runs to completion even if the +caller is later cancelled or times out — in-flight threads cannot be killed. +Callers that need to stop a doomed run early must gate submission themselves. +""" + +import logging +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import TypeVar + +logger = logging.getLogger(__name__) + +T = TypeVar("T") +R = TypeVar("R") + +# Cap on how many per-item failures are logged with a full traceback before +# switching to one-line records — stops a dead adapter from emitting hundreds +# of identical stack traces (and Sentry events) for one logical failure. +_MAX_TRACEBACKS = 5 + + +def parallel_map( + items: list[T], + worker: Callable[[T], R], + *, + max_workers: int, + on_error: Callable[[int, T, Exception], R] | None = None, + label: str = "", +) -> list[R | None]: + """Apply ``worker`` to each item with bounded concurrency, order preserved. + + Args: + items: Inputs to process. An empty list returns ``[]``. + worker: Called as ``worker(item)`` for each item. + max_workers: Thread-pool cap. ``<= 1`` runs effectively sequentially. + on_error: Called as ``on_error(index, item, exception)`` to produce a + fallback result when a worker raises. If omitted, a failed item's + slot is left as ``None`` (still counted — length is preserved). + Prefer passing ``on_error`` so a failed slot is distinguishable + from a legitimate ``None`` result. + label: Optional label for the progress log line. + + Returns: + A list the SAME LENGTH as ``items``, results in input order. A failed + item holds the ``on_error`` fallback, or ``None`` when ``on_error`` is + omitted — hence the ``R | None`` element type. + """ + if not items: + return [] + + n = len(items) + effective_workers = max(max_workers, 1) + results: list[R | None] = [None] * n + failures = 0 + + if effective_workers > 1: + suffix = f" ({label})" if label else "" + logger.info( + "parallel_map: %d items across up to %d workers%s", + n, + effective_workers, + suffix, + ) + + # The as_completed loop body runs on the calling thread (only worker() runs + # in the pool), so results/logging here need no locking. + with ThreadPoolExecutor(max_workers=effective_workers) as executor: + future_to_idx = { + executor.submit(worker, item): idx for idx, item in enumerate(items) + } + for future in as_completed(future_to_idx): + idx = future_to_idx[future] + try: + results[idx] = future.result() + except Exception as e: + failures += 1 + # Full traceback for the first few, then one-liners so a + # systemic failure doesn't flood logs. + logger.error( + "parallel_map: item %d/%d failed: %s", + idx + 1, + n, + e, + exc_info=failures <= _MAX_TRACEBACKS, + ) + if on_error is not None: + # Guard on_error itself: if it raises, it would propagate + # through the pool's __exit__ (shutdown(wait=True)) and hang + # the caller until every in-flight task drains, with no log. + try: + results[idx] = on_error(idx, items[idx], e) + except Exception: + logger.error( + "parallel_map: on_error raised for item %d/%d", + idx + 1, + n, + exc_info=True, + ) + + if failures: + logger.warning( + "parallel_map: %d/%d items failed%s", + failures, + n, + f" ({label})" if label else "", + ) + return results