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:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user