# #region AdaptiveBatchSizer [C:3] [TYPE Module] [SEMANTICS translate, batch, sizing, token-budget] # @BRIEF Adaptive batch sizing for LLM translation — splits source rows into variable-sized # batches based on actual content length and token budget estimates. # @LAYER Domain # @RELATION DEPENDS_ON -> [estimate_token_budget] # @RELATION DEPENDS_ON -> [LLMProviderService] # @RELATION DEPENDS_ON -> [TranslationJob] # @RATIONALE Extracted from TranslationExecutor to comply with INV_7 (module < 400 lines). # Fixed batch_size of 50 wastes LLM context for short rows and overflows for # long rows. Variable sizing maximizes throughput while preventing truncation. # @REJECTED Fixed batch_size of 50 — causes truncation on long-content rows. # Single monolithic batch — would lose all progress on any failure. from sqlalchemy.orm import Session from ...core.config_manager import ConfigManager from ...core.logger import belief_scope, logger from ...models.translate import TranslationJob from ...services.llm_provider import LLMProviderService from ._token_budget import ( JSON_OVERHEAD_PER_ROW, MAX_OUTPUT_HEADROOM, OUTPUT_PER_ROW_PER_LANG, PROMPT_BASE_TOKENS, REASONING_OVERHEAD, estimate_token_budget, ) from ._utils import estimate_row_tokens # #region AdaptiveBatchSizer [C:3] [TYPE Class] # @BRIEF Split source rows into auto-sized batches based on token budget estimates. class AdaptiveBatchSizer: """Split source rows into auto-sized batches based on token budget estimates. Each batch is sized so that its total estimated tokens fit within the available context window (input budget), accounting for prompt overhead, dictionary entries, and output tokens. """ def __init__(self, db: Session, config_manager: ConfigManager) -> None: self.db = db self.config_manager = config_manager # #region resolve_provider_model [C:2] [TYPE Function] [SEMANTICS llm, provider, model] # @BRIEF Resolve the LLM provider model name for token budget estimation. # @POST Returns model name string or None if resolution fails. # @SIDE_EFFECT DB query to LLM provider table. def resolve_provider_model(self, job: TranslationJob) -> str | None: """Resolve the provider model name for token budget estimation.""" if not job.provider_id: return None try: p_svc = LLMProviderService(self.db) p = p_svc.get_provider(job.provider_id) if p: return p.default_model or "gpt-4o-mini" except Exception: pass return None # #endregion resolve_provider_model # #region auto_size_batches [C:3] [TYPE Function] [SEMANTICS translate, batch, sizing] # @BRIEF Split source rows into variable-sized batches based on content length. # @PRE source_rows is non-empty. job has valid config. # @POST Returns list of batches, each batch is a list of row dicts. # Each batch fits within the estimated token budget for its rows. # @SIDE_EFFECT DB query to resolve provider model. Logs batch statistics. def auto_size_batches( self, job: TranslationJob, source_rows: list[dict], target_languages: list[str], provider_info: str | None = None, ) -> list[list[dict]]: """Split source rows into auto-sized batches based on content length. Each batch is sized so that its total estimated tokens fit within the available context window (input budget), accounting for prompt overhead, dictionary entries, and output tokens. """ with belief_scope("AdaptiveBatchSizer.auto_size_batches"): if not source_rows: return [] if provider_info is None: provider_info = self.resolve_provider_model(job) # 1. Estimate per-row token counts row_tokens: list[int] = [] for row in source_rows: source_text = row.get("source_text", "") source_data = row.get("source_data") tokens = estimate_row_tokens(source_text, source_data, job) row_tokens.append(tokens) # 2. Get budget recommendation budget = estimate_token_budget( source_rows=source_rows, target_languages=target_languages, source_column=job.translation_column or "source_text", context_columns=job.context_columns, batch_size=len(source_rows), provider_info=provider_info, ) recommended = budget.get("batch_size_adjusted", 0) # Fallback: if budget calculation fails, use fixed size if recommended <= 0: fallback_size = job.batch_size or 50 logger.explore("Token budget returned zero — falling back to fixed batch size", { "fallback_size": fallback_size, "total_rows": len(source_rows), }) return [ source_rows[i:i + fallback_size] for i in range(0, len(source_rows), fallback_size) ] # 3. Compute per-batch row-content budget # Use the ACTUAL available input capacity (context_window - max_output_tokens), # NOT the sum of the first N rows. This prevents long rows from being # incorrectly placed in 1-row batches when the context window has plenty of room. available_input = budget.get("available_input_budget") if available_input is not None: # New-style: use actual available input capacity per_batch_budget = available_input - PROMPT_BASE_TOKENS else: # Fallback for tests: use estimated_input (sum of first N rows) estimated_input = budget.get("estimated_input_tokens", 50000) per_batch_budget = estimated_input - PROMPT_BASE_TOKENS if per_batch_budget <= 0: fallback_size = job.batch_size or 50 logger.explore("Per-batch budget collapsed — falling back to fixed batch size", { "per_batch_budget": per_batch_budget, "fallback_size": fallback_size, "total_rows": len(source_rows), }) return [ source_rows[i:i + fallback_size] for i in range(0, len(source_rows), fallback_size) ] # 4. Greedy batch splitting # Compute max rows per batch from OUTPUT constraint. # This prevents output truncation (finish_reason=length) when batching # many short rows within the large input budget. num_languages = len(target_languages) max_output_tokens_val = budget.get("max_output_tokens") if max_output_tokens_val is not None: output_per_row = num_languages * OUTPUT_PER_ROW_PER_LANG + JSON_OVERHEAD_PER_ROW available_output = max_output_tokens_val - REASONING_OVERHEAD - MAX_OUTPUT_HEADROOM max_rows_by_output = max(available_output // output_per_row, 1) if output_per_row > 0 else 20 max_rows_hard_cap = max_rows_by_output else: # Fallback for tests: old formula max_rows_hard_cap = max(recommended * 2, 20) # 5. Respect job.batch_size as the absolute maximum rows per batch. # User-configured batch_size overrides model-based estimates to # prevent LLM quality degradation on large batches. if job.batch_size: max_rows_hard_cap = min(max_rows_hard_cap, job.batch_size) batches: list[list[dict]] = [] current_batch: list[dict] = [] current_tokens = 0 for i, row in enumerate(source_rows): rt = row_tokens[i] if rt > per_batch_budget: if current_batch: batches.append(current_batch) current_batch = [] current_tokens = 0 logger.reason("Single row exceeds per-batch token budget — placing in own batch", { "row_index": i, "row_tokens": rt, "per_batch_budget": per_batch_budget, }) batches.append([row]) continue should_split = bool( current_batch and (current_tokens + rt > per_batch_budget or len(current_batch) >= max_rows_hard_cap) ) if should_split: batches.append(current_batch) current_batch = [row] current_tokens = rt else: current_batch.append(row) current_tokens += rt if current_batch: batches.append(current_batch) # 5. Log adaptive batch statistics if batches: avg_size = sum(len(b) for b in batches) / max(1, len(batches)) max_size = max(len(b) for b in batches) logger.reason("Auto-sized batches", { "total_rows": len(source_rows), "num_batches": len(batches), "avg_batch_size": round(avg_size, 1), "max_batch_size": max_size, "recommended_batch_size": recommended, "per_batch_budget": per_batch_budget, }) return batches # #endregion auto_size_batches # #endregion AdaptiveBatchSizer # #endregion AdaptiveBatchSizer