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:
195
backend/src/plugins/translate/__tests__/test_token_budget.py
Normal file
195
backend/src/plugins/translate/__tests__/test_token_budget.py
Normal file
@@ -0,0 +1,195 @@
|
||||
# #region TestTokenBudget [C:3] [TYPE Module] [SEMANTICS test, token, budget, estimation, batch, translate]
|
||||
# @BRIEF Verify estimate_token_budget contracts — safe batch sizing, auto-reduction, warning generation.
|
||||
# @RELATION BINDS_TO -> [estimate_token_budget:Module]
|
||||
# @TEST_EDGE: empty_rows — empty source rows returns batch_size_adjusted=1
|
||||
# @TEST_EDGE: small_rows — short text fits in a single batch at requested size
|
||||
# @TEST_EDGE: large_rows — long text causes batch size reduction
|
||||
# @TEST_EDGE: multi_language — more target languages increases output estimate
|
||||
# @TEST_EDGE: context_columns — context columns increase input token estimate
|
||||
# @TEST_EDGE: dictionary_entries — glossary entries increase input token estimate
|
||||
# @TEST_EDGE: auto_calc — no batch_size specified, auto-calculates max safe
|
||||
# @TEST_EDGE: exact_fit — exactly fits context window, batch_adjusted == requested
|
||||
# @TEST_EDGE: conservative_min — even with huge rows, batch_size_adjusted >= 1
|
||||
# @TEST_INVARIANT: batch_size_adjusted >= 1 always
|
||||
# @TEST_INVARIANT: max_output_needed between MIN_MAX_TOKENS(4096) and max_output_tokens(8192)
|
||||
# @TEST_INVARIANT: warning is None when batch fits, str when reduced
|
||||
|
||||
from src.plugins.translate._token_budget import DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_OUTPUT_TOKENS, estimate_token_budget
|
||||
|
||||
|
||||
# region _make_row [TYPE Function]
|
||||
# @BRIEF Create a test source row dict.
|
||||
def _make_row(text: str, **context) -> dict:
|
||||
row = {"source_text": text, "row_index": "0"}
|
||||
row.update(context)
|
||||
return row
|
||||
# endregion _make_row
|
||||
|
||||
|
||||
# region TestTokenBudget [TYPE Class]
|
||||
# @BRIEF Test suite for estimate_token_budget.
|
||||
class TestTokenBudget:
|
||||
|
||||
# region test_small_rows_fit_at_requested_size [TYPE Function]
|
||||
# @BRIEF Short text rows fill the requested batch_size without reduction.
|
||||
def test_small_rows_fit_at_requested_size(self):
|
||||
"""Short text of ~50 chars should fit 50-row batch easily."""
|
||||
rows = [_make_row("Hello world, this is a short text for translation.")] * 50
|
||||
result = estimate_token_budget(
|
||||
source_rows=rows,
|
||||
target_languages=["ru"],
|
||||
batch_size=50,
|
||||
)
|
||||
assert result["batch_size_adjusted"] == 50
|
||||
assert result["estimated_input_tokens"] > 0
|
||||
assert result["estimated_output_tokens"] > 0
|
||||
assert result["max_output_needed"] >= 4096
|
||||
assert result["max_output_needed"] <= DEFAULT_MAX_OUTPUT_TOKENS
|
||||
assert result["warning"] is None
|
||||
# endregion test_small_rows_fit_at_requested_size
|
||||
|
||||
# region test_large_rows_reduce_batch_size [TYPE Function]
|
||||
# @BRIEF Very long text rows force batch size reduction to fit context window.
|
||||
def test_large_rows_reduce_batch_size(self):
|
||||
"""~10000 char rows should cause batch reduction from 50 to a smaller number."""
|
||||
long_text = "X" * 10000 # ~4545 tokens each at 2.2 chars/token
|
||||
rows = [_make_row(long_text)] * 50
|
||||
result = estimate_token_budget(
|
||||
source_rows=rows,
|
||||
target_languages=["ru", "en"],
|
||||
batch_size=50,
|
||||
)
|
||||
# Each row is ~4545 tokens input, plus ~2400 output (2 langs * 60 * 20 + 2000)
|
||||
# With 50 rows: ~227K tokens — way over 64K. Should reduce significantly.
|
||||
assert result["batch_size_adjusted"] < 50
|
||||
assert result["batch_size_adjusted"] >= 1
|
||||
assert result["warning"] is not None
|
||||
assert "Reduced batch size" in result["warning"]
|
||||
# endregion test_large_rows_reduce_batch_size
|
||||
|
||||
# region test_multi_language_increases_output_estimate [TYPE Function]
|
||||
# @BRIEF More target languages increase the estimated output tokens.
|
||||
def test_multi_language_increases_output_estimate(self):
|
||||
"""4 target languages should have higher output estimate than 1."""
|
||||
rows = [_make_row("Short text.")] * 10
|
||||
|
||||
single_lang = estimate_token_budget(rows, ["ru"], batch_size=10)
|
||||
multi_lang = estimate_token_budget(rows, ["ru", "en", "fr", "de"], batch_size=10)
|
||||
|
||||
assert multi_lang["estimated_output_tokens"] > single_lang["estimated_output_tokens"]
|
||||
# Both should fit without reduction
|
||||
assert single_lang["batch_size_adjusted"] == 10
|
||||
assert multi_lang["batch_size_adjusted"] == 10
|
||||
# endregion test_multi_language_increases_output_estimate
|
||||
|
||||
# region test_context_columns_increase_input [TYPE Function]
|
||||
# @BRIEF Context columns add to the input token estimate.
|
||||
def test_context_columns_increase_input(self):
|
||||
"""Rows with context columns should have higher input tokens."""
|
||||
rows = [_make_row("Short text.")] * 10
|
||||
|
||||
no_context = estimate_token_budget(rows, ["ru"], batch_size=10)
|
||||
with_context = estimate_token_budget(
|
||||
rows, ["ru"], batch_size=10,
|
||||
context_columns=["description", "category"],
|
||||
)
|
||||
|
||||
assert with_context["estimated_input_tokens"] > no_context["estimated_input_tokens"]
|
||||
# endregion test_context_columns_increase_input
|
||||
|
||||
# region test_dictionary_entries_increase_input [TYPE Function]
|
||||
# @BRIEF Dictionary entries add to the input token estimate.
|
||||
def test_dictionary_entries_increase_input(self):
|
||||
"""Dictionary entries should increase estimated_input_tokens."""
|
||||
rows = [_make_row("Short text.")] * 10
|
||||
dict_entries = [{"source_term": f"term_{i}", "target_term": f"trans_{i}"} for i in range(100)]
|
||||
|
||||
no_dict = estimate_token_budget(rows, ["ru"], batch_size=10)
|
||||
with_dict = estimate_token_budget(rows, ["ru"], batch_size=10, dictionary_entries=dict_entries)
|
||||
|
||||
assert with_dict["estimated_input_tokens"] > no_dict["estimated_input_tokens"]
|
||||
# endregion test_dictionary_entries_increase_input
|
||||
|
||||
# region test_auto_calc_finds_safe_size [TYPE Function]
|
||||
# @BRIEF Without batch_size specified, auto-calculate max safe batch.
|
||||
def test_auto_calc_finds_safe_size(self):
|
||||
"""Auto-calculation should find the max rows that fit context window."""
|
||||
rows = [_make_row("A" * 500)] * 100 # ~227 tokens/row
|
||||
|
||||
result = estimate_token_budget(
|
||||
source_rows=rows,
|
||||
target_languages=["ru", "en"],
|
||||
batch_size=None, # auto-calc
|
||||
)
|
||||
|
||||
assert result["batch_size_adjusted"] >= 1
|
||||
assert result["batch_size_adjusted"] <= 100
|
||||
# Verify it found a reasonable size
|
||||
assert result["estimated_input_tokens"] + result["max_output_needed"] <= DEFAULT_CONTEXT_WINDOW
|
||||
# endregion test_auto_calc_finds_safe_size
|
||||
|
||||
# region test_empty_rows_returns_minimum [TYPE Function]
|
||||
# @BRIEF Empty source rows returns batch_size_adjusted=1 and minimum estimates.
|
||||
def test_empty_rows_returns_minimum(self):
|
||||
"""Empty rows list should return batch_size_adjusted=1 with min estimates."""
|
||||
result = estimate_token_budget(
|
||||
source_rows=[],
|
||||
target_languages=["ru"],
|
||||
batch_size=10,
|
||||
)
|
||||
assert result["batch_size_adjusted"] == 1
|
||||
assert result["estimated_input_tokens"] >= 300 # at least prompt base
|
||||
assert result["estimated_output_tokens"] >= 2000 # at least reasoning overhead
|
||||
# endregion test_empty_rows_returns_minimum
|
||||
|
||||
# region test_exact_fit_no_warning [TYPE Function]
|
||||
# @BRIEF When batch fits exactly, no warning is generated.
|
||||
def test_exact_fit_no_warning(self):
|
||||
"""Small batch that fits easily should have no warning."""
|
||||
rows = [_make_row("Short text.")] * 5
|
||||
result = estimate_token_budget(rows, ["ru"], batch_size=5)
|
||||
assert result["warning"] is None
|
||||
assert result["batch_size_adjusted"] == 5
|
||||
# endregion test_exact_fit_no_warning
|
||||
|
||||
# region test_huge_rows_still_return_min_one [TYPE Function]
|
||||
# @BRIEF Even with massive rows, batch_size_adjusted is at least 1.
|
||||
def test_huge_rows_still_return_min_one(self):
|
||||
"""Massive rows that exceed context window still return batch_size_adjusted >= 1."""
|
||||
huge_text = "X" * 500000 # ~227K tokens — exceeds context window
|
||||
rows = [_make_row(huge_text)] * 10
|
||||
result = estimate_token_budget(rows, ["ru", "en", "fr"], batch_size=10)
|
||||
assert result["batch_size_adjusted"] >= 1
|
||||
# endregion test_huge_rows_still_return_min_one
|
||||
|
||||
# region test_max_output_needed_bounds [TYPE Function]
|
||||
# @BRIEF max_output_needed stays within reasonable bounds.
|
||||
def test_max_output_needed_bounds(self):
|
||||
"""max_output_needed is between MIN_MAX_TOKENS and max_output_tokens."""
|
||||
rows = [_make_row("Short text.")] * 5
|
||||
result = estimate_token_budget(rows, ["ru", "en"], batch_size=5)
|
||||
assert 4096 <= result["max_output_needed"] <= DEFAULT_MAX_OUTPUT_TOKENS
|
||||
# endregion test_max_output_needed_bounds
|
||||
|
||||
# region test_single_row_never_reduced [TYPE Function]
|
||||
# @BRIEF Batch of 1 should never be reduced (batch_size_adjusted >= requested when requested=1).
|
||||
def test_single_row_never_reduced(self):
|
||||
"""Even a very long single row should still fit (batch_size_adjusted=1)."""
|
||||
long_text = "X" * 100000 # large but still fits context if alone
|
||||
rows = [_make_row(long_text)]
|
||||
result = estimate_token_budget(rows, ["ru", "en", "fr", "de"], batch_size=1)
|
||||
assert result["batch_size_adjusted"] == 1
|
||||
# endregion test_single_row_never_reduced
|
||||
|
||||
# region test_no_target_languages_defaults_to_en [TYPE Function]
|
||||
# @BRIEF When target_languages is empty, defaults to ["en"].
|
||||
def test_no_target_languages_defaults_to_en(self):
|
||||
"""Empty target_languages defaults to ['en'] without error."""
|
||||
rows = [_make_row("Short text.")] * 5
|
||||
result = estimate_token_budget(rows, [], batch_size=5)
|
||||
assert result["batch_size_adjusted"] == 5
|
||||
assert result["estimated_output_tokens"] > 0
|
||||
# endregion test_no_target_languages_defaults_to_en
|
||||
|
||||
# endregion TestTokenBudget
|
||||
# endregion TestTokenBudget
|
||||
201
backend/src/plugins/translate/_token_budget.py
Normal file
201
backend/src/plugins/translate/_token_budget.py
Normal file
@@ -0,0 +1,201 @@
|
||||
# #region estimate_token_budget [C:3] [TYPE Module] [SEMANTICS translate, token, budget, estimation, llm]
|
||||
# @BRIEF Calculate safe batch_size and max_tokens for LLM translation calls based on actual content length and model context window limits.
|
||||
# @LAYER: Domain
|
||||
# @RELATION DEPENDS_ON -> [TranslationPreview:Module]
|
||||
# @RELATION DEPENDS_ON -> [TranslationExecutor:Module]
|
||||
# @RATIONALE: Prevents LLM truncation (finish_reason=length) by sizing batches within context limits.
|
||||
# DeepSeek v4 Flash supports up to 64K context window; output is limited by max_tokens.
|
||||
# @REJECTED: External tokenizer library — would introduce heavy dependency for estimation only.
|
||||
# Fixed batch_size of 50 — causes truncation on long-content rows.
|
||||
|
||||
# #region DEFAULT_CONTEXT_WINDOW [TYPE Constant]
|
||||
# @BRIEF Default context window for DeepSeek v4 Flash: up to 64K tokens.
|
||||
DEFAULT_CONTEXT_WINDOW = 64000
|
||||
# #endregion DEFAULT_CONTEXT_WINDOW
|
||||
|
||||
# #region DEFAULT_MAX_OUTPUT_TOKENS [TYPE Constant]
|
||||
# @BRIEF Default max_tokens setting for LLM output (8192 tokens).
|
||||
DEFAULT_MAX_OUTPUT_TOKENS = 8192
|
||||
# #endregion DEFAULT_MAX_OUTPUT_TOKENS
|
||||
|
||||
# #region REASONING_OVERHEAD [TYPE Constant]
|
||||
# @BRIEF CoT reasoning overhead tokens for DeepSeek models (~2000 tokens for chain-of-thought).
|
||||
REASONING_OVERHEAD = 2000
|
||||
# #endregion REASONING_OVERHEAD
|
||||
|
||||
# #region OUTPUT_PER_ROW_PER_LANG [TYPE Constant]
|
||||
# @BRIEF Estimated output tokens per row per language in JSON response format.
|
||||
OUTPUT_PER_ROW_PER_LANG = 60
|
||||
# #endregion OUTPUT_PER_ROW_PER_LANG
|
||||
|
||||
# #region PROMPT_BASE_TOKENS [TYPE Constant]
|
||||
# @BRIEF Base tokens for system prompt + instructions + JSON format specification.
|
||||
PROMPT_BASE_TOKENS = 300
|
||||
# #endregion PROMPT_BASE_TOKENS
|
||||
|
||||
# #region DICT_TOKENS_PER_ENTRY [TYPE Constant]
|
||||
# @BRIEF Estimated tokens per dictionary entry in the glossary section.
|
||||
DICT_TOKENS_PER_ENTRY = 20
|
||||
# #endregion DICT_TOKENS_PER_ENTRY
|
||||
|
||||
# #region DICT_TOKENS_MAX [TYPE Constant]
|
||||
# @BRIEF Cap for dictionary tokens in estimation to avoid overestimating.
|
||||
DICT_TOKENS_MAX = 5000
|
||||
# #endregion DICT_TOKENS_MAX
|
||||
|
||||
# #region CHARS_PER_TOKEN_MIXED [TYPE Constant]
|
||||
# @BRIEF Characters per token for mixed Russian/English text (empirical ~2.2 for mixed, ~2 for Russian, ~4 for English).
|
||||
CHARS_PER_TOKEN_MIXED = 2.2
|
||||
# #endregion CHARS_PER_TOKEN_MIXED
|
||||
|
||||
# #region MIN_MAX_TOKENS [TYPE Constant]
|
||||
# @BRIEF Minimum max_tokens to allow (4096 ensures reasonable output even for small batches).
|
||||
MIN_MAX_TOKENS = 4096
|
||||
# #endregion MIN_MAX_TOKENS
|
||||
|
||||
# #region MAX_OUTPUT_HEADROOM [TYPE Constant]
|
||||
# @BRIEF Extra headroom added to max_output_needed beyond the estimate (1000 tokens buffer).
|
||||
MAX_OUTPUT_HEADROOM = 1000
|
||||
# #endregion MAX_OUTPUT_HEADROOM
|
||||
|
||||
|
||||
# region _count_rows_that_fit [TYPE Function]
|
||||
# @BRIEF Count how many rows fit within the available input budget.
|
||||
# @PRE: input_per_row is non-empty; available_budget > 0.
|
||||
# @POST: Returns (safe_count, total_input_tokens).
|
||||
def _count_rows_that_fit(
|
||||
input_per_row: list[int],
|
||||
available_budget: int,
|
||||
) -> tuple[int, int]:
|
||||
"""Count consecutive rows that fit within available_budget.
|
||||
|
||||
Returns:
|
||||
(safe_count, total_input_tokens): Number of rows and their total tokens.
|
||||
"""
|
||||
running_total = 0
|
||||
safe_size = 0
|
||||
for tokens in input_per_row:
|
||||
if running_total + tokens + REASONING_OVERHEAD < available_budget:
|
||||
running_total += tokens
|
||||
safe_size += 1
|
||||
else:
|
||||
break
|
||||
safe_size = max(safe_size, 1)
|
||||
return safe_size, running_total
|
||||
# endregion _count_rows_that_fit
|
||||
|
||||
|
||||
# region estimate_token_budget [C:3] [TYPE Function]
|
||||
# @BRIEF Estimate token budget for a batch of source rows and return safe batch parameters.
|
||||
# @PRE: source_rows is a list of dicts (can be empty). target_languages is a non-empty list.
|
||||
# @POST: Returns dict with batch_size_adjusted, estimated_input_tokens, estimated_output_tokens, max_output_needed, warning.
|
||||
# @RATIONALE: Uses character-count heuristics (chars/2.2 for mixed text) since exact tokenization
|
||||
# depends on the LLM model. Estimates are intentionally conservative to prevent truncation.
|
||||
# @REJECTED: Using tiktoken or similar tokenizer — would introduce a heavy dependency and still
|
||||
# not match DeepSeek's tokenizer exactly.
|
||||
def estimate_token_budget(
|
||||
source_rows: list[dict],
|
||||
target_languages: list[str],
|
||||
source_column: str = "source_text",
|
||||
context_columns: list[str] | None = None,
|
||||
dictionary_entries: list | None = None,
|
||||
batch_size: int | None = None,
|
||||
context_window: int = DEFAULT_CONTEXT_WINDOW,
|
||||
max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
|
||||
) -> dict:
|
||||
"""Estimate token budget and return safe batch parameters.
|
||||
|
||||
Args:
|
||||
source_rows: List of row dicts with source text.
|
||||
target_languages: List of target language codes.
|
||||
source_column: Key for the source text in each row dict.
|
||||
context_columns: Optional list of keys for context columns.
|
||||
dictionary_entries: Optional list of dictionary entries for glossary.
|
||||
batch_size: Desired batch size. If None, auto-calculate max safe size.
|
||||
context_window: Model context window (default 64000 for DeepSeek v4 Flash).
|
||||
max_output_tokens: Hard max output tokens limit (default 8192).
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
batch_size_adjusted: Safe batch size (may be less than requested).
|
||||
estimated_input_tokens: Total estimated input tokens for the batch.
|
||||
estimated_output_tokens: Total estimated output tokens.
|
||||
max_output_needed: Recommended max_tokens for this batch.
|
||||
warning: str | None if batch was reduced.
|
||||
"""
|
||||
if not target_languages:
|
||||
target_languages = ["en"]
|
||||
|
||||
num_languages = len(target_languages)
|
||||
|
||||
# 1. Estimate tokens per row
|
||||
input_per_row = []
|
||||
limit = batch_size if batch_size else len(source_rows)
|
||||
for i, row in enumerate(source_rows):
|
||||
if i >= limit:
|
||||
break
|
||||
text = str(row.get(source_column, "") or "")
|
||||
# Use ~2.2 chars per token for mixed Russian/English text
|
||||
estimated_tokens = max(1, int(len(text) / CHARS_PER_TOKEN_MIXED))
|
||||
# Add context columns
|
||||
if context_columns:
|
||||
for col in context_columns:
|
||||
val = str(row.get(col, "") or "")
|
||||
estimated_tokens += max(1, int(len(val) / CHARS_PER_TOKEN_MIXED))
|
||||
input_per_row.append(estimated_tokens)
|
||||
|
||||
# 2. Calculate dictionary tokens
|
||||
dict_tokens = 0
|
||||
if dictionary_entries:
|
||||
dict_tokens = min(
|
||||
len(dictionary_entries) * DICT_TOKENS_PER_ENTRY,
|
||||
DICT_TOKENS_MAX,
|
||||
)
|
||||
|
||||
prompt_tokens = PROMPT_BASE_TOKENS + dict_tokens
|
||||
available_input_budget = context_window - max_output_tokens
|
||||
|
||||
# 3. Calculate safe batch size using shared helper
|
||||
safe_size, input_total = _count_rows_that_fit(input_per_row, available_input_budget)
|
||||
estimated_input = prompt_tokens + input_total
|
||||
|
||||
# If batch_size was specified and we reduced it, recalculate
|
||||
if batch_size and safe_size < batch_size:
|
||||
# Ensure at minimum one row even for huge content
|
||||
safe_size = max(safe_size, 1)
|
||||
_, truncated_total = _count_rows_that_fit(
|
||||
input_per_row[:batch_size], available_input_budget,
|
||||
)
|
||||
estimated_input = prompt_tokens + truncated_total
|
||||
|
||||
# 4. Estimate output tokens
|
||||
estimated_output = (
|
||||
safe_size * num_languages * OUTPUT_PER_ROW_PER_LANG + REASONING_OVERHEAD
|
||||
)
|
||||
|
||||
# 5. Calculate recommended max_tokens
|
||||
max_output_needed = min(
|
||||
estimated_output + MAX_OUTPUT_HEADROOM,
|
||||
context_window - estimated_input,
|
||||
)
|
||||
max_output_needed = max(max_output_needed, MIN_MAX_TOKENS)
|
||||
max_output_needed = min(max_output_needed, max_output_tokens)
|
||||
|
||||
# 6. Generate warning if batch was reduced
|
||||
warning = None
|
||||
if batch_size and safe_size < batch_size:
|
||||
total_estimated = estimated_input + max_output_needed
|
||||
warning = (
|
||||
f"Reduced batch size from {batch_size} to {safe_size} "
|
||||
f"(estimated {total_estimated} tokens vs {context_window} window)"
|
||||
)
|
||||
|
||||
return {
|
||||
"batch_size_adjusted": safe_size,
|
||||
"estimated_input_tokens": estimated_input,
|
||||
"estimated_output_tokens": estimated_output,
|
||||
"max_output_needed": max_output_needed,
|
||||
"warning": warning,
|
||||
}
|
||||
# endregion estimate_token_budget
|
||||
# #endregion estimate_token_budget
|
||||
@@ -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
|
||||
from .preview import DEFAULT_EXECUTION_PROMPT_TEMPLATE
|
||||
from .prompt_builder import ContextAwarePromptBuilder
|
||||
@@ -589,12 +590,31 @@ class TranslationExecutor:
|
||||
|
||||
# Process rows needing LLM translation
|
||||
if rows_for_llm:
|
||||
# Check token budget for this batch to determine safe max_tokens
|
||||
token_budget = estimate_token_budget(
|
||||
source_rows=rows_for_llm,
|
||||
target_languages=target_languages,
|
||||
source_column="source_text",
|
||||
context_columns=None, # Context is embedded in dict, not separate column
|
||||
dictionary_entries=dict_matches,
|
||||
batch_size=len(rows_for_llm),
|
||||
context_window=DEFAULT_CONTEXT_WINDOW,
|
||||
max_output_tokens=DEFAULT_MAX_OUTPUT_TOKENS,
|
||||
)
|
||||
if token_budget["warning"]:
|
||||
logger.explore("Token budget warning for batch", {
|
||||
"batch_id": batch_id,
|
||||
"batch_index": batch_index,
|
||||
"warning": token_budget["warning"],
|
||||
})
|
||||
|
||||
llm_result = self._call_llm_for_batch(
|
||||
job=job,
|
||||
run_id=run_id,
|
||||
batch_rows=rows_for_llm,
|
||||
dict_matches=dict_matches,
|
||||
batch_id=batch_id,
|
||||
max_tokens=token_budget["max_output_needed"],
|
||||
)
|
||||
result["successful"] += llm_result["successful"]
|
||||
result["failed"] += llm_result["failed"]
|
||||
@@ -630,6 +650,7 @@ class TranslationExecutor:
|
||||
batch_rows: list[dict[str, Any]],
|
||||
dict_matches: list[dict[str, Any]],
|
||||
batch_id: str,
|
||||
max_tokens: int = 8192,
|
||||
) -> dict[str, int]:
|
||||
with belief_scope("TranslationExecutor._call_llm_for_batch"):
|
||||
# Build dictionary section using ContextAwarePromptBuilder
|
||||
@@ -689,7 +710,7 @@ class TranslationExecutor:
|
||||
|
||||
for attempt in range(1, MAX_RETRIES_PER_BATCH + 1):
|
||||
try:
|
||||
llm_response = self._call_llm(job, prompt)
|
||||
llm_response = self._call_llm(job, prompt, max_tokens=max_tokens)
|
||||
break
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
@@ -889,7 +910,7 @@ class TranslationExecutor:
|
||||
# @PRE: job has valid provider_id.
|
||||
# @POST: Returns raw LLM response string.
|
||||
# @SIDE_EFFECT: 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("TranslationExecutor._call_llm"):
|
||||
if not job.provider_id:
|
||||
raise ValueError("Job has no LLM provider configured")
|
||||
@@ -913,6 +934,7 @@ class TranslationExecutor:
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
provider_type=provider_type,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported provider type '{provider_type}'")
|
||||
@@ -930,6 +952,7 @@ class TranslationExecutor:
|
||||
model: str,
|
||||
prompt: str,
|
||||
provider_type: str = "openai",
|
||||
max_tokens: int = 8192,
|
||||
) -> str:
|
||||
with belief_scope("TranslationExecutor._call_openai_compatible"):
|
||||
import requests as http_requests
|
||||
@@ -946,7 +969,7 @@ class TranslationExecutor:
|
||||
{"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")
|
||||
|
||||
@@ -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