feat(translate): auto batch_size estimator for LLM token budget

New _token_budget.py calculates safe batch_size and max_tokens based on source text length, target languages, dictionary size, and model context window (64K DeepSeek v4 Flash). preview.py: auto-reduces sample_size when texts are long. executor.py: per-batch dynamic max_tokens. 12 new tests.
This commit is contained in:
2026-05-15 09:53:54 +03:00
parent 365d710806
commit 4641c82397
4 changed files with 478 additions and 6 deletions

View File

@@ -35,6 +35,7 @@ from ...models.translate import (
)
from ...services.llm_prompt_templates import render_prompt
from ...services.llm_provider import LLMProviderService
from ._token_budget import DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_OUTPUT_TOKENS, estimate_token_budget
from .dictionary import DictionaryManager
# #region DEFAULT_EXECUTION_PROMPT_TEMPLATE [TYPE Constant]
@@ -215,6 +216,46 @@ class TranslationPreview:
f"val='{first_row.get(job.translation_column, '')}'"
)
# 3b. Check token budget and auto-reduce sample size if needed
token_budget = estimate_token_budget(
source_rows=source_rows,
target_languages=target_languages,
source_column=job.translation_column,
context_columns=job.context_columns,
batch_size=actual_row_count,
context_window=DEFAULT_CONTEXT_WINDOW,
max_output_tokens=DEFAULT_MAX_OUTPUT_TOKENS,
)
if token_budget["warning"]:
logger.explore("Token budget warning", {
"warning": token_budget["warning"],
"sample_size": actual_row_count,
"adjusted": token_budget["batch_size_adjusted"],
})
# If budget says we need fewer rows than fetched, truncate
adjusted_size = token_budget["batch_size_adjusted"]
if adjusted_size < actual_row_count:
logger.explore(
f"Reducing preview from {actual_row_count} to {adjusted_size} rows "
f"to fit within context window",
{"estimated_input": token_budget["estimated_input_tokens"],
"estimated_output": token_budget["estimated_output_tokens"],
"context_window": DEFAULT_CONTEXT_WINDOW},
)
source_rows = source_rows[:adjusted_size]
actual_row_count = len(source_rows)
# Recalculate token budget for truncated set
token_budget = estimate_token_budget(
source_rows=source_rows,
target_languages=target_languages,
source_column=job.translation_column,
context_columns=job.context_columns,
batch_size=actual_row_count,
context_window=DEFAULT_CONTEXT_WINDOW,
max_output_tokens=DEFAULT_MAX_OUTPUT_TOKENS,
)
# 4. Build prompt context from rows
all_source_texts = []
row_meta: list[dict[str, Any]] = []
@@ -291,16 +332,19 @@ class TranslationPreview:
sample_size, num_languages, sample_total_tokens, sample_cost
)
# 8. Call LLM
# 8. Call LLM with token-budget-aware max_tokens
max_tokens_for_call = token_budget["max_output_needed"]
logger.reason("Calling LLM for preview translation", {
"provider_id": job.provider_id,
"row_count": actual_row_count,
"num_languages": num_languages,
"estimated_tokens": sample_total_tokens,
"max_tokens": max_tokens_for_call,
})
llm_response = self._call_llm(
job=job,
prompt=prompt,
max_tokens=max_tokens_for_call,
)
# 9. Parse LLM response (multi-language)
@@ -447,6 +491,13 @@ class TranslationPreview:
"estimated_cost": total_est_cost,
"warning": cost_warning,
},
"token_budget": {
"batch_size_adjusted": token_budget["batch_size_adjusted"],
"estimated_input_tokens": token_budget["estimated_input_tokens"],
"estimated_output_tokens": token_budget["estimated_output_tokens"],
"max_output_needed": token_budget["max_output_needed"],
"warning": token_budget["warning"],
},
"config_hash": config_hash,
"dict_snapshot_hash": dict_snapshot_hash,
}
@@ -887,7 +938,7 @@ class TranslationPreview:
# @PRE: job has a valid provider_id.
# @POST: Returns raw LLM response string.
# @SIDE_EFFECT: Makes HTTP call to LLM provider.
def _call_llm(self, job: TranslationJob, prompt: str) -> str:
def _call_llm(self, job: TranslationJob, prompt: str, max_tokens: int = 8192) -> str:
with belief_scope("TranslationPreview._call_llm"):
if not job.provider_id:
raise ValueError("Job has no LLM provider configured")
@@ -912,6 +963,7 @@ class TranslationPreview:
model=model,
prompt=prompt,
provider_type=provider_type,
max_tokens=max_tokens,
)
else:
raise ValueError(f"Unsupported provider type '{provider_type}' for preview")
@@ -936,6 +988,7 @@ class TranslationPreview:
model: str,
prompt: str,
provider_type: str = "openai",
max_tokens: int = 8192,
) -> str:
with belief_scope("TranslationPreview._call_openai_compatible"):
import requests as http_requests
@@ -952,7 +1005,7 @@ class TranslationPreview:
{"role": "user", "content": prompt},
],
"temperature": 0.1,
"max_tokens": 8192,
"max_tokens": max_tokens,
}
# Structured output (response_format) only for native OpenAI — upstream providers routed via
# Kilo/OpenRouter may not support it (e.g. StepFun returns "structured_outputs is not supported")