api/chat_agent.py

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

from json_repair import repair_json
from langchain_openai import ChatOpenAI

from utils.components import (
    Lexicon,
    get_punctuation,
    join_tokens,
    load_config,
    normalize_candidates,
    read_sentences_input,
    strip_punctuation,
)
from utils.surprisal import SurprisalScorer, SurprisalThresholdConfig, SurprisalThresholdFilter
from utils.distractor_utils import (
    build_dummy_output,
    ensure_distractor_count,
    extract_forced_target,
    reattach_punctuation_to_output,
)
from utils.maze_prompt import DistractorGeneratorPrompt, MazeChatPrompt
from utils.runtime_config import resolve_runtime_paths


@dataclass
class ChatAgentConfig:
    lexicon_path: Optional[str] = None
    language_code: Optional[str] = None
    punctuations: Sequence[str] = ()
    min_candidates: int = 10
    max_candidates: int = 20
    num_distractors: int = 3
    token_joiner: str = ""
    lexicon_mode: bool = False
    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
    model_id: str = "gpt-4o-mini"
    gen_temperature: float = 0.7
    sel_temperature: float = 0.2


class ChatAgent:
    def __init__(
        self,
        cfg: ChatAgentConfig,
        selector_prompt: Any,
        generator_prompt: Optional[Any] = None,
        *,
        lexicon_mode: Optional[bool] = None,
        chat_gen: Optional[ChatOpenAI] = None,
        chat_sel: Optional[ChatOpenAI] = None,
    ) -> None:
        self.cfg = cfg
        self.selector_prompt = selector_prompt
        self.generator_prompt = generator_prompt
        if not 1 <= int(self.cfg.num_distractors) <= 10:
            raise ValueError("num_distractors must be in [1, 10].")
        requested_lexicon_mode = bool(self.cfg.lexicon_mode) if lexicon_mode is None else bool(lexicon_mode)
        self.use_lexicon = False
        self.puncts = set(self.cfg.punctuations)
        self.chat_gen = chat_gen or ChatOpenAI(model=self.cfg.model_id, temperature=self.cfg.gen_temperature)
        self.chat_sel = chat_sel or ChatOpenAI(model=self.cfg.model_id, temperature=self.cfg.sel_temperature)

        self.lexicon_obj: Optional[Lexicon] = None
        self._neighbor_cache: dict[str, tuple[str, ...]] = {}
        self._threshold_filter = SurprisalThresholdFilter(
            config=SurprisalThresholdConfig(enabled=False)
        )

        if requested_lexicon_mode:
            path = self._resolve_lexicon_path()
            if path is not None:
                self.lexicon_obj = Lexicon(str(path))
                self.use_lexicon = True
                self.cfg.lexicon_path = str(path)
                print(f"[ChatAgent] Lexicon mode enabled: {path}")
            else:
                print("[ChatAgent] Lexicon mode requested but no lexicon file found. Falling back to non-lexicon mode.")

        if not self.use_lexicon and self.generator_prompt is None:
            raise ValueError("Non-lexicon mode requires generator_prompt (DistractorGeneratorPrompt).")

        if self.cfg.apply_surprisal_threshold:
            try:
                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,
                )
            except Exception as exc:
                print(
                    "[ChatAgent] WARNING: Failed to initialize surprisal scorer "
                    f"for model '{self.cfg.model_id}'. Surprisal filtering disabled. ({exc})"
                )

    # -------------------------
    # Helpers
    # -------------------------

    def _resolve_lexicon_path(self) -> Optional[Path]:
        candidates: List[Path] = []
        if self.cfg.lexicon_path:
            candidates.append(Path(self.cfg.lexicon_path))
        if self.cfg.language_code:
            root = Path(__file__).resolve().parents[1]
            candidates.append(root / "data" / "lexicon" / f"lexicon_{self.cfg.language_code}.txt")
        for p in candidates:
            if p.exists():
                return p
        return None

    def _fallback_maze_out(self, core: str) -> Dict[str, Any]:
        return build_dummy_output(
            source=core,
            dummy_len=max(1, len(core) if core else 1),
            num_distractors=self.cfg.num_distractors,
        )

    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)

    def _get_cached_neighbors(self, core: str) -> tuple[str, ...]:
        if not self.lexicon_obj:
            raise RuntimeError("Lexicon object not initialized (lexicon mode is off).")

        if core not in self._neighbor_cache:
            raw = self.lexicon_obj.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 and w != core][: self.cfg.max_candidates]
            self._neighbor_cache[core] = tuple(clean)
        return self._neighbor_cache[core]

    def _parse_distractor_pool_lenient(self, text: str) -> List[str]:
        """Return list[str] or [] if anything goes wrong."""
        try:
            fixed = repair_json(text)
            obj = json.loads(fixed)
            words = obj.get("distractors") if isinstance(obj, dict) else obj if isinstance(obj, list) else []
            return [w for w in words if isinstance(w, str)]
        except Exception:
            return []

    def _get_candidates(
        self,
        target_word: str,
        sentence_prefix: str,
    ) -> tuple[str, str, str, Sequence[str], Optional[int]]:
        pfx, core, sfx = strip_punctuation(target_word, self.puncts)
        if not core:
            return pfx, core, sfx, ("X" * len(target_word),), None

        forced_target = extract_forced_target(core)
        if forced_target is not None:
            return pfx, forced_target, sfx, tuple(), max(1, len(forced_target))

        if self.use_lexicon:
            cands = list(self._get_cached_neighbors(core))
        else:
            if self.generator_prompt is None:
                raise RuntimeError("generator_prompt is required in non-lexicon mode.")
            gen_msgs = self.generator_prompt.render_messages(sentence_prefix=sentence_prefix, word=core)
            gen_resp = self.chat_gen.invoke(gen_msgs)
            words = self._parse_distractor_pool_lenient(gen_resp.content)
            cands = normalize_candidates(words)
            cands = [c for c in cands if c != core]

        # enforce size + fallback if too few
        cands = cands[: self.cfg.max_candidates]
        if len(cands) < self.cfg.min_candidates:
            # keep pipeline running (fallback placeholders)
            pad = ["X" * len(core)] * max(0, self.cfg.min_candidates - len(cands))
            cands = (cands + pad)[: self.cfg.max_candidates]

        return pfx, core, sfx, cands, None

    def _select_distractors(self, sentence_prefix: str, core: str, candidates: Sequence[str]) -> Dict[str, Any]:
        sel_msgs = self.selector_prompt.render_messages(sentence_prefix=sentence_prefix, word=core, candidates=candidates)
        sel_resp = self.chat_sel.invoke(sel_msgs)
        text = sel_resp.content

        # 1) strict parse
        try:
            out = self.selector_prompt.parser.parse(text).model_dump()
            return out
        except Exception:
            pass

        # 2) repair + lenient dict + fill keys
        try:
            fixed = repair_json(text)
            obj = json.loads(fixed)
            if isinstance(obj, list):
                obj = next((x for x in reversed(obj) if isinstance(x, dict)), {})
            if not isinstance(obj, dict):
                return self._fallback_maze_out(core)

            obj.setdefault("source", core)
            obj.setdefault("distractor1", "X" * (len(core) if core else 1))
            obj.setdefault("distractor2", "X" * (len(core) if core else 1))
            obj.setdefault("distractor3", None)

            out = self.selector_prompt.parser.pydantic_object.model_validate(obj).model_dump()
            return out
        except Exception:
            return self._fallback_maze_out(core)

    # -------------------------
    # Jobs + Run
    # -------------------------

    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)):
                yield {
                    "sentence_index": s_idx,
                    "word_index": 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),
                    "target_word": tokens[i],
                }

    def run(self, input_data: Sequence[Sequence[str]], limit: Optional[int] = None) -> List[Dict[str, Any]]:
        run_start = time.perf_counter()
        total_sentences = len(input_data)
        results: List[Dict[str, Any]] = []
        current_sentence_index: Optional[int] = None
        current_words: List[Dict[str, Any]] = []

        for k, job in enumerate(self.iter_jobs(input_data)):
            if limit is not None and k >= limit:
                break

            if current_sentence_index is None:
                current_sentence_index = job["sentence_index"]
                print(f"[ChatAgent] Processing sentence: {current_sentence_index + 1}/{total_sentences}...")

            if job["sentence_index"] != current_sentence_index:
                results.append({"sentence_index": current_sentence_index, "words": current_words})
                current_sentence_index = job["sentence_index"]
                current_words = []
                print(f"[ChatAgent] Processing sentence: {current_sentence_index + 1}/{total_sentences}...")

            pfx, core, sfx, cands, forced_len = self._get_candidates(job["target_word"], job["sentence_prefix"])
            if job["word_index"] == 0 or forced_len is not None:
                dummy_len = int(forced_len if forced_len is not None else len(job["target_word"]))
                out = build_dummy_output(
                    source=core,
                    dummy_len=max(1, dummy_len),
                    num_distractors=self.cfg.num_distractors,
                )
                out = ensure_distractor_count(
                    out,
                    num_distractors=self.cfg.num_distractors,
                    target_word=job["target_word"],
                    target_core=core,
                    candidate_pool=cands,
                    puncts=self.puncts,
                    allow_candidate_fill=False,
                )
                if forced_len is not None:
                    out = self._reattach_punct(out, pfx, sfx)
            else:
                out = self._select_distractors(job["sentence_prefix"], core, cands)
                out = ensure_distractor_count(
                    out,
                    num_distractors=self.cfg.num_distractors,
                    target_word=job["target_word"],
                    target_core=core,
                    candidate_pool=cands,
                    puncts=self.puncts,
                    allow_candidate_fill=True,
                )
                out = self._threshold_filter.enforce_output_threshold(
                    out=out,
                    context_prefix=job["context_prefix"],
                    target_core=core,
                    candidate_pool=cands,
                    puncts=self.puncts,
                    token_joiner=self.cfg.token_joiner,
                )
                out = self._reattach_punct(out, pfx, sfx)

            out.update(
                sentence_index=job["sentence_index"],
                word_index=job["word_index"],
                sentence_prefix=job["sentence_prefix"],
                target_word=job["target_word"],
            )
            current_words.append(out)

        if current_sentence_index is not None:
            results.append({"sentence_index": current_sentence_index, "words": current_words})

        print(f"[ChatAgent] Timing summary: total={time.perf_counter() - run_start:.2f}s")
        return results

    @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 test for ChatAgent + save outputs
# -------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chat Maze distractor generation (generator + 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(
        "--processing-mode",
        default=None,
        choices=["naturalistic_reading", "controlled_experiment"],
        help="Processing mode (chat currently supports naturalistic_reading only).",
    )
    parser.add_argument("--model-id", default=None, help="Override MODEL_ID from config.")
    parser.add_argument("--word-separator", default=None, help="Override WORD_SEPARATOR from config.")
    parser.add_argument("--output-path", default=None, help="Optional explicit output JSON path.")
    parser.add_argument("--input-data-path", default=None, help="Optional explicit input data path override.")
    parser.add_argument("--template-dir", default=None, help="Optional explicit template directory override.")
    parser.add_argument("--limit", type=int, default=None, help="Optional max number of word-jobs to run.")
    parser.add_argument("--num-distractors", type=int, default=None, help="Number of distractors per word (1-10).")
    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("--lexicon-path", default=None, help="Optional explicit lexicon path override.")
    parser.add_argument(
        "--lexicon-mode",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="If true, use lexicon-mode when lexicon file exists; else fallback to non-lexicon mode.",
    )
    args = parser.parse_args()

    # Load API keys from common .env locations before initializing ChatOpenAI.
    try:
        from dotenv import load_dotenv

        load_dotenv()  # cwd + parent discovery
        project_env = Path(__file__).resolve().parents[1] / ".env"  # llmmaze/.env
        if project_env.exists():
            load_dotenv(project_env, override=False)
        config_env = Path(args.config_path).resolve().parent / ".env"
        if config_env.exists():
            load_dotenv(config_env, override=False)
    except Exception:
        # Keep behavior non-fatal if python-dotenv is unavailable.
        pass

    configs = load_config(args.config_path)
    language_code = str(args.language_code or configs.get("LANGUAGE_CODE", "en"))
    processing_mode = str(args.processing_mode or configs.get("PROCESSING_MODE", "naturalistic_reading")).strip().lower()
    if processing_mode != "naturalistic_reading":
        raise ValueError("chat_agent currently supports only naturalistic_reading mode.")
    model_id = args.model_id or str(configs.get("MODEL_ID", "gpt-4o-mini"))
    word_separator = args.word_separator if args.word_separator is not None else str(configs.get("WORD_SEPARATOR", " "))
    resolved_paths = resolve_runtime_paths(
        language_code=language_code,
        processing_mode=processing_mode,
        agent_type="chat",
        model_id=model_id,
        lexicon_path=args.lexicon_path or configs.get("LEXICON_PATH"),
        input_data_path=args.input_data_path or configs.get("INPUT_DATA_PATH"),
        template_dir=args.template_dir or configs.get("TEMPLATE_DIR"),
        output_path=args.output_path or configs.get("OUTPUT_PATH"),
    )
    input_data = read_sentences_input(str(resolved_paths["input_data_path"]), split_on=word_separator)
    puncts = get_punctuation(language_code)
    num_distractors = int(args.num_distractors if args.num_distractors is not None else configs.get("NUM_DISTRACTORS", 3))
    min_abs = args.min_abs if args.min_abs is not None else configs.get("SURPRISAL_MIN_ABS")
    min_delta = float(args.min_delta if args.min_delta is not None else configs.get("SURPRISAL_MIN_DELTA", 0.0) or 0.0)
    if args.absolute_threshold_only is None:
        absolute_threshold_only = bool(configs.get("SURPRISAL_ABSOLUTE_THRESHOLD_ONLY", False))
    else:
        absolute_threshold_only = bool(args.absolute_threshold_only)
    surprisal_device = args.surprisal_device if args.surprisal_device is not None else configs.get("SURPRISAL_DEVICE")

    template_dir = Path(resolved_paths["template_dir"])
    if not template_dir.exists():
        raise ValueError(f"Template directory not found for LANGUAGE_CODE='{language_code}': {template_dir}")
    gen_prompt = DistractorGeneratorPrompt(
        path_to_user_template=template_dir / "chat_distractor_gen_base.txt",
        path_to_extension_template=template_dir / "chat_distractor_gen_extension.txt",
        path_to_system_template=template_dir / "system.txt",
    )
    selector_prompt = MazeChatPrompt(
        path_to_user_template=template_dir / "base.txt",
        path_to_extension_template=template_dir / "extension.txt",
        path_to_system_template=template_dir / "system.txt",
    )

    cfg = ChatAgentConfig(
        lexicon_path=resolved_paths["lexicon_path"],
        language_code=language_code,
        punctuations=puncts,
        num_distractors=num_distractors,
        token_joiner="" if str(word_separator).strip() == "" else str(word_separator),
        apply_surprisal_threshold=(
            min_abs is not None
            or min_delta != 0.0
            or absolute_threshold_only
        ),
        min_abs=min_abs,
        min_delta=min_delta,
        absolute_threshold_only=absolute_threshold_only,
        surprisal_device=surprisal_device,
        model_id=model_id,
    )
    agent = ChatAgent(
        cfg=cfg,
        selector_prompt=selector_prompt,
        generator_prompt=gen_prompt,
        lexicon_mode=args.lexicon_mode,
    )

    outputs = agent.run(input_data, limit=args.limit)
    output_path = args.output_path or resolved_paths["output_path"]
    agent.save_pretty_json(outputs, output_path)
    print(f"Saved output to: {output_path}")