56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
"""Offline-safe, ID-correlated provider batch execution."""
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
from .provider import ProviderTransportError
|
|
|
|
|
|
@dataclass
|
|
class BatchRunResult:
|
|
results_by_work_item_id: dict[str, dict] = field(default_factory=dict)
|
|
errors: dict[str, str] = field(default_factory=dict)
|
|
provider_calls: int = 0
|
|
retries: int = 0
|
|
usage: dict = field(default_factory=dict)
|
|
|
|
|
|
def run_batches(items, provider, batch_size: int, max_retries: int = 2, *, scenario: str | None = None) -> BatchRunResult:
|
|
if batch_size <= 0:
|
|
raise ValueError("batch_size must be positive")
|
|
if max_retries < 0:
|
|
raise ValueError("max_retries must not be negative")
|
|
|
|
result = BatchRunResult()
|
|
for start in range(0, len(items), batch_size):
|
|
batch = list(items[start:start + batch_size])
|
|
expected_ids = [str(item.work_item_id) for item in batch]
|
|
answers = None
|
|
for attempt in range(max_retries + 1):
|
|
result.provider_calls += 1
|
|
try:
|
|
kwargs = {} if scenario is None else {"scenario": scenario}
|
|
answers = provider.select_batch(batch, **kwargs)
|
|
break
|
|
except ProviderTransportError as error:
|
|
if attempt == max_retries:
|
|
for work_item_id in expected_ids:
|
|
result.errors[work_item_id] = str(error)
|
|
else:
|
|
result.retries += 1
|
|
except Exception as error:
|
|
for work_item_id in expected_ids:
|
|
result.errors[work_item_id] = str(error)
|
|
break
|
|
if answers is None:
|
|
continue
|
|
answer_ids = [str(answer.get("work_item_id")) for answer in answers]
|
|
if len(answer_ids) != len(expected_ids) or set(answer_ids) != set(expected_ids) or len(set(answer_ids)) != len(answer_ids):
|
|
error = "provider response did not contain exactly one answer per work_item_id"
|
|
for work_item_id in expected_ids:
|
|
result.errors[work_item_id] = error
|
|
continue
|
|
result.results_by_work_item_id.update({str(answer["work_item_id"]): answer for answer in answers})
|
|
|
|
result.usage = dict(getattr(provider, "usage", {}) or {})
|
|
return result
|