"""CoSER given-circumstance acting (GCA): Python reference for the platform protocol coser-gca-1.0.

The verification matrix is docs/复现规范/CoSER-核查矩阵.md. This file has three parts:

1. Generators. `strings` extracts every prompt, template, suffix and constant the platform needs from the pinned
   official sources with `ast` (never retyped) and writes worker/src/benchmarks/coser/gca/strings.gen.ts with the
   SHA-256 of each source blob; `punkt` writes punkt.gen.ts from nltk_data punkt_tab/english (the parameters
   nltk 3.9.1 word_tokenize loads).
2. Verbatim copies of the official pure functions, between `# >>> VERBATIM <commit>:<path> <name>` and
   `# <<< VERBATIM` lines. coser_spec_check.py compares each block, byte for byte, with the source it names (utils.py,
   agent.py and prompts.py have CRLF line endings; they are compared as Python reads source, CRLF as LF). Their
   dependencies (API client, cache, logger, retrieval, `random`) are replaced by the stubs in part 2a, so that the
   official code runs with scripted model outputs and fixed random choices and never calls an API.
3. Drivers that run the official functions the way the platform protocol prescribes (paper-era speaker-name prefixes,
   C16; 624efc5 NSP prompt, parsing and history, C12-C14; motivations from `thought`, C07; the platform's seeded
   choices, C11). Lines transcribed from the official scripts (main.py is not importable) are marked with their source
   lines.

    python3 -I tools/experiments/coser_reference.py strings <CoSER repo at 624efc5> [out.ts]
    python3 -I tools/experiments/coser_reference.py punkt <nltk_data>/tokenizers/punkt_tab/english [out.ts]

The repository is untrusted data: its modules are never imported; sources are read with `git show <commit>:<path>`.
"""
import ast
import contextlib
import copy
import hashlib
import io
import json
import os
import re
import subprocess
import sys
import types
from typing import Dict, List

try:  # only in the isolated environment with nltk 3.9.1 (the BLEU checks); the rest needs nothing outside the stdlib
    import nltk
    from nltk.tokenize import word_tokenize
    from nltk.translate.bleu_score import sentence_bleu
except ImportError:
    nltk = word_tokenize = sentence_bleu = None

ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
GEN_DIR = os.path.join(ROOT, "worker", "src", "benchmarks", "coser", "gca")
PINNED = "624efc50db7be3a4e37cec6a1c58ed434dc9f2ce"
PAPER_ERA = "93c147462592e372dba03afbdf37f5d0d667f8cc"
PROTOCOL = "coser-gca-1.1"
ENV, NSP_NAME = "Environment", "NSP"

# ---------------------------------------------------------------------------------------------------------------------
# 1. Generators


def git_show(repo, commit, path):
    return subprocess.run(["git", "-C", repo, "show", f"{commit}:{path}"], capture_output=True, check=True).stdout


def sha256(b):
    return hashlib.sha256(b).hexdigest()


def find_def(tree, name, kind=(ast.FunctionDef, ast.ClassDef)):
    found = [n for n in ast.walk(tree) if isinstance(n, kind) and n.name == name]
    assert len(found) == 1, (name, len(found))
    return found[0]


def source_block(src, node):
    """The source lines of a def/class (with its decorators), CRLF read as LF, with the def line's own indentation
    removed from every line (a nested def becomes top-level; nothing else changes): what a verbatim copy must equal."""
    lines = src.replace("\r\n", "\n").split("\n")
    first = min([node.lineno] + [d.lineno for d in getattr(node, "decorator_list", [])])
    prefix = lines[node.lineno - 1][:node.col_offset]
    return "\n".join(l[len(prefix):] if l.startswith(prefix) else l for l in lines[first - 1:node.end_lineno])


def parts(node):
    """An f-string as a list of literal strings and {"v": <expression source>} placeholders."""
    assert isinstance(node, ast.JoinedStr), ast.dump(node)
    out = []
    for v in node.values:
        if isinstance(v, ast.Constant):
            out.append(v.value)
        else:
            assert isinstance(v, ast.FormattedValue) and v.conversion == -1 and v.format_spec is None, ast.dump(v)
            out.append({"v": ast.unparse(v.value)})
    return out


def const_eval(node):
    """Literal strings and ints combined with + and * (for `'*'*30 + '\\n\\n'`)."""
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.BinOp) and isinstance(node.op, (ast.Add, ast.Mult)):
        a, b = const_eval(node.left), const_eval(node.right)
        return a + b if isinstance(node.op, ast.Add) else a * b
    if isinstance(node, (ast.List, ast.Tuple)):
        return [const_eval(e) for e in node.elts]
    raise ValueError(ast.dump(node))


def walk(node):
    """ast.walk in source order."""
    return sorted((n for n in ast.walk(node) if hasattr(n, "lineno")), key=lambda n: (n.lineno, n.col_offset))


def assigns(fn, name):
    return [n for n in walk(fn) if isinstance(n, ast.Assign) and any(isinstance(t, ast.Name) and t.id == name for t in n.targets)]


def extract_strings(repo):
    """Everything the platform takes from the official code, read from the pinned blobs."""
    files = {
        "utils": (PINNED, "gca_evaluation/utils.py"),
        "agent": (PINNED, "gca_evaluation/agent.py"),
        "main": (PINNED, "gca_evaluation/main.py"),
        "prompts": (PINNED, "gca_evaluation/prompts.py"),
        "utils_paper": (PAPER_ERA, "gca_evaluation/utils.py"),
        "eval_reproduce_paper": (PAPER_ERA, "gca_evaluation/eval_reproduce.py"),
        "test_set": (PINNED, "data/test/test_set.json"),
    }
    blobs = {k: git_show(repo, c, p) for k, (c, p) in files.items()}
    sources = {k: {"commit": c, "path": p, "sha256": sha256(blobs[k])} for k, (c, p) in files.items()}
    trees = {k: ast.parse(blobs[k].decode("utf-8")) for k in files if k != "test_set"}
    out = {"SOURCES": sources}

    # utils.py: get_character_prompt (fixed_template branch, used with exclude_plot_summary=True, add_rag=None).
    fn = find_def(trees["utils"], "get_character_prompt")
    of = [const_eval(a.value) for a in assigns(fn, "output_format")]
    assert len(of) == 3
    others = [n for n in fn.body if isinstance(n, ast.If) and ast.unparse(n.test) == "other_character_profiles"]
    assert len(others) == 1
    choice = [n for n in walk(others[0]) if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name) and n.targets[0].id == "decorator"]
    assert len(choice) == 1 and ast.unparse(choice[0].value.func) == "random.choice"
    aug = [n for n in walk(others[0]) if isinstance(n, ast.AugAssign) and isinstance(n.target, ast.Name) and n.target.id == "other_character_profiles_str"]
    fixed = [n for n in fn.body if isinstance(n, ast.If) and ast.unparse(n.test) == "fixed_template"]
    assert len(aug) == 1 and len(fixed) == 1
    fb = fixed[0].body
    assert ast.unparse(fb[0].test) == "motivation" and ast.unparse(fb[1].test) == "other_character_profiles_str"
    assert ast.unparse(fb[2].targets[0]) == "system_prompt" and ast.unparse(fb[3].test) == "add_rag"
    assert isinstance(fb[4], ast.AugAssign) and ast.unparse(fb[4].target) == "system_prompt"
    out["CHARACTER_PROMPT"] = {
        "outputFormat": {"speechAction": of[0], "thought": of[1], "thoughtWithExample": of[2]},
        "decorators": const_eval(choice[0].value.args[0]),
        "otherEntry": parts(aug[0].value),
        "motivationSection": parts(fb[0].body[0].value),
        "othersSection": parts(fb[1].body[0].value),
        "prompt": parts(fb[2].value),
        "requirements": parts(fb[4].value),
    }

    fn = find_def(trees["utils"], "get_environment_prompt")
    prompt = assigns(fn, "prompt")
    out["ENVIRONMENT_PROMPT"] = {
        "environment": const_eval(assigns(fn, "ENVIRONMENT")[0].value),
        "modelRoles": const_eval(assigns(fn, "model_roles")[0].value),
        "prompt": parts(prompt[0].value),
    }
    fn = find_def(trees["utils"], "get_nsp_prompt")
    out["NSP_PROMPT"] = {"environment": const_eval(assigns(fn, "ENVIRONMENT")[0].value), "prompt": parts(assigns(fn, "prompt")[0].value)}
    fn = find_def(trees["utils_paper"], "get_nsp_prompt")
    out["PAPER_ERA_NSP_PROMPT"] = {"environment": const_eval(assigns(fn, "ENVIRONMENT")[0].value), "prompt": parts(assigns(fn, "prompt")[0].value)}

    fn = find_def(trees["utils"], "extract_nsp")
    rstrip = [n for n in walk(fn) if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "rstrip"]
    repl = [n for n in walk(fn) if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "replace"]
    out["NSP_PARSE"] = {
        "pattern": const_eval(assigns(fn, "pattern")[0].value),
        "rstripChars": const_eval(rstrip[0].args[0]),
        "remove": const_eval(repl[0].args[0]),
    }

    fn = find_def(trees["utils"], "add_speaker_name")
    fs = [parts(n) for n in walk(fn) if isinstance(n, ast.JoinedStr)]
    assert len(fs) == 3
    out["SPEAKER_NAME"] = {"lineStarts": fs[:2], "prefix": fs[2]}

    fn = find_def(trees["utils"], "extract_json")
    fix = find_def(fn, "_fix_json")
    out["JSON_FIX_PROMPT"] = parts(assigns(fix, "prompt")[0].value)

    fn = find_def(trees["utils"], "get_response")
    mt = [n for n in walk(fn) if isinstance(n, ast.If) and ast.unparse(n.test) == "max_tokens is None"]
    temps = {ast.unparse(k.value) for n in walk(fn) if isinstance(n, ast.Call) for k in n.keywords if k.arg == "temperature"}
    assert temps == {"0 if nth_generation == 0 else 1"}, temps
    fn2 = find_def(trees["utils"], "get_response_json")
    retry = [n for n in walk(fn2) if isinstance(n, ast.Call) and ast.unparse(n) .startswith("kwargs.get('max_retry'")]
    out["REQUEST"] = {
        "defaultMaxTokens": const_eval(mt[0].body[0].value),
        "temperatureFirst": 0,
        "temperatureRetry": 1,
        "maxRetry": const_eval(retry[0].args[1]),
    }

    cls = find_def(trees["agent"], "Agent", ast.ClassDef)
    init = find_def(cls, "__init__")
    adds = [n.right for n in walk(init) if isinstance(n, ast.BinOp) and isinstance(n.op, ast.Add) and ast.unparse(n.left) == "self.system_prompt"]
    assert len(adds) == 2
    chat = find_def(cls, "chat")
    mt = [k.value for n in walk(chat) if isinstance(n, ast.Call) for k in n.keywords if k.arg == "max_tokens"]
    out["AGENT"] = {"concise": const_eval(adds[0]), "namePrefix": parts(adds[1]), "maxTokens": const_eval(mt[0])}

    sim = find_def(trees["main"], "gca_simulation")
    start = [n for n in walk(sim) if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == "update"
             and ast.unparse(n.func.value) == "character_agent"]
    end_cmp = [n for n in walk(sim) if isinstance(n, ast.Compare) and ast.unparse(n.left) == "i_round" and isinstance(n.ops[0], ast.GtE)]
    end_tok = [n for n in walk(sim) if isinstance(n, ast.Compare) and ast.unparse(n.left) == "next_actor" and isinstance(n.ops[0], ast.Eq)]
    mod = trees["main"]
    out["LOOP"] = {
        "conversationStart": const_eval(start[0].args[1]),
        "maxRounds": const_eval(assigns(sim, "max_rounds")[0].value),
        "endMinRound": const_eval(end_cmp[0].comparators[0]),
        "endToken": const_eval(end_tok[0].comparators[0]),
        "environment": const_eval(assigns(mod, "ENVIRONMENT")[0].value),
        "nsp": const_eval(assigns(mod, "NSP")[0].value),
    }

    judging = find_def(trees["main"], "gca_judging")
    cp = assigns(judging, "critic_prompt")[0].value
    chain = []
    node = cp
    while isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) and node.func.attr == "replace":
        chain.append({"placeholder": const_eval(node.args[0]), "value": ast.unparse(node.args[1])})
        node = node.func.value
    assert ast.unparse(node) == "critic_prompts['self-play-deduct-template']"
    score = [n for n in walk(judging) if isinstance(n, ast.Assign) and ast.unparse(n.targets[0]) == "res[dimension]['score']"]
    prompts_assign = [n for n in trees["prompts"].body if isinstance(n, ast.Assign) and n.targets[0].id == "critic_prompts"][0]
    critic = ast.literal_eval(prompts_assign.value)
    out["JUDGE"] = {
        "dimensions": const_eval(assigns(judging, "dimensions")[0].value),
        "template": critic["self-play-deduct-template"],
        "details": critic["dimension_details"],
        "replaceChain": chain[::-1],
        "additionalInstructions": const_eval(assigns(judging, "additional_instructions")[0].value),
        "scoreExpression": ast.unparse(score[0].value),
    }

    # Paper-era code (eval_reproduce.py@93c1474): the lines C16 adopts, recorded as source text.
    gen = find_def(trees["eval_reproduce_paper"], "generate_self_play_conversations")
    paper_lines = sorted({ast.unparse(n) for n in walk(gen) if isinstance(n, (ast.Assign, ast.Expr))
                          and ("add_speaker_name(response, actor)" in ast.unparse(n) or "remove_inner_thoughts(response)" in ast.unparse(n)
                               or "find_thought" in ast.unparse(n) or "next_actor.split(':')" in ast.unparse(n) or ".update('assistant', next_actor)" in ast.unparse(n))})
    judging_paper = find_def(trees["eval_reproduce_paper"], "evaluate_self_play_conversations")
    paper_lines += sorted({ast.unparse(n) for n in walk(judging_paper) if isinstance(n, ast.Assign) and "'score'" in ast.unparse(n) and "severity" in ast.unparse(n)})
    out["PAPER_ERA_LINES"] = paper_lines
    return out


TS_HEADER = """// GENERATED by tools/experiments/coser_reference.py {cmd}; do not edit. coser_spec_check.py regenerates this file
// and compares it byte for byte.
"""


def ts_value(v, indent=0):
    return json.dumps(v, ensure_ascii=True, indent=2).replace("\n", "\n" + " " * indent)


def render_strings(out):
    s = TS_HEADER.format(cmd="strings")
    s += "// Read from the pinned official sources with Python `ast`, never retyped. A literal string is a JSON string; an\n"
    s += "// f-string is a list of literal parts and { v: <Python expression> } placeholders (the expression as in the source).\n\n"
    s += "export type Part = string | { v: string };\n\n"
    docs = {
        "SOURCES": "Source blobs (git show <commit>:<path>) and their SHA-256.",
        "CHARACTER_PROMPT": "utils.py get_character_prompt, fixed_template branch (624efc5; identical strings at 93c1474).",
        "ENVIRONMENT_PROMPT": "utils.py get_environment_prompt (624efc5).",
        "NSP_PROMPT": "utils.py get_nsp_prompt (624efc5, with the Output Format section); recorded only, not used (C12).",
        "PAPER_ERA_NSP_PROMPT": "utils.py get_nsp_prompt at 93c1474 (the paper's table); the protocol's next-speaker prompt (C12).",
        "NSP_PARSE": "utils.py extract_nsp (624efc5): the regular expression (Python syntax) and the cleanup arguments; recorded only, not used (C13).",
        "SPEAKER_NAME": "utils.py add_speaker_name: the two line-start tests and the prefix.",
        "JSON_FIX_PROMPT": "utils.py extract_json._fix_json: the one-message prompt sent to the judge model.",
        "REQUEST": "utils.py get_response / get_response_json: default max_tokens, temperatures, retries after the first.",
        "AGENT": "agent.py Agent.__init__ suffixes and Agent.chat max_tokens (624efc5).",
        "LOOP": "main.py gca_simulation constants (624efc5).",
        "JUDGE": "prompts.py critic_prompts and main.py gca_judging (624efc5): dimensions, template, criteria, the order of the str.replace chain, the score expression.",
        "PAPER_ERA_LINES": "eval_reproduce.py@93c1474 statements that C07, C13, C14, C16 and C24 adopt (as source text, for the record).",
    }
    for k, v in out.items():
        s += f"/** {docs[k]} */\nexport const {k} = {ts_value(v)} as const;\n\n"
    return s


def punkt_params(lang_dir):
    """load_punkt_params of nltk 3.9.1 (nltk/tokenize/punkt.py, nltk/tabdata.py), reading the same files the same way."""
    def rm_nl(s):
        return s[:-1] if s[-1] == "\n" else s
    def lines(name):
        with open(os.path.join(lang_dir, name), encoding="utf-8") as f:
            return [rm_nl(x) for x in f]
    return {
        "collocations": [tuple(x.split("\t")) for x in lines("collocations.tab")],
        "sent_starters": lines("sent_starters.txt"),
        "abbrev_types": lines("abbrev_types.txt"),
        "ortho_context": [(a, int(b)) for a, b in (tuple(x.split("\t")) for x in lines("ortho_context.tab"))],
    }


def render_punkt(lang_dir):
    p = punkt_params(lang_dir)
    files = {n: sha256(open(os.path.join(lang_dir, n), "rb").read()) for n in ["abbrev_types.txt", "collocations.tab", "ortho_context.tab", "sent_starters.txt"]}
    s = TS_HEADER.format(cmd="punkt")
    s += "// nltk_data tokenizers/punkt_tab/english, read as nltk 3.9.1 load_punkt_params reads it (sets as sorted lists).\n\n"
    s += f"export const PUNKT_FILES = {ts_value(files)} as const;\n\n"
    s += f"export const ABBREV_TYPES: readonly string[] = {json.dumps(sorted(set(p['abbrev_types'])), ensure_ascii=True)};\n\n"
    s += f"export const SENT_STARTERS: readonly string[] = {json.dumps(sorted(set(p['sent_starters'])), ensure_ascii=True)};\n\n"
    s += f"export const COLLOCATIONS: readonly (readonly [string, string])[] = {json.dumps(sorted(set(p['collocations'])), ensure_ascii=True)};\n\n"
    s += "/** type<TAB>flags lines, in file order (a later line for the same type wins, as in the dict). */\n"
    s += f"export const ORTHO_CONTEXT = {json.dumps(chr(10).join(f'{a}{chr(9)}{b}' for a, b in p['ortho_context']), ensure_ascii=True)};\n"
    return s


# ---------------------------------------------------------------------------------------------------------------------
# 2a. Stubs for the official functions' dependencies (no network, no files, no global random stream).


class _NullLogger:
    def info(self, *a, **k):
        pass

    error = warning = debug = info


logger = _NullLogger()
config = {"api_key": "stub", "base_url": "stub"}  # not a credential: the openai module below is a stub
streaming = False
special_characters = [ENV, NSP_NAME]  # agent.py:18 (624efc5)


def cached(func):
    """Stands in for utils.cached: the platform keeps no cross-call response cache (N4)."""
    return func


def num_tokens_from_string(string, encoding_name="cl100k_base"):
    return 0


def build_rag_corpus(*a, **k):
    raise AssertionError("retrieval is not part of the protocol")


def rag(*a, **k):
    raise AssertionError("retrieval is not part of the protocol")


class FixedRandom:
    """Replaces the `random` module inside the copied functions: each choice() must be one the driver announced."""

    def __init__(self):
        self.queue = []

    def expect(self, options, chosen):
        self.queue.append((list(options), chosen))

    def choice(self, seq):
        assert self.queue, f"unexpected random.choice({seq!r})"
        options, chosen = self.queue.pop(0)
        assert list(seq) == options, (seq, options)
        assert chosen in options
        return chosen


random = FixedRandom()


class Rouge:
    """utils.calculate_bleu_rouge also computes ROUGE-L, which the platform does not report (C29)."""

    def get_scores(self, hyp, ref):
        return [{"rouge-l": {"f": 0.0}}]


class ScriptedServer:
    """The model endpoint behind the stub openai module: pops scripted outputs per model and records each request."""

    def __init__(self):
        self.scripts = {}
        self.requests = []
        self.errors = []

    def reply(self, model, messages, max_tokens, temperature):
        self.requests.append({"model": model, "messages": copy.deepcopy(messages), "max_tokens": max_tokens, "temperature": temperature})
        queue = self.scripts.get(model)
        if not queue:  # never raise here: get_response would turn it into None, and get_response_json retries None forever
            self.errors.append(f"no scripted output left for {model}")
            return ""
        return queue.pop(0)


SERVER = ScriptedServer()


def _openai_stub():
    mod = types.ModuleType("openai")

    class _Completions:
        def create(self, model, messages, max_tokens, temperature, timeout, stream=False):
            assert not stream
            content = SERVER.reply(model, messages, max_tokens, temperature)
            msg = types.SimpleNamespace(content=content)
            return types.SimpleNamespace(choices=[types.SimpleNamespace(message=msg)])

    class OpenAI:
        def __init__(self, api_key=None, base_url=None, timeout=None):
            self.chat = types.SimpleNamespace(completions=_Completions())

    mod.OpenAI = OpenAI
    return mod


class official_env:
    """Installs the stub `openai` module and a `utils` module (for `from utils import ...` inside the copied code)."""

    def __enter__(self):
        # Agent.chat prints its caught exceptions (e.g. the IndexError of a reply without the agent's own line)
        self.quiet = [contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO())]
        for q in self.quiet:
            q.__enter__()
        self.saved = {k: sys.modules.get(k) for k in ("openai", "utils")}
        sys.modules["openai"] = _openai_stub()
        utils = types.ModuleType("utils")
        for name in ("get_response_with_retry", "add_speaker_name", "remove_inner_thoughts", "fix_repeation", "get_response"):
            setattr(utils, name, globals()[name])
        sys.modules["utils"] = utils
        return self

    def __exit__(self, *exc):
        for k, v in self.saved.items():
            if v is None:
                sys.modules.pop(k, None)
            else:
                sys.modules[k] = v
        for q in reversed(self.quiet):
            q.__exit__(*exc)


# ---------------------------------------------------------------------------------------------------------------------
# 2b. Verbatim copies (624efc5 unless the marker says otherwise).

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_response (lines 102-178)
@cached
def get_response(model, messages, max_tokens=None, nth_generation=0):
	# if messages is str
	if isinstance(messages, str):
		messages = [{"role": "user", "content": messages}]
		
	# correct 'system' to 'user'
	if model.startswith('claude') and messages and messages[0]['role'] == 'system': messages[0]['role'] = 'user'	

	# merge adjacent user messages
	merged_messages = []
	for message in messages:
		if message['role'] == 'user' and merged_messages and merged_messages[-1]['role'] == 'user':
			merged_messages[-1]['content'] += message['content']
		else:
			merged_messages.append(copy.deepcopy(message))

	messages = merged_messages
	
	import openai 
	try:
		import openai 
		client = openai.OpenAI(api_key=config['api_key'], base_url=config['base_url'], timeout=180)

		# you may need to adjust the default max_tokens for your models 
		if max_tokens is None: max_tokens = 8196

		if streaming: 
			stream = client.chat.completions.create(
				model=model,
				messages=messages,
				stream=True,
				max_tokens=max_tokens,
				temperature=0 if nth_generation == 0 else 1,
				timeout=180
			)

			response = ""
			for chunk in stream:
				try:
					if chunk.choices[0].delta.content is not None:
						response += chunk.choices[0].delta.content
				except:
					if len(response) == 0:
						return None

					if len(chunk.choices) == 0 and response.strip()[-1] == '}':
						break 
		else:
			completion = client.chat.completions.create(
				model=model,
				messages=messages,
				max_tokens=max_tokens,
				temperature=0 if nth_generation == 0 else 1,
				timeout=180
			)
			response = completion.choices[0].message.content
		
		return response
	# logger.error error info
	except Exception as e:
		import traceback 
		logger.error(f'Prompt: {messages[-1]["content"][:500]}')
		logger.error(f"Error in get_response: {str(e)} from model {model}")

		try:
			if hasattr(response, 'text'):
				logger.error(f"Response: {response.text}")
			else:
				logger.error(f"Response: {response}")
		except Exception as e:
			logger.error(f"Could not print response: {e}")
		
		logger.error(f"Number of input tokens: {num_tokens_from_string(messages[0]['content'])}")

		traceback.print_exc()
		return None
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py remove_inner_thoughts (lines 196-207)
def remove_inner_thoughts(dialogue: str) -> str:
	cleaned_dialogue = re.sub(r'\[.*?\]', '', dialogue)

	cleaned_dialogue = '\n'.join(line.strip() for line in cleaned_dialogue.split('\n'))
	
	cleaned_dialogue = re.sub(r'\n+', '\n', cleaned_dialogue)

	cleaned_dialogue = re.sub(r'\s*<\s*system[_\s]+think(?:ing)?\s*>.*?</\s*system[_\s]+think(?:ing)?\s*>\s*', '', cleaned_dialogue, flags=re.S).strip()
	cleaned_dialogue = re.sub(r'\s*<\s*role[_\s]+think(?:ing)?\s*>.*?</\s*role[_\s]+think(?:ing)?\s*>\s*', '', cleaned_dialogue, flags=re.S).strip()
	cleaned_dialogue = re.sub(r'\s*<\s*[_\s]+think(?:ing)?\s*>.*?</\s*[_\s]+think(?:ing)?\s*>\s*', '', cleaned_dialogue, flags=re.S).strip()
	
	return cleaned_dialogue.strip()
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py add_speaker_name (lines 209-215)
def add_speaker_name(dialogue: str, speaker: str) -> str:
	# Check if the dialogue already contains a speaker prefix at the beginning of any line
	if any(line.strip().startswith(f"{speaker}:") or line.strip().startswith(f"{speaker}：") for line in dialogue.split('\n')):
		return dialogue
	
	# Add the speaker name at the beginning
	return f"{speaker}: {dialogue}"
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_character_prompt (lines 222-359)
def get_character_prompt(book_name, character, character_profile, background, scenario, motivation, thoughtless=False, other_character_profiles=None, exclude_plot_summary=False, fixed_template=False, add_output_example=False, add_rag=False):

	if thoughtless:
		output_format = "Your output should include **speech** and **action**. Use (your action) for actions, which others can see."
	else:
		output_format = "Your output should include **thought**, **speech**, and **action**. Use [your thought] for thoughts, which others can't see. Use (your action) for actions, which others can see."

		if add_output_example:
			output_format = "Your output should include **thought**, **speech**, and **action**. Use [your thought] for thoughts, which others can't see, e.g. [I'm terrified, but I must appear strong.]. Use (your action) for actions, which others can see, such as (watches silently, trying to control her fear and anger)."

	if other_character_profiles:
		assert isinstance(other_character_profiles, Dict)
		other_character_profiles_str = ''

		decorator = random.choice(['*'*30 + '\n\n', '*'*20 + '\n\n', '\n\n', '\n', ''])
		for other_character, profile in other_character_profiles.items():
			if other_character != character:
				other_character_profiles_str += f"{decorator}{other_character}: {profile}\n\n"
	else:
		other_character_profiles_str = ''
	
	if fixed_template:
		if motivation: motivation = f"===Your Inner Thoughts===\n{motivation}\n\n"
		if other_character_profiles_str: other_character_profiles_str = f"===Information about the other Characters===\n{other_character_profiles_str}\n\n"

		system_prompt = f"You are {character} from {book_name}.\n\n==={character}'s Profile===\n{character_profile}\n\n===Current Scenario===\n{scenario}\n\n{other_character_profiles_str}{motivation}\n\n"
		
		if add_rag:
			system_prompt += "===Relevant Background Information==={retrieved_knowledge}\n\n"
		
		system_prompt += f"===Requirements===\n{output_format}\n\n"

		return system_prompt
	
	styles = ['natural'] * 40 + ['='] * 30 + ['#'] * 20 + ['*'] * 10

	templates = {
		"begin": [f"You are {character}.", f"Play the role of {character}.", f"Imagine you are {character}.", f"Think, speak, and act like {character}.", f"Step into the shoes of {character}.", f"Immerse yourself in the character of {character}.", f"You are roleplaying as {character}.", f"You will be portraying {character}.", f"Roleplay as {character}.", f"Your role is to be {character}.", f"You are {character} from {book_name}.", f"Play the role of {character} from {book_name}.", f"Imagine you are {character} from {book_name}.", f"Think, speak, and act like {character} from {book_name}.", f"Step into the shoes of {character} from {book_name}.", f"Immerse yourself in the character of {character} from {book_name}.", f"You are roleplaying as {character} from {book_name}.", f"You will be portraying {character} from {book_name}.", f"Roleplay as {character} from {book_name}.", f"Your role is to be {character} from {book_name}."],
		"natural": {
			"character_profile": [f"The profile of {character} is as follows:\n{character_profile}", f"Here is the profile of {character}:\n{character_profile}", f"Your profile is: \n{character_profile}", f"Here is some information about {character}:\n{character_profile}", f"The background of {character} is as follows:\n{character_profile}"],
			"current_scenario": [f"The current scenario is:\n{scenario}", f"Current scenario:\n{scenario}", f"The situation you are in is:\n{scenario}", f"Here is the situation you are in:\n{scenario}"],
			"current_scenario_with_plot_summary": [f"The current scenario and its background are:\nBackground: {background}\nCurrently: {scenario}", f"Current scenario and the background:\nScenario: {scenario}\nMore Background: {background}", f"The situation you are in is:\nStory arc summary: {background}\nCurrent scenario: {scenario}", f"Here is the situation you are in:\nSummary of relevant plots: {background}\nScenario: {scenario}"],
			"other_characters_profile": [f"Here is the your knowledge about the other characters:\n{other_character_profiles_str}", f"Information about other characters:\n{other_character_profiles_str}", f"The background of other characters is as follows:\n{other_character_profiles_str}"],
			"thought": [f"Your thoughts are:\n{motivation}", f"Your thoughts in this situation are:\n{motivation}", f"Your inner thoughts are:\n{motivation}", f"Your inner monologue is:\n{motivation}", f"Your inner thoughts in the scenario are:\n{motivation}"],
			"requirements": [output_format, "" if thoughtless else output_format],
		},
		"=": {
			"decorator": ["==={}===", "=={}==", "={}="],
		},
		"#": {
			"decorator": ["#{}", "# {}", "## {}", "### {}"],
		}, 
		"*": {
			"decorator": ["**{}**", "*{}*", "***{}***"],
		},
		"pieces":{
			"character_profile": ["Character Profile", f"The profile of {character}", f"{character}'s profile"],
			"current_scenario": ["Current Scenario", "The situation you are in", "Scenario"],
			"plot_summary": ["Summary of Relevant Plots", "Background", "Story Arc", "Plot Summary"],
			"thought": [f"{character}'s Thought", "Your thoughts", "Your inner thoughts", "Your inner monologue"],
			"other_characters_profile": [f"Information about other characters", f"The background of other characters", f"Other characters' profiles"],
			"requirements": ["Requirements", "Instructions for roleplaying"],
		}
	}

	# Randomly select a style
	current_style = random.choice(styles)
	
	# Start with a random beginning template
	system_prompt = random.choice(templates["begin"]) + "\n\n"
	
	# Add decorated sections based on style
	if current_style == 'natural':
		# Natural style without decorators
		system_prompt += random.choice(templates["natural"]["character_profile"]) + "\n\n"

		if exclude_plot_summary or random.random() < 0.5:
			system_prompt += random.choice(templates["natural"]["current_scenario"]) + "\n\n"
		else:
			# use Plot Summary in 50% cases
			system_prompt += random.choice(templates["natural"]["current_scenario_with_plot_summary"]) + "\n\n"

		if other_character_profiles_str:
			system_prompt += random.choice(templates["natural"]["other_characters_profile"]) + "\n\n"

		if motivation:
			system_prompt += random.choice(templates["natural"]["thought"]) + "\n\n"
		
		if add_rag:
			system_prompt += "Relevant Background Information: \n{retrieved_knowledge}\n\n"

		system_prompt += random.choice(templates["natural"]["requirements"]) + "\n\n"
	else:
		# Styled with decorators
		decorator = random.choice(templates[current_style]["decorator"])
		
		# Character profile section
		section_title = random.choice(templates["pieces"]["character_profile"])
		system_prompt += decorator.format(section_title) + "\n"
		system_prompt += character_profile + "\n\n"
		
		if not exclude_plot_summary and random.random() < 0.5:
			# use Plot Summary in 50% cases
			# Plot summary section
			section_title = random.choice(templates["pieces"]["plot_summary"])
			system_prompt += decorator.format(section_title) + "\n"
			system_prompt += background + "\n\n"

		# Current scenario section
		section_title = random.choice(templates["pieces"]["current_scenario"])
		system_prompt += decorator.format(section_title) + "\n"
		system_prompt += f"{scenario}\n\n"

		if other_character_profiles_str:
			section_title = random.choice(templates["pieces"]["other_characters_profile"])
			system_prompt += decorator.format(section_title) + "\n"
			system_prompt += other_character_profiles_str + "\n\n"

		# Thought section (if not empty)
		if motivation:
			section_title = random.choice(templates["pieces"]["thought"])
			system_prompt += decorator.format(section_title) + "\n"
			system_prompt += motivation + "\n\n"
		
		if add_rag:
			section_title = "Relevant Background Information"
			system_prompt += decorator.format(section_title) + "\n"
			system_prompt += "{retrieved_knowledge}" + "\n\n"

		# Requirements section (if not empty)
		requirements = random.choice(templates["natural"]["requirements"])
		if requirements:
			section_title = random.choice(templates["pieces"]["requirements"])
			system_prompt += decorator.format(section_title) + "\n"
			system_prompt += requirements + "\n\n"
		

	return system_prompt
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_environment_prompt (lines 361-389)
def get_environment_prompt(major_characters, scenario):
	ENVIRONMENT = "Environment"
	major_characters = [c for c in major_characters if c != ENVIRONMENT]

	model_roles = [
		"an environment model",
		"a world model",
		"a world simulator",
		"an environment simulator"
	]

	prompt = f"""You are {random.choice(model_roles)} for a role-playing game. Your task is to provide the environmental feedback: Based on the characters' interactions, dialogues, and actions, describe the resulting changes in the environment. This includes:
   - Physical changes in the setting
   - Reactions of background characters or crowds
   - Ambient sounds, weather changes, or atmospheric shifts
   - Any other relevant environmental details

Your descriptions should be vivid and help set the scene, but avoid dictating the actions or dialogue of the main characters (including {major_characters}).

Important notes:
- You may include actions and reactions of minor characters or crowds, as long as they're not main characters (including {major_characters}).
- Keep your environmental descriptions concise but impactful, typically 1-3 sentences.
- Respond to subtle cues in the characters' interactions to create a dynamic, reactive environment.
- Your output should match the tone, setting, and cultural context of the scenario.

===The scenario is as follows===
{scenario}"""

	return prompt
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_nsp_prompt (lines 391-404)
def get_nsp_prompt(all_characters, scenario):
	ENVIRONMENT = "Environment"

	prompt = f"""Your task is to predict the next speaker for a role-playing game. That is, you need to determine which character (or the {ENVIRONMENT}) might act next based on their previous interactions. The {ENVIRONMENT} is a special role that provides the environmental feedback. Choose a name from this list: {all_characters}. If it's unclear who should act next, output "random". If you believe the scene or conversation should conclude, output "<END CHAT>".

===The scenario is as follows===
{scenario}

===Output Format===
You must structure your response as follows:
1. Reasoning: Explain your thought process.
2. Next Speaker: Select exactly one character name from the list: {all_characters}.
"""
	return prompt
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py extract_nsp (lines 407-420)
def extract_nsp(response):
	pattern = r'(?:\*\*|#)?\s*Next\s+Speaker\s*(?:\*\*|#)?\s*(?:[:\-]|\bis\b)?\s*\*?([^\n\*]+)'
	
	matches = re.findall(pattern, response, re.IGNORECASE)
	
	if matches:
		last_match = matches[-1]
		
		clean_name = last_match.strip().rstrip('.,!"\'').replace('*', '')
		
		if clean_name:
			return clean_name
			
	return None
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py extract_json (lines 455-509)
def extract_json(text, **kwargs):
	def _fix_json(json_response):
		prompt = f'''I will provide you with a JSON string that contains errors, making it unparseable by `json.loads`. The most common issue is the presence of unescaped double quotes inside strings. Your task is to output the corrected JSON string. The JSON string to be corrected is:
{json_response}
'''

		response = get_response(model=kwargs['model'], messages=[{"role": "user", "content": prompt}])

		logger.info(f'fixed json: {response}')	

		return response

	def _extract_json(text):
		# Use regular expressions to find all content within curly braces
		orig_text = text

		text = re.sub(r'"([^"\\]*(\\.[^"\\]*)*)"', lambda m: m.group().replace('\n', r'\\n'), text) 
		
		#json_objects = re.findall(r'(\{[^{}]*\}|\[[^\[\]]*\])', text, re.DOTALL)

		def parse_json_safely(text):
			try:
				result = json.loads(text)
				return result
			except json.JSONDecodeError:
				results = []
				start = 0
				while start < len(text):
					try:
						obj, end = json.JSONDecoder().raw_decode(text[start:])
						results.append(obj)
						start += end
					except json.JSONDecodeError:
						start += 1
				
				if results:
					longest_json = max(results, key=lambda x: len(json.dumps(x)))
					return longest_json
				else:
					return None
		
		extracted_json = parse_json_safely(text)
		
		if extracted_json:
			return extracted_json
		else:
			logger.error('Error parsing response: ', orig_text)
			return None

	res = _extract_json(text)

	if res:
		return res
	else:
		return _extract_json(_fix_json(text))
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_response_with_retry (lines 539-541)
def get_response_with_retry(**kwargs):

	return get_response_json([], **kwargs)
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py get_response_json (lines 543-584)
def get_response_json(post_processing_funcs=[extract_json], **kwargs):
	"""
    Get and process a response from an LLM with retries and error handling.
    
    This function handles:
    1. Getting responses from the LLM with retries
    2. Processing responses through a pipeline of post-processing functions
    3. Fallback handling for parsing failures
    
    Args:
        post_processing_funcs (list): List of functions to process the LLM response, defaults to [extract_json]
        **kwargs: Additional arguments passed to get_response(), including:
            - messages: List of message dicts for the LLM
            - model: Name of LLM model to use
            - max_retry: Max number of retry attempts (default 5)
            
    Returns:
        dict: Processed JSON response from the LLM, or error dict if parsing fails
    """
	nth_generation = 0

	while (True):
		response = get_response(**kwargs, nth_generation=nth_generation)

		#print(f'{nth_generation}th response: {response}')

		if response is None:
			continue 

		for post_processing_func in post_processing_funcs:
			response = post_processing_func(response, **kwargs)
		#print(f'parse results: {json_response}')
		json_response = response 
		
		if json_response:
			break 
		else:
			nth_generation += 1
			if nth_generation > kwargs.get('max_retry', 5):				
				break	
	
	return json_response
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py tokenize_words (lines 598-610)
def tokenize_words(text):
	import regex
	pattern = r'\b\w+\b|[\u4e00-\u9fff]|[\u3040-\u309F\u30A0-\u30FF]|\d|[\p{P}\p{S}]'
	tokens = regex.findall(pattern, text)


	tokens_expanded = []
	for token in tokens:
		if re.match(r'[\u4e00-\u9fff]|[\u3040-\u309F\u30A0-\u30FF]', token):
			tokens_expanded.extend(list(token))
		else:
			tokens_expanded.append(token)
	return tokens_expanded
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py fix_repeation (lines 612-743)
def fix_repeation(response):
	"""
	Fix repetitive text patterns in the response by detecting and removing repetitions.
	
	This function handles three types of repetition detection:
	1. Long letter substrings (100+ characters)
	2. Consecutive repetitions of token sequences
	3. Non-consecutive repetitions of token sequences
	
	Args:
		response (str): The text response to check for repetitions
		
	Returns:
		str: The fixed text with repetitions removed if repetitions were found
		False: If no repetitions were detected
	"""

	def detect_repetitions(tokens, min_length=5, max_length=30, threshold=0.1):
		"""Check for consecutive repetitions of token sequences"""
		total_length = len(tokens)
		repetitions = 0
	
		# Try different lengths of subsequences
		for length in range(min_length, min(max_length + 1, total_length + 1)):
			for i in range(total_length - length + 1):
				substr = tokens[i:i + length] 

				# Check if this subsequence repeats consecutively up to 4 times
				is_repeated = True
				for repeat_idx in range(1, 5):
					check_pos = i + (repeat_idx * length)

					if tokens[check_pos:check_pos + length] != substr:
						is_repeated = False
						break
				
				if is_repeated:
					return tokens[:i + length]  # Return text up to first repetition

		return False

	def detect_repetitions2(tokens, min_length=15, max_length=30, threshold=0.1):
		"""Check for non-consecutive repetitions of token sequences"""
		total_length = len(tokens)
		repetitions = 0
		
		first_repeat_idx = 999999999999999
		first_start_idx = {}

		# Try different lengths of subsequences
		for length in range(min_length, min(max_length + 1, total_length + 1)):
			substr_count = {}

			for i in range(total_length - length + 1):
				substr = tuple(tokens[i:i + length]) 
				if substr_count.get(substr, 0) > 0:
					# Found a repeat - check if it's far enough from first occurrence
					if i - first_start_idx[substr] >= length:			
						first_repeat_idx = min(first_repeat_idx, i)
				else:
					first_start_idx[substr] = i

				substr_count[substr] = substr_count.get(substr, 0) + 1
			
			repetitions += sum(count > 1 for count in substr_count.values())

		repetition_rate = repetitions / total_length if total_length else 0
	
		if first_repeat_idx < 999999999999999:
			return tokens[:first_repeat_idx]  # Return text up to first repetition
		else:
			return False

	def concatenate_tokens(tokens):
		"""Reconstruct text from tokens with proper spacing and punctuation"""
		text = ""
		last_type = None
		
		for token in tokens:
			# Determine token type (CJK, punctuation, or other)
			current_type = 'CJK' if re.match(r'[\u4e00-\u9fff]|[\u3040-\u309F\u30A0-\u30FF]', token) else 'Other'
			import string
			if token in string.punctuation:
				current_type = 'P'

			# Add space between certain token types
			if last_type in ['Other', 'P'] and current_type == 'Other':
				text += " " + token
			else:
				text += token

			last_type = current_type
		
		# Add appropriate ending punctuation based on last character
		if re.match(r'[a-zA-Z0-9]+$', text[-1]):
			text += '.'
		if re.match(r'[\u4e00-\u9fff]+$', text[-1]):
			text += '。'
		if re.match(r'[\u3040-\u309F\u30A0-\u30FF]+$', text[-1]):
			text += '。'

		return text

	def find_long_letter_substrings(s):
		"""Find substrings of letters that are 100+ characters long"""
		pattern = r'[a-zA-Z]{100,}'
		matches = re.findall(pattern, s)
		return matches

	repeat_sign = False 

	# First check for very long letter sequences
	_ = find_long_letter_substrings(response)
	if _:
		for substr in _:
			response = response.replace(substr, substr[:20])
		repeat_sign = True

	# Then check for token sequence repetitions
	tokens = tokenize_words(response)
	_ = detect_repetitions(tokens)
	if _ == False:  # If no consecutive repetitions found
		_ = detect_repetitions2(tokens)  # Check for non-consecutive repetitions

	if _:
		response = concatenate_tokens(_)
		repeat_sign = True

	if repeat_sign:
		return response  # Return fixed text if repetitions were found
	else:
		return False  # Return False if no repetitions detected
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/utils.py calculate_bleu_rouge (lines 769-782)
def calculate_bleu_rouge(reference, simulation):

    simulation_str = '\n\n'.join([m['content'].strip('\n') for m in simulation])
    reference_str = '\n\n'.join([f"{m['character']}: {m['message']}".strip('\n') for m in reference])

    # remove the speaker name
    reference_tokens = word_tokenize(reference_str.lower())
    simulation_tokens = word_tokenize(simulation_str.lower())
    
    bleu = sentence_bleu([reference_tokens], simulation_tokens)
    
    rouge_l = Rouge().get_scores(simulation_str, reference_str)[0]['rouge-l']['f']
    
    return bleu, rouge_l
# <<< VERBATIM

# >>> VERBATIM 624efc5:gca_evaluation/agent.py Agent (lines 299-462)
class Agent:
    """
    A conversational agent class that manages dialogue interactions using language models.
    
    This agent can engage in conversations while maintaining context, retrieve relevant information 
    from a knowledge database, and generate appropriate responses based on its configuration.

    Attributes:
        model (str): The language model to use for generating responses
        name (str): Name of the agent/character
        database (Dict): Knowledge database for retrieval
        scene (Dict): Scene context information
        system_prompt (str): Initial system prompt to guide agent behavior
        retrievers (Dict): RAG retrievers for different content types
        system_role (str): Role type for system messages ('user' or 'system')
        messages (List): Conversation history as a list of message dictionaries
    """
    def __init__(self, model: str, name, database: Dict, system_prompt: str = None, scene: Dict = None, retrieval_target: str = 'conversation'):
        # Initialize basic agent properties
        self.model = model 
        self.name = name 
        self.database = database
        self.scene = scene

        self.system_prompt = system_prompt 
        
        # Clean up system prompt by removing trailing newlines
        self.system_prompt = self.system_prompt.strip('\n')
        
        # Set up RAG retrievers if database is provided
        if retrieval_target and database:
            self.retrievers = build_rag_corpus(name, database, retrieval_target)
        else:
            self.retrievers = None

        # Add conciseness instruction for non-special characters
        if self.name not in special_characters:
            self.system_prompt = self.system_prompt + '\n\nSpeak concisely as humans, instead of being verbose. Limit your response to 60 words.\n\n'

        # Add name prefix instruction for specific models
        if self.model.startswith('llama') or self.model.startswith('step'):
            self.system_prompt = self.system_prompt + f'Start your response with "{name}: ". Avoid speaking as other characters.\n\n'
        
        # Set appropriate system role based on model type
        if self.model.startswith('claude') or self.model.startswith('o1'):
            self.system_role = 'user'
        else:
            self.system_role = 'system'

        # Initialize conversation history with system prompt
        self.messages = [{"role": self.system_role, "content": self.system_prompt}]
        

    def chat(self) -> str:
        """
        Generates a response based on the conversation history and available knowledge.
        
        This method:
        1. Retrieves relevant information from the knowledge database if available
        2. Generates a response using the configured language model
        3. Processes and cleans the response based on model-specific requirements
        
        Returns:
            str: The generated response text, or empty string if an error occurs
        """
        try:
            messages = self.messages
            if self.retrievers:
                # Retrieve relevant information from recent context (last 3 messages)
                contexts = self.messages[1:]
                contexts = contexts[-3:]

                # Gather knowledge from all configured retrievers
                knowledge = ''
                for target_type, retriever in self.retrievers.items():
                    knowledge += rag(contexts, retriever, target_type)

                # Insert retrieved knowledge into system prompt
                messages = copy.deepcopy(self.messages)
                messages[0]['content'] = messages[0]['content'].replace('{retrieved_knowledge}', '<begin of background information>\n\n' + knowledge + '\n\n<end of background information>\n\n')

            from utils import get_response_with_retry
            response = get_response_with_retry(model=self.model, messages=messages, max_tokens=512)

            def parse_response(response: str, character_name: str) -> str:
                """
                Extracts the utterance of a specific character from a (unexpected) multi-character response.
                
                Args:
                    response (str): Full response text
                    character_name (str): Name of character whose utterance to extract
                
                Returns:
                    str: Extracted utterance for the specified character
                """
                lines = response.split('\n')
                current_character = None
                current_utterance = ""
                parsed_utterances = []

                for line in lines:
                    # Check for character name at start of line
                    if ':' in line:
                        character = line.split(':', 1)[0].strip()
                        
                        if current_character != character:
                            # Save previous character's utterance and start new one
                            if current_utterance:
                                parsed_utterances.append((current_character, current_utterance))
                            current_character = character
                            current_utterance = ""

                    current_utterance += line + "\n"
            
                # Save final utterance
                if current_utterance:
                    parsed_utterances.append((current_character, current_utterance))
                
                parsed_utterances = [utterance for character, utterance in parsed_utterances if character == character_name][0]

                return parsed_utterances

            # Process response based on model type and agent name
            if (self.model.startswith('llama') or self.model.startswith('step')) and self.name != 'NSP':
                response = parse_response(response, self.name)
            
            # Fix repetition issues for certain models
            if not any(self.model.lower().startswith(model_type) for model_type in ['gpt', 'claude']) and self.name != 'NSP':
                res = fix_repeation(response)
                if res:
                    logger.info(f'{self.model} Repetition found and fixed: Original: "{response}" Fixed: "{res}"')
                    response = res

            return response

        except Exception as e:
            import traceback
            print(f"Error getting response: {e}")
            traceback.print_exc()
            
            return ""
    
    def update(self, role: str, message: str):
        """
        Updates the conversation history with a new message.
        
        Args:
            role (str): Role of the message sender
            message (str): Content of the message
        """
        if message:
            # Append message to last message if same role, otherwise add new message
            if self.messages and self.messages[-1]['role'] == role:
                self.messages[-1]['content'] = self.messages[-1]['content'] + '\n\n' + message
            else:
                self.messages.append({"role": role, "content": message})

        return

    def reset(self):
        """
        Resets the conversation history to initial state with only system prompt.
        """
        self.messages = self.messages[:1]
# <<< VERBATIM

dimensions = ['Storyline Consistency', 'Anthropomorphism', 'Character Fidelity', 'Storyline Quality']  # main.py:393, read by parse_response
# >>> VERBATIM 624efc5:gca_evaluation/main.py parse_response (lines 438-451)
def parse_response(response, **kwargs):
    try:
        assert isinstance(response, dict)
        for k, v in response.items():
            assert k in dimensions
            assert 'flaws' in v

            for f in v['flaws']:
                if f.get('severity', None) is None:
                    f['severity'] = 1

        return response
    except:
        return False
# <<< VERBATIM


remove_inner_thoughts_624efc5, get_nsp_prompt_624efc5 = remove_inner_thoughts, get_nsp_prompt

# The paper-era versions where they differ (recorded; the protocol uses 624efc5's remove_inner_thoughts and NSP prompt).
# >>> VERBATIM 93c1474:gca_evaluation/utils.py remove_inner_thoughts (lines 206-216)
def remove_inner_thoughts(dialogue: str) -> str:

	cleaned_dialogue = re.sub(r'\[.*?\]', '', dialogue)
	
	cleaned_dialogue = '\n'.join(line.strip() for line in cleaned_dialogue.split('\n'))
	
	cleaned_dialogue = re.sub(r'\n+', '\n', cleaned_dialogue)


	
	return cleaned_dialogue.strip()
# <<< VERBATIM

# >>> VERBATIM 93c1474:gca_evaluation/utils.py get_nsp_prompt (lines 392-400)
def get_nsp_prompt(all_characters, scenario):
	ENVIRONMENT = "Environment"

	prompt = f"""Your task is to predict the next speaker for a role-playing game. That is, you need to determine which character (or the {ENVIRONMENT}) might act next based on their previous interactions. The {ENVIRONMENT} is a special role that provides the environmental feedback. Choose a name from this list: {all_characters}. If it's unclear who should act next, output "random". If you believe the scene or conversation should conclude, output "<END CHAT>".

===The scenario is as follows===
{scenario}"""
	
	return prompt
# <<< VERBATIM


remove_inner_thoughts_93c1474, get_nsp_prompt_93c1474 = remove_inner_thoughts, get_nsp_prompt
# C21: the 624efc5 version of remove_inner_thoughts; C12: the 93c1474 (paper) version of the next-speaker prompt.
remove_inner_thoughts, get_nsp_prompt = remove_inner_thoughts_624efc5, get_nsp_prompt_93c1474

# ---------------------------------------------------------------------------------------------------------------------
# 3. Drivers.

# Official model-name strings that make the copied code take each family's branches (agent.py and utils.py test the
# model name with startswith / `in`): the platform's family table must map its model ids to the same behaviours.
OFFICIAL_NAME = {
    "gpt": "gpt-4o",
    "claude": "claude-3-5-sonnet-20240620",
    "o1": "o1-mini",
    "llama31-instruct": "llama3p1-8b-instruct",
    "step": "step-1-8k",
    "coser": "coser-70b",
    "other": "qwen2.5-72b-instruct",
}


def platform_rng_index(sample_id, purpose, n, k=0, protocol=PROTOCOL):
    """The platform's seeded choice (C11), implemented independently of the TypeScript."""
    h = hashlib.sha256("\x1f".join(["coser-gca", protocol, sample_id, purpose, str(k)]).encode("utf-8")).digest()
    return int.from_bytes(h[:8], "big") % n


def sample_id(circ):
    return f"{circ['book']}-{circ['plot']['i_p']}-{circ['i_c']}"


def ordered_fallback(major_characters, speaking_characters_w_env, current_speaker):
    """main.py:300-303 (624efc5) with the set replaced by its first-occurrence order (C11)."""
    seen = []
    for c in major_characters + [ENV]:
        if c not in seen and c != current_speaker:
            seen.append(c)
    if not seen:
        for c in speaking_characters_w_env:
            if c not in seen and c != current_speaker:
                seen.append(c)
    return seen


def nsp_decide(response, i_round, current_speaker, speaking_characters_w_env, major_characters, sid):
    """eval_reproduce.py:231-245 (93c1474): the text before the first colon, END from round 5, validity, and the fallback
    speaker drawn from the ordered candidates with the platform's seeded index (C11). Returns (next_actor, next speaker, invalid)."""
    next_actor = response
    if ':' in next_actor:
        next_actor = next_actor.split(':')[0].strip()
    invalid_nsp = False
    if next_actor == "<END CHAT>" and i_round >= 5:
        current_speaker = "<END CHAT>"
    elif next_actor in speaking_characters_w_env and next_actor != current_speaker:
        current_speaker = next_actor
    else:
        invalid_nsp = True
        candidates = ordered_fallback(major_characters, speaking_characters_w_env, current_speaker)
        current_speaker = candidates[platform_rng_index(sid, f"nsp-fallback:{i_round}", len(candidates))]
    return next_actor, current_speaker, invalid_nsp


def build_agents(circumstance, actor_model, env_model, nsp_model, add_output_example, decorators, env_role):
    """main.py:155-257 (624efc5) without retrieval; the motivation is read from `thought` as eval_reproduce.py:163
    (93c1474) does (C07); add_output_example comes from the family table (C08). `decorators` maps each character to the
    separator the platform drew, `env_role` is its environment phrase: they are fed to the copied functions' random.choice."""
    circumstance = copy.deepcopy(circumstance)
    book_title = circumstance['book']
    plot = circumstance['plot']
    conversation = circumstance
    character_profiles = circumstance['character_profiles']
    plot_characters = [c['name'] for c in plot['key_characters']]
    speaking_characters_w_env = conversation['speaking_characters_w_env']
    if ENV not in speaking_characters_w_env:
        speaking_characters_w_env.append(ENV)
    major_characters = conversation['major_characters']
    character_agents = {}
    involved_character_profiles = {}
    for character in speaking_characters_w_env:
        if character == ENV:
            continue
        character_profile = character_profiles.get(character, '')
        if character in plot_characters:
            character_info = [c for c in plot['key_characters'] if c.get('name', '') == character][0]
            if 'description' in character_info:
                character_profile = character_info.get('description', '').strip('\n') + '\n\n' + character_profile.strip('\n')
        character_profile = character_profile.strip(' \n')
        if character_profile != '':
            involved_character_profiles[character] = character_profile
    deco_options = None
    for character in speaking_characters_w_env + [NSP_NAME]:
        if character == NSP_NAME:
            system_prompt = get_nsp_prompt(speaking_characters_w_env, conversation['scenario'])
        elif character == ENV:
            random.expect(["an environment model", "a world model", "a world simulator", "an environment simulator"], env_role)
            system_prompt = get_environment_prompt(major_characters, conversation['scenario'])
        else:
            character_profile = involved_character_profiles.get(character, '')
            character_profile = character_profile.strip(' \n')
            find_thought = [c.get('thought', '') for c in conversation['key_characters'] if c.get('name', '') == character]
            thought = find_thought[0] if find_thought else ''
            if involved_character_profiles:
                deco_options = ['*' * 30 + '\n\n', '*' * 20 + '\n\n', '\n\n', '\n', '']
                random.expect(deco_options, decorators[character])
            system_prompt = get_character_prompt(
                book_title, character, character_profile, plot["summary"],
                conversation["scenario"], thought, thoughtless=False,
                other_character_profiles=involved_character_profiles,
                exclude_plot_summary=True, fixed_template=True,
                add_output_example=add_output_example, add_rag=None
            )
        if character not in special_characters:
            character_model = actor_model
        elif character == ENV:
            character_model = env_model
        else:
            character_model = nsp_model
        character_agent = Agent(character_model, character, None, system_prompt=system_prompt, retrieval_target=None)
        character_agent.update('user', "===Conversation Start===\n\n")
        character_agents[character] = character_agent
    assert not random.queue
    return {
        "agents": character_agents,
        "speaking": speaking_characters_w_env,
        "major": major_characters,
        "involved": involved_character_profiles,
    }


def simulate(circumstance, actor_model, env_model, nsp_model, add_output_example, decorators, env_role, scripts):
    """eval_reproduce.py:227-266 (93c1474, the same rules as main.py at d7cee90): for the NSP the paper's prompt, the text
    before the first colon and the parsed name in its history (C12-C14); `response = add_speaker_name(response, actor)`
    before recording, and the other agents receive remove_inner_thoughts(response) (C16). The fallback speaker is drawn from the ordered candidates with
    the platform's seeded index (C11). Every model call goes through the copied Agent.chat -> get_response_with_retry ->
    get_response_json -> get_response, and is answered by SERVER from `scripts` ({model: [outputs]}).
    Returns the record and the requests each agent sent."""
    SERVER.scripts = copy.deepcopy(scripts)
    SERVER.requests = []
    SERVER.errors = []
    sid = sample_id(circumstance)
    with official_env():
        built = build_agents(circumstance, actor_model, env_model, nsp_model, add_output_example, decorators, env_role)
        character_agents = built["agents"]
        speaking_characters_w_env = built["speaking"]
        major_characters = built["major"]
        requests = []
        max_rounds = 20
        agent_conversations = []
        current_speaker = speaking_characters_w_env[0]
        for i_round in range(max_rounds):
            if current_speaker == "<END CHAT>":
                break
            for actor in [current_speaker, "NSP"]:
                current_agent = character_agents[actor]
                n_before = len(SERVER.requests)
                response = current_agent.chat()
                requests.append({"agent": actor, "calls": SERVER.requests[n_before:]})
                if actor == "NSP":
                    next_actor, current_speaker, invalid_nsp = nsp_decide(
                        response, i_round, current_speaker, speaking_characters_w_env, major_characters, sid)
                    agent_conversations.append({"role": actor, "content": next_actor})
                    current_agent.update('assistant', next_actor)
                else:
                    response = add_speaker_name(response, actor)
                    agent_conversations.append({"role": actor, "content": response})
                    for other_actor, other_agent in character_agents.items():
                        if other_actor == actor:
                            other_agent.update('assistant', response)
                        else:
                            other_agent.update('user', remove_inner_thoughts(response))
    leftover = {m: q for m, q in SERVER.scripts.items() if q}
    return {"simulation": agent_conversations, "requests": requests, "involved": built["involved"], "leftover": leftover, "errors": SERVER.errors}


def judge_inputs(circumstance, simulation, involved_character_profiles):
    """main.py:408-457 (624efc5), continue_from = 0."""
    simulation = [m for m in simulation if m['role'] != NSP_NAME]
    reference = circumstance['dialogues']
    simulation = [m if m['role'] == ENV else {**m, 'content': remove_inner_thoughts(m['content'])} for m in simulation]
    reference = [m if m['character'] == ENV else {**m, 'message': remove_inner_thoughts(m['message'])} for m in reference]
    simulation_str = '\n\n'.join([m['content'].strip('\n') for m in simulation])
    reference_str = '\n\n'.join([f"{m['character']}: {m['message']}".strip('\n') for m in reference])
    character_profile_str = '\n\n'.join([f"### {character}\n\n{profile.strip('')}" for character, profile in involved_character_profiles.items()])
    actor_rounds = len([m for m in simulation if m['role'] != ENV])
    return {"simulation": simulation, "reference": reference, "simulation_str": simulation_str, "reference_str": reference_str,
            "character_profile_str": character_profile_str, "actor_rounds": actor_rounds}


def judge_messages(circumstance, simulation, involved_character_profiles, dimension):
    """main.py:462-464 (624efc5)."""
    ji = judge_inputs(circumstance, simulation, involved_character_profiles)
    book_title = circumstance['book']
    scenario_str = circumstance['scenario']
    major_characters = circumstance['major_characters']
    additional_instructions = ''
    critic_prompts = PROMPTS_CRITIC
    critic_prompt = critic_prompts['self-play-deduct-template'].replace('{book}', book_title).replace('{plot_summary}', circumstance['plot']['summary']).replace('{scenario}', scenario_str).replace('{character_profiles}', ji['character_profile_str']).replace('{original_conversation}', ji['reference_str']).replace('{major_characters}', ', '.join(major_characters)).replace('{additional_instructions}', additional_instructions).replace('{dimension_name}', dimension).replace('{dimension_brief}', critic_prompts['dimension_details'][dimension]['dimension_brief']).replace('{dimension_criteria}', critic_prompts['dimension_details'][dimension]['dimension_criteria'])
    return [{"role": "system", "content": critic_prompt}, {"role": "user", "content": ji['simulation_str']}]


PROMPTS_CRITIC = None  # set by load_critic_prompts(repo): prompts.py is a pure literal, read with ast.literal_eval


def load_critic_prompts(repo):
    global PROMPTS_CRITIC
    tree = ast.parse(git_show(repo, PINNED, "gca_evaluation/prompts.py").decode("utf-8"))
    node = [n for n in tree.body if isinstance(n, ast.Assign) and n.targets[0].id == "critic_prompts"][0]
    PROMPTS_CRITIC = ast.literal_eval(node.value)


def read_judge(text, judge_model, fix_output):
    """get_response_json([extract_json, parse_response], ...) for one generation: extract_json (with its fix call
    answered by `fix_output`), then parse_response. Returns (value, fix_requests)."""
    SERVER.scripts = {judge_model: [fix_output]} if fix_output is not None else {}
    SERVER.requests = []
    SERVER.errors = []
    with official_env():
        try:
            res = extract_json(text, model=judge_model, messages=[])
            res = parse_response(res)
        except Exception as e:  # the official run would stop here
            return ("crash", type(e).__name__), SERVER.requests
    return res, SERVER.requests


def score(res, dimension, actor_rounds):
    """eval_reproduce.py:409 (93c1474), floored at 0 only; main.py:471 (624efc5) with its ceiling of 100, the line as written;
    the unclipped value; the integer severity sum."""
    res = copy.deepcopy(res)
    res[dimension]['score'] = max(0, min(100 - (sum([f['severity'] for f in res[dimension]['flaws'] if isinstance(f['severity'], int)]) - 0.3 * actor_rounds) * 5, 100) )
    upper = res[dimension]['score']
    s = sum([f['severity'] for f in res[dimension]['flaws'] if isinstance(f['severity'], int)])
    unclipped = 100 - (s - 0.3 * actor_rounds) * 5
    return max(0, unclipped), unclipped, s, upper


def tagged(v):
    """A JSON-like value with Python's types made explicit, for comparison with the TypeScript PyVal."""
    if isinstance(v, bool):
        return {"bool": v}
    if isinstance(v, int):
        return {"int": str(v)}
    if isinstance(v, float):
        return {"float": repr(v)}
    if isinstance(v, str):
        return {"str": v}
    if v is None:
        return {"none": True}
    if isinstance(v, list):
        return {"list": [tagged(x) for x in v]}
    if isinstance(v, dict):
        return {"dict": [[k, tagged(x)] for k, x in v.items()]}
    raise TypeError(type(v))


def main(argv):
    if len(argv) >= 3 and argv[1] == "strings":
        text = render_strings(extract_strings(argv[2]))
        path = argv[3] if len(argv) > 3 else os.path.join(GEN_DIR, "strings.gen.ts")
    elif len(argv) >= 3 and argv[1] == "punkt":
        text = render_punkt(argv[2])
        path = argv[3] if len(argv) > 3 else os.path.join(GEN_DIR, "punkt.gen.ts")
    else:
        print(__doc__)
        return 2
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
    print(f"wrote {path} ({len(text.encode('utf-8'))} bytes)")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))