api/llm_agent.py

from __future__ import annotations

import argparse
import json
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Union

from json_repair import repair_json
from langchain_huggingface import HuggingFacePipeline

from utils.components import Lexicon, get_punctuation, join_tokens, strip_punctuation, normalize_candidates
from utils.controlled_mode import ControlledModeService, read_controlled_input as read_controlled_input_file
from utils.distractor_utils import (
    ensure_distractor_count,
    extract_forced_target,
    reattach_punctuation_to_output,
    sorted_distractor_keys,
)
from utils.runtime_config import build_llm_runtime_config
from utils.surprisal import SurprisalScorer, SurprisalThresholdConfig, SurprisalThresholdFilter


@dataclass
class AgentConfig:
    model_id: str
    lexicon_path: str
    punctuations: Sequence[str]
    # generation params
    max_new_tokens: int = 128
    temperature: float = 0.2
    top_p: float = 0.9
    do_sample: bool = True
    return_full_text: bool = False
    # lexicon params
    min_candidates: int = 10
    max_candidates: int = 20
    num_distractors: int = 3
    num_workers: int = 1
    # surprisal threshold params
    apply_surprisal_threshold: bool = False
    min_abs: Optional[float] = None
    min_delta: float = 0.0
    absolute_threshold_only: bool = False
    surprisal_device: Optional[str] = None
    # token joining strategy (e.g., "" for CJK char tokens, " " for space-separated)
    token_joiner: str = ""


class LLMAgent:
    """
    For each sentence (list of tokens), iterate i=1..len-1:
      - prefix_sentence = cfg.token_joiner.join(tokens[:i+1])  (ends at target word)
      - target_word = tokens[i]
      - candidates = lexicon neighbors for target_word (punct preserved)
      - LLM chooses distractors from candidates and returns schema-valid JSON
      - optional: enforce surprisal thresholds over chosen distractors
    """

    def __init__(self, prompt: Any, config: AgentConfig, llm: Optional[HuggingFacePipeline] = None) -> None:
        self.prompt = prompt
        self.cfg = config
        if not 1 <= int(self.cfg.num_distractors) <= 10:
            raise ValueError("num_distractors must be in [1, 10].")
        if int(self.cfg.num_workers) < 1:
            raise ValueError("num_workers must be >= 1.")

        self.lexicon = Lexicon(self.cfg.lexicon_path)
        self.puncts = set(self.cfg.punctuations)
        self._neighbor_cache: dict[str, tuple[str, ...]] = {}

        self.llm = llm or HuggingFacePipeline.from_model_id(
            model_id=self.cfg.model_id,
            task="text-generation",
            pipeline_kwargs={
                "max_new_tokens": self.cfg.max_new_tokens,
                "temperature": self.cfg.temperature,
                "top_p": self.cfg.top_p,
                "do_sample": self.cfg.do_sample,
                "return_full_text": self.cfg.return_full_text,
            },
        )
        self._configure_generation_padding()
        self._threshold_filter = SurprisalThresholdFilter(
            config=SurprisalThresholdConfig(enabled=False)
        )
        if self.cfg.apply_surprisal_threshold:
            scorer = SurprisalScorer(
                model_id=self.cfg.model_id,
                device=self.cfg.surprisal_device,
            )
            self._threshold_filter = SurprisalThresholdFilter(
                config=SurprisalThresholdConfig(
                    enabled=True,
                    min_abs=self.cfg.min_abs,
                    min_delta=self.cfg.min_delta,
                    absolute_threshold_only=self.cfg.absolute_threshold_only,
                ),
                scorer=scorer,
            )

    def _configure_generation_padding(self) -> None:
        """Avoid right-padding warnings for decoder-only batched generation."""
        pipe = getattr(self.llm, "pipeline", None)
        if pipe is None:
            return
        tokenizer = getattr(pipe, "tokenizer", None)
        model = getattr(pipe, "model", None)
        if tokenizer is None:
            return

        model_cfg = getattr(model, "config", None)
        is_decoder_only = not bool(getattr(model_cfg, "is_encoder_decoder", False))
        if not is_decoder_only:
            return

        if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
            tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "left"
        if model_cfg is not None and getattr(model_cfg, "pad_token_id", None) is None:
            model_cfg.pad_token_id = tokenizer.pad_token_id

    # ---------- helper function ----------
    def _get_cached_neighbors(self, core: str) -> tuple[str, ...]:
        if core not in self._neighbor_cache:
            raw = self.lexicon.get_neighbor(core, min_size=self.cfg.min_candidates, max_size=self.cfg.max_candidates)
            clean = normalize_candidates(raw)
            clean = [w for w in clean if w != core]
            self._neighbor_cache[core] = tuple(clean)
        return self._neighbor_cache[core]

    def _reattach_punct(self, out: Dict[str, Any], prefix: str, suffix: str) -> Dict[str, Any]:
        return reattach_punctuation_to_output(out, prefix=prefix, suffix=suffix)

    @staticmethod
    def _sorted_distractor_keys(out: Dict[str, Any]) -> List[str]:
        return sorted_distractor_keys(out)

    def _ensure_distractor_count(
        self,
        out: Dict[str, Any],
        *,
        target_word: str,
        target_core: str,
        candidate_pool: Sequence[str],
        allow_candidate_fill: bool = True,
    ) -> Dict[str, Any]:
        return ensure_distractor_count(
            out,
            num_distractors=self.cfg.num_distractors,
            target_word=target_word,
            target_core=target_core,
            candidate_pool=candidate_pool,
            puncts=self.puncts,
            allow_candidate_fill=allow_candidate_fill,
        )

    @staticmethod
    def _extract_forced_target(core: str) -> Optional[str]:
        return extract_forced_target(core)

    # ---------- candidate generation ----------

    def get_candidates(self, word: str):
        prefix, core, suffix = strip_punctuation(word, self.puncts)
        if not core:
            return prefix, core, suffix, ("X" * len(word),), None
        forced_target = self._extract_forced_target(core)
        if forced_target is not None:
            forced_len = max(1, len(forced_target))
            # Do not query lexicon / LLM for explicitly marked target tokens.
            return prefix, forced_target, suffix, tuple(), forced_len
        neighbors = self._get_cached_neighbors(core)
        return prefix, core, suffix, neighbors, None

    # ---------- LLM + parsing ----------

    def _invoke_raw(self, sentence_prefix: str, target_word: str, candidates: Sequence[str]) -> str:
        prompt_text = self.prompt.render_text(
            sentence_prefix=sentence_prefix,
            word=target_word,
            candidates=candidates,
        )
        resp = self.llm.invoke(prompt_text)
        return self._extract_generated_text(resp)

    @staticmethod
    def _extract_generated_text(resp: Any) -> str:
        if isinstance(resp, str):
            return resp
        if isinstance(resp, dict):
            for key in ("generated_text", "text", "content"):
                if key in resp and resp[key] is not None:
                    return str(resp[key])
        if isinstance(resp, list):
            if not resp:
                return ""
            first = resp[0]
            if isinstance(first, dict):
                return LLMAgent._extract_generated_text(first)
            if isinstance(first, list):
                return LLMAgent._extract_generated_text(first[0] if first else "")
        return str(resp)

    def _invoke_raw_batch(
        self,
        prompts: Sequence[str],
        *,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> List[str]:
        """Batch invoke HF pipeline; prefer dataset-style streaming on GPU."""
        prompt_list = list(prompts)
        total = len(prompt_list)
        if total == 0:
            return []

        def _report(done: int) -> None:
            if progress_callback is not None:
                progress_callback(done, total)

        pipe = getattr(self.llm, "pipeline", None)
        if pipe is not None:
            try:
                from transformers.pipelines.pt_utils import KeyDataset

                class _PromptDataset:
                    def __init__(self, rows: Sequence[str]) -> None:
                        self.rows = rows

                    def __len__(self) -> int:
                        return len(self.rows)

                    def __getitem__(self, idx: int) -> Dict[str, str]:
                        return {"text": self.rows[idx]}

                dataset = _PromptDataset(prompt_list)
                streamed = pipe(
                    KeyDataset(dataset, "text"),
                    batch_size=max(1, int(self.cfg.num_workers)),
                )
                out: List[str] = []
                for i, item in enumerate(streamed, start=1):
                    out.append(self._extract_generated_text(item))
                    _report(i)
                return out
            except Exception:
                pass

            try:
                batch_out = pipe(
                    prompt_list,
                    batch_size=max(1, int(self.cfg.num_workers)),
                )
                out = [self._extract_generated_text(item) for item in batch_out]
                for i in range(1, len(out) + 1):
                    _report(i)
                return out
            except Exception:
                pass

        out: List[str] = []
        for i, p in enumerate(prompt_list, start=1):
            out.append(self._extract_generated_text(self.llm.invoke(p)))
            _report(i)
        return out

    def _parse_to_dict(self, raw: str, repair: bool = True) -> Dict[str, Any]:
        try:
            obj = self.prompt.parser.parse(raw)
            return obj.model_dump()
        except Exception:
            if not repair:
                raise
        try:
            fixed = repair_json(raw)
            data = json.loads(fixed)
        except Exception:
            return {}
        if isinstance(data, list):
            for x in reversed(data):
                if isinstance(x, dict):
                    data = x
                    break
        if not isinstance(data, dict):
            return {}
        data.setdefault("source", "")
        data.setdefault("distractor1", "X")
        data.setdefault("distractor2", "X")
        data.setdefault("distractor3", None)
        try:
            obj = self.prompt.parser.pydantic_object.model_validate(data)
            return obj.model_dump()
        except Exception:
            return {}


    # ---------- public API ----------

    def iter_jobs(self, input_data: Sequence[Sequence[str]]) -> Iterable[Dict[str, Any]]:
        for s_idx, tokens in enumerate(input_data):
            for i in range(len(tokens)):
                target = tokens[i]
                sentence_prefix = join_tokens(tokens[: i + 1], join_with=self.cfg.token_joiner, puncts=self.puncts)
                context_prefix = join_tokens(tokens[:i], join_with=self.cfg.token_joiner, puncts=self.puncts)
                pfx, core, sfx, candidates, forced_len = self.get_candidates(target)
                yield {
                    "sentence_index": s_idx,
                    "word_index": i,
                    "sentence_prefix": sentence_prefix,
                    "context_prefix": context_prefix,
                    "target_word": target,
                    "target_core": core,
                    "punct_prefix": pfx,
                    "punct_suffix": sfx,
                    "candidates": candidates,
                    "forced_dummy_len": forced_len,
                }

    def read_controlled_input(self, path: Union[str, Path], split_on: Optional[str]) -> List[Dict[str, Any]]:
        return read_controlled_input_file(path, split_on)

    def _process_job(self, job: Dict[str, Any], *, repair: bool = True) -> Dict[str, Any]:
        forced_len = job.get("forced_dummy_len")
        if job["word_index"] == 0 or forced_len is not None:
            dummy = "X" * int(forced_len if forced_len is not None else len(job["target_word"]))
            out = {"source": job["target_core"], "distractor1": dummy}
            out = self._ensure_distractor_count(
                out,
                target_word=job["target_word"],
                target_core=job["target_core"],
                candidate_pool=job["candidates"],
                allow_candidate_fill=False,
            )
            if forced_len is not None:
                out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"])
        else:
            raw = self._invoke_raw(job["sentence_prefix"], job["target_core"], job["candidates"])
            out = self._parse_to_dict(raw, repair=repair)
            # Fallback if parsing totally failed
            if not out:
                dummy = "X" * len(job["target_word"])
                out = {
                    "source": job["target_core"],
                    "distractor1": dummy,
                    "distractor2": dummy,
                    "distractor3": None,
                }
            out = self._ensure_distractor_count(
                out,
                target_word=job["target_word"],
                target_core=job["target_core"],
                candidate_pool=job["candidates"],
                allow_candidate_fill=True,
            )
            out = self._threshold_filter.enforce_output_threshold(
                out=out,
                context_prefix=job["context_prefix"],
                target_core=job["target_core"],
                candidate_pool=job["candidates"],
                puncts=self.puncts,
                token_joiner=self.cfg.token_joiner,
            )
            out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"])

        # attach metadata (useful for reconstructing maze items)
        out.update(
            sentence_index=job["sentence_index"],
            word_index=job["word_index"],
            sentence_prefix=job["sentence_prefix"],
            target_word=job["target_word"],
        )
        return out

    def _process_job_with_raw(self, job: Dict[str, Any], raw: str, *, repair: bool = True) -> Dict[str, Any]:
        out = self._parse_to_dict(raw, repair=repair)
        if not out:
            dummy = "X" * len(job["target_word"])
            out = {
                "source": job["target_core"],
                "distractor1": dummy,
                "distractor2": dummy,
                "distractor3": None,
            }
        out = self._ensure_distractor_count(
            out,
            target_word=job["target_word"],
            target_core=job["target_core"],
            candidate_pool=job["candidates"],
            allow_candidate_fill=True,
        )
        out = self._threshold_filter.enforce_output_threshold(
            out=out,
            context_prefix=job["context_prefix"],
            target_core=job["target_core"],
            candidate_pool=job["candidates"],
            puncts=self.puncts,
            token_joiner=self.cfg.token_joiner,
        )
        out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"])
        out.update(
            sentence_index=job["sentence_index"],
            word_index=job["word_index"],
            sentence_prefix=job["sentence_prefix"],
            target_word=job["target_word"],
        )
        return out

    def run(
        self,
        input_data: Sequence[Sequence[str]],
        repair: bool = True,
        limit: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """Run over input_data and return a sentence-grouped list of JSON dict outputs."""
        run_start = time.perf_counter()
        jobs = list(self.iter_jobs(input_data))
        if limit is not None:
            jobs = jobs[:limit]
        if not jobs:
            return []

        total_sentences = len(input_data)
        print(f"[LLMAgent] Processing {len(jobs)} word-jobs from {total_sentences} sentences...")
        print(f"[LLMAgent] Generation batch size: {self.cfg.num_workers}")

        outputs: List[Dict[str, Any]] = []
        llm_jobs: List[Dict[str, Any]] = []
        llm_prompts: List[str] = []
        total_llm_seconds = 0.0

        for job in jobs:
            if job["word_index"] == 0 or job.get("forced_dummy_len") is not None:
                outputs.append(self._process_job(job, repair=repair))
            else:
                llm_jobs.append(job)
                llm_prompts.append(
                    self.prompt.render_text(
                        sentence_prefix=job["sentence_prefix"],
                        word=job["target_core"],
                        candidates=job["candidates"],
                    )
                )

        if llm_jobs:
            batch_size = max(1, int(self.cfg.num_workers))
            total_llm_jobs = len(llm_jobs)
            print(
                f"[LLMAgent] LLM generation jobs: {total_llm_jobs} "
                f"(dataset-style, batch_size={batch_size})"
            )

            llm_start = time.perf_counter()
            report_every = max(1, total_llm_jobs // 20)

            def _progress(done: int, total: int) -> None:
                if done == 1 or done == total or done % report_every == 0:
                    elapsed = time.perf_counter() - llm_start
                    avg = elapsed / max(1, done)
                    eta = avg * (total - done)
                    print(
                        f"[LLMAgent] Progress: {done}/{total} LLM jobs done, "
                        f"elapsed={elapsed:.2f}s, avg={avg:.2f}s/job, eta={eta:.2f}s"
                    )

            raw_outputs = self._invoke_raw_batch(llm_prompts, progress_callback=_progress)
            if len(raw_outputs) != len(llm_jobs):
                raw_outputs = [
                    self._invoke_raw(j["sentence_prefix"], j["target_core"], j["candidates"])
                    for j in llm_jobs
                ]
            total_llm_seconds = time.perf_counter() - llm_start

            for job, raw in zip(llm_jobs, raw_outputs):
                outputs.append(self._process_job_with_raw(job, raw, repair=repair))

        outputs.sort(key=lambda x: (x["sentence_index"], x["word_index"]))
        grouped: Dict[int, List[Dict[str, Any]]] = {}
        for out in outputs:
            grouped.setdefault(int(out["sentence_index"]), []).append(out)
        run_seconds = time.perf_counter() - run_start
        if llm_jobs:
            print(
                f"[LLMAgent] Timing summary: total={run_seconds:.2f}s, "
                f"llm_calls={total_llm_seconds:.2f}s, "
                f"avg_llm_call={total_llm_seconds / max(1, len(llm_jobs)):.2f}s"
            )
        else:
            print(f"[LLMAgent] Timing summary: total={run_seconds:.2f}s")
        return [{"sentence_index": s_idx, "words": grouped[s_idx]} for s_idx in sorted(grouped.keys())]

    def run_controlled_experiment(
        self,
        controlled_items: Sequence[Dict[str, Any]],
        repair: bool = True,
        limit: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """
        Controlled mode:
        - one shared distractor set per (item_id, word_index)
        - all conditions of the item reuse this shared set
        """
        service = ControlledModeService(
            lexicon=self.lexicon,
            puncts=self.puncts,
            num_distractors=self.cfg.num_distractors,
            token_joiner=self.cfg.token_joiner,
            min_candidates=self.cfg.min_candidates,
            max_candidates=self.cfg.max_candidates,
            apply_surprisal_threshold=self.cfg.apply_surprisal_threshold,
            threshold_filter=self._threshold_filter,
            get_candidates=self.get_candidates,
            invoke_raw=self._invoke_raw,
            parse_to_dict=self._parse_to_dict,
            ensure_distractor_count=self._ensure_distractor_count,
            reattach_punct=self._reattach_punct,
            sorted_distractor_keys=self._sorted_distractor_keys,
        )
        return service.run(controlled_items=controlled_items, repair=repair, limit=limit)

    @staticmethod
    def save_jsonl(records: Sequence[Dict[str, Any]], path: Union[str, Path]) -> None:
        """Machine-friendly JSONL (one object per line)."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as f:
            for r in records:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")

    @staticmethod
    def save_pretty_json(records: Sequence[Dict[str, Any]], path: Union[str, Path]) -> None:
        """Human-friendly JSON (pretty-printed)."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as f:
            json.dump(records, f, ensure_ascii=False, indent=2)

# -------------------------
# minimal usage example
# -------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="LLM Maze distractor generation (lexicon + LLM selector).")
    parser.add_argument("--config-path", default="config.yaml", help="Path to YAML config file.")
    parser.add_argument("--language-code", default=None, help="Override LANGUAGE_CODE from config.")
    parser.add_argument("--min-abs", type=float, default=None, help="Minimum absolute distractor surprisal.")
    parser.add_argument("--min-delta", type=float, default=None, help="Minimum surprisal delta over target.")
    parser.add_argument(
        "--absolute-threshold-only",
        action=argparse.BooleanOptionalAction,
        default=None,
        help="Use only min_abs as hard threshold (ignore target + min_delta).",
    )
    parser.add_argument("--surprisal-device", default=None, help='Device for surprisal scorer, e.g. "cuda" or "cpu".')
    parser.add_argument("--limit", type=int, default=None, help="Optional max number of word-jobs to run.")
    parser.add_argument("--output-path", default=None, help="Optional explicit output JSON path.")
    parser.add_argument("--model-id", default=None, help="Override MODEL_ID from config.")
    parser.add_argument("--input-data-path", default=None, help="Optional explicit input data path override.")
    parser.add_argument("--lexicon-path", default=None, help="Optional explicit lexicon path override.")
    parser.add_argument("--template-dir", default=None, help="Optional explicit template directory override.")
    parser.add_argument(
        "--processing-mode",
        default=None,
        choices=["naturalistic_reading", "controlled_experiment"],
        help="Processing mode: naturalistic_reading or controlled_experiment.",
    )
    parser.add_argument("--word-separator", default=None, help="Override WORD_SEPARATOR from config.")
    parser.add_argument("--num-distractors", type=int, default=None, help="Number of distractors per word (1-10).")
    parser.add_argument("--num-workers", type=int, default=None, help="Parallel workers for word-level processing (>=1).")
    args = parser.parse_args()

    from utils.components import load_config, read_sentences_input
    from utils.maze_prompt import MazeLLMPrompt  

    yaml_cfg = load_config(args.config_path)
    runtime_cfg = build_llm_runtime_config(
        yaml_cfg=yaml_cfg,
        overrides={
            "AGENT_TYPE": "llm",
            "LANGUAGE_CODE": args.language_code,
            "MODEL_ID": args.model_id,
            "INPUT_DATA_PATH": args.input_data_path,
            "LEXICON_PATH": args.lexicon_path,
            "TEMPLATE_DIR": args.template_dir,
            "OUTPUT_PATH": args.output_path,
            "PROCESSING_MODE": args.processing_mode,
            "WORD_SEPARATOR": args.word_separator,
            "NUM_DISTRACTORS": args.num_distractors,
            "NUM_WORKERS": args.num_workers,
            "SURPRISAL_MIN_ABS": args.min_abs,
            "SURPRISAL_MIN_DELTA": args.min_delta,
            "SURPRISAL_ABSOLUTE_THRESHOLD_ONLY": args.absolute_threshold_only,
            "SURPRISAL_DEVICE": args.surprisal_device,
        },
    )

    language_code = runtime_cfg.language_code
    agent_type = runtime_cfg.agent_type
    model_id = runtime_cfg.model_id
    processing_mode = runtime_cfg.processing_mode
    main_path = Path(runtime_cfg.template_dir)
    word_separator = runtime_cfg.word_separator
    puncts = get_punctuation(language_code)
    if not main_path.exists():
        raise ValueError(f"Template directory not found for LANGUAGE_CODE='{language_code}': {main_path}")
    if not Path(runtime_cfg.lexicon_path).exists():
        raise ValueError(f"Lexicon file not found: {runtime_cfg.lexicon_path}")
    if not Path(runtime_cfg.input_data_path).exists():
        raise ValueError(f"Input data file not found: {runtime_cfg.input_data_path}")

    maze_prompt = MazeLLMPrompt(
        path_to_user_template=main_path / "base.txt",
        path_to_extension_template=main_path / "extension.txt",
        path_to_system_template=main_path / "system.txt",
    )

    agent_cfg = AgentConfig(
        model_id=model_id ,
        lexicon_path=runtime_cfg.lexicon_path,
        punctuations=puncts,
        num_distractors=runtime_cfg.num_distractors,
        num_workers=runtime_cfg.num_workers,
        apply_surprisal_threshold=runtime_cfg.apply_surprisal_threshold,
        min_abs=runtime_cfg.min_abs,
        min_delta=runtime_cfg.min_delta,
        absolute_threshold_only=runtime_cfg.absolute_threshold_only,
        surprisal_device=runtime_cfg.surprisal_device,
        token_joiner="" if str(word_separator).strip() == "" else str(word_separator),
    )

    agent = LLMAgent(prompt=maze_prompt, config=agent_cfg)

    main_start = time.perf_counter()
    if processing_mode == "controlled_experiment":
        controlled_items = agent.read_controlled_input(
            runtime_cfg.input_data_path,
            split_on=word_separator,
        )
        outputs = agent.run_controlled_experiment(controlled_items, limit=args.limit)
    else:
        input_data = read_sentences_input(runtime_cfg.input_data_path, split_on=word_separator)
        outputs = agent.run(input_data, limit=args.limit)
    generation_elapsed = time.perf_counter() - main_start
    print(
        f"[LLMAgent] Generation finished in {generation_elapsed:.2f}s "
        f"(mode={processing_mode}, records={len(outputs)})."
    )
    output_path = args.output_path or runtime_cfg.output_path
    agent.save_pretty_json(outputs, output_path)
    total_elapsed = time.perf_counter() - main_start
    print(f"Saved output to: {output_path}")
    print(f"[LLMAgent] Total elapsed (generation + save): {total_elapsed:.2f}s")