agent/translate.py

"""agent/translate.py — segment-map translation of a research brief.

Key design: extract only prose strings into a flat {id: text} map, translate
that map in ONE LLM call, then re-inject by path.  Enums, numbers, tickers,
evidence_snippet, and quarter_deltas are structurally never touched.
"""
from __future__ import annotations

import copy
import json
import sys

from langchain_core.messages import HumanMessage, SystemMessage

from agent.cost_log import compute_run_cost
from agent.graph import _extract_json
from agent.llm import RunConfig, default_config, make_chat_model
from agent.prompts import (
    NEVER_TRANSLATE_FIELDS,
    PROSE_FIELDS,
    PROSE_LIST_FIELDS,
    SKIP_SUBTREES,
)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _collect_segments(node, path: tuple = ()) -> list[tuple[tuple, str]]:
    """Walk *node* recursively and return a list of (path, text) pairs for
    every prose string that should be translated.

    Rules (applied in order for dict keys):
      - key in SKIP_SUBTREES          → skip subtree entirely
      - key in NEVER_TRANSLATE_FIELDS → skip (don't collect, don't recurse)
      - key in PROSE_FIELDS           → collect if value is a non-empty string
      - key in PROSE_LIST_FIELDS      → collect each non-empty string item
      - otherwise                     → recurse into value
    """
    if isinstance(node, dict):
        result: list[tuple[tuple, str]] = []
        for key, value in node.items():
            if key in SKIP_SUBTREES:
                continue
            if key in NEVER_TRANSLATE_FIELDS:
                continue
            if key in PROSE_FIELDS:
                if isinstance(value, str) and value:
                    result.append((path + (key,), value))
            elif key in PROSE_LIST_FIELDS:
                if isinstance(value, list):
                    for i, item in enumerate(value):
                        if isinstance(item, str) and item:
                            result.append((path + (key, i), item))
            else:
                result.extend(_collect_segments(value, path + (key,)))
        return result

    if isinstance(node, list):
        result = []
        for i, item in enumerate(node):
            result.extend(_collect_segments(item, path + (i,)))
        return result

    return []


def _set_at_path(obj: dict, path: tuple, value: str) -> None:
    """Navigate *obj* using *path* (mixed str/int keys) and set the leaf.

    Silently returns without error if any intermediate key is missing.
    Example: path=("bull_points", 0, "text") → obj["bull_points"][0]["text"] = value
    """
    if not path:
        return
    current = obj
    for step in path[:-1]:
        try:
            current = current[step]
        except (KeyError, IndexError, TypeError):
            return
    try:
        current[path[-1]] = value
    except (KeyError, IndexError, TypeError):
        return


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def translate_brief(
    brief: dict, target_lang: str, config: RunConfig | None = None
) -> tuple[dict, dict | None]:
    """Translate all prose fields in *brief* into *target_lang*.

    Returns ``(translated_brief, usage_dict)`` on success, or
    ``(original_brief, None)`` on any failure.  Never raises.
    """
    try:
        # Step 1 — collect prose segments
        segments = _collect_segments(brief)
        if not segments:
            return (brief, None)

        # Step 2 — build flat id→text payload
        payload: dict[str, str] = {
            str(i): text for i, (_, text) in enumerate(segments)
        }

        # Step 3 — create LLM (max_tokens=8192 is mandatory to avoid truncation)
        cfg = config or default_config()
        llm = make_chat_model(cfg, max_retries=2, max_tokens=8192)

        # Step 4 — system prompt
        system_prompt = (
            f"You are a financial translator. The user message is a JSON object "
            f"mapping ids to English equity-research text.\n"
            f"Translate every value into {target_lang}.\n"
            f"Rules:\n"
            f"- Return ONLY a JSON object with exactly the same keys; values are the translations.\n"
            f"- Professional equity-research register "
            f"(for French: vouvoiement, standard finance vocabulary).\n"
            f"- Keep verbatim: all numbers, percentages, currency amounts, ticker symbols, "
            f"company names, period labels (e.g. \"Q3 2025\"), and untranslatable finance terms "
            f"(EPS, FCF, capex, guidance, EBIT, EBITDA, YoY, QoQ).\n"
            f"- Plain text only — no HTML, no markdown, no commentary, no code fences."
        )

        # Step 5 — invoke
        resp = llm.invoke([
            SystemMessage(content=system_prompt),
            HumanMessage(content=json.dumps(payload, ensure_ascii=False)),
        ])

        # Step 6 — parse response
        raw_data: dict = json.loads(_extract_json(resp.content))

        # Step 7 — quality gate: at least 80% of ids must have a non-empty translation
        expected = len(payload)
        present = sum(
            1 for k in payload
            if isinstance(raw_data.get(k), str) and raw_data[k]
        )
        if present < 0.8 * expected:
            print(
                f"[translate] quality gate failed: {present}/{expected} ids translated "
                f"(threshold 80%)",
                file=sys.stderr,
            )
            return (brief, None)

        # Step 8 — deep-copy the brief
        translated = copy.deepcopy(brief)

        # Step 9 — re-inject translations by path
        for i, (path, _) in enumerate(segments):
            translation = raw_data.get(str(i))
            if isinstance(translation, str) and translation:
                _set_at_path(translated, path, translation)
            # missing ids → keep original English text (silent fallthrough)

        # Step 10 — tag the language
        translated["language"] = target_lang

        # Step 11 — compute cost
        usage = compute_run_cost([resp], cfg.model)

        # Step 12 — return
        return (translated, usage)

    except Exception as e:
        print(f"[translate] {e}", file=sys.stderr)
        return (brief, None)