model.py

import re
import torch
from queue import Empty
from threading import Thread
from typing import Callable, Generator, List, Dict, Optional

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer,
    StoppingCriteria,
    StoppingCriteriaList,
)

from prompts import REASONING_SYSTEM_PROMPT
from tokens_control import protocol as token_protocol


CHAT_CONTROL_TOKENS = token_protocol.CHAT_CONTROL_TOKENS


def strip_chat_control_tokens(text: str | None, trim: bool = True) -> str:
    if text is None:
        return ""

    cleaned = str(text)
    for token in CHAT_CONTROL_TOKENS:
        cleaned = cleaned.replace(token, "")

    return cleaned.strip() if trim else cleaned


class StopOnSubstrings(StoppingCriteria):
    """
    Stops generation once any stop string appears in the decoded tail.
    """

    def __init__(
        self,
        tokenizer,
        stop_strings: List[str],
        prompt_length: int = 0,
        max_check_tokens: int = 80,
    ):
        self.tokenizer = tokenizer
        self.stop_strings = stop_strings
        self.prompt_length = prompt_length
        self.max_check_tokens = max_check_tokens

    def __call__(self, input_ids, scores, **kwargs) -> bool:
        generated_ids = input_ids[0][self.prompt_length:]
        if generated_ids.numel() == 0:
            return False

        tail_ids = generated_ids[-self.max_check_tokens:]
        tail = self.tokenizer.decode(tail_ids, skip_special_tokens=False)

        return any(stop in tail for stop in self.stop_strings)


class StopOnThinkingTokenLimit(StoppingCriteria):
    """
    Stops generation when the currently open <think> block reaches a token cap.
    """

    def __init__(
        self,
        tokenizer,
        token_limit: int,
        prompt_length: int = 0,
        initial_in_think: bool = False,
    ):
        self.tokenizer = tokenizer
        self.token_limit = max(0, int(token_limit))
        self.prompt_length = prompt_length
        self.initial_in_think = initial_in_think
        self.triggered = False

    def __call__(self, input_ids, scores, **kwargs) -> bool:
        generated_ids = input_ids[0][self.prompt_length:]
        if generated_ids.numel() == 0:
            return False

        generated = self.tokenizer.decode(generated_ids, skip_special_tokens=False)
        think_start = generated.rfind("<think>")
        think_end = generated.rfind("</think>")

        if think_start >= 0 and think_end < think_start:
            content = generated[think_start + len("<think>"):]
        elif self.initial_in_think and think_end < 0:
            content = generated
        else:
            return False

        count = len(
            self.tokenizer(
                content,
                add_special_tokens=False,
            )["input_ids"]
        )

        if count >= self.token_limit:
            self.triggered = True
            return True

        return False


class ResultTagSuppressor:
    """
    Suppresses model-generated <result> blocks from the visible stream.
    Runtime-injected <result> blocks are handled separately.
    """

    START = "<result>"
    END = "</result>"
    TAIL = max(len(START), len(END)) - 1

    def __init__(self):
        self.buffer = ""
        self.in_result = False
        self.suppressed = False

    def feed(self, chunk: str) -> str:
        self.buffer += chunk
        output = []

        while self.buffer:
            if self.in_result:
                end_index = self.buffer.find(self.END)

                if end_index < 0:
                    self.buffer = self.buffer[-self.TAIL:]
                    break

                self.buffer = self.buffer[end_index + len(self.END):]
                self.in_result = False
                continue

            start_index = self.buffer.find(self.START)

            if start_index >= 0:
                output.append(self.buffer[:start_index])
                self.buffer = self.buffer[start_index + len(self.START):]
                self.in_result = True
                self.suppressed = True
                continue

            if len(self.buffer) <= self.TAIL:
                break

            emit_upto = len(self.buffer) - self.TAIL
            output.append(self.buffer[:emit_upto])
            self.buffer = self.buffer[emit_upto:]
            break

        return "".join(output)

    def flush(self) -> str:
        if self.in_result:
            self.buffer = ""
            self.in_result = False
            return ""

        output = self.buffer
        self.buffer = ""
        return output


class EarlyAnswerSuppressor:
    """
    While retrieval has started but min_tool_rounds has not been reached,
    hides early <answer> attempts from the UI.
    """

    START = "<answer>"
    TAIL = len(START) - 1

    def __init__(self):
        self.buffer = ""
        self.blocked = False

    def feed(self, chunk: str) -> str:
        if self.blocked:
            return ""

        self.buffer += chunk
        start_index = self.buffer.find(self.START)

        if start_index >= 0:
            output = self.buffer[:start_index]
            self.buffer = ""
            self.blocked = True
            return output

        if len(self.buffer) <= self.TAIL:
            return ""

        emit_upto = len(self.buffer) - self.TAIL
        output = self.buffer[:emit_upto]
        self.buffer = self.buffer[emit_upto:]

        return output

    def flush(self) -> str:
        if self.blocked:
            self.buffer = ""
            return ""

        output = self.buffer
        self.buffer = ""
        return output


class StartTagSuppressor:
    """
    Suppresses a forbidden opening tag and everything generated after it.
    Used after retrieval is exhausted so leftover <query> spam never reaches UI.
    """

    def __init__(self, start_tags: List[str]):
        self.start_tags = start_tags
        self.tail = max((len(tag) for tag in start_tags), default=1) - 1
        self.buffer = ""
        self.blocked = False

    def feed(self, chunk: str) -> str:
        if self.blocked:
            return ""

        self.buffer += chunk
        starts = [
            index
            for tag in self.start_tags
            for index in [self.buffer.find(tag)]
            if index >= 0
        ]

        if starts:
            start_index = min(starts)
            output = self.buffer[:start_index]
            self.buffer = ""
            self.blocked = True
            return output

        if len(self.buffer) <= self.tail:
            return ""

        emit_upto = len(self.buffer) - self.tail
        output = self.buffer[:emit_upto]
        self.buffer = self.buffer[emit_upto:]

        return output

    def flush(self) -> str:
        if self.blocked:
            self.buffer = ""
            return ""

        output = self.buffer
        self.buffer = ""
        return output


class MultimodalQwen:
    QUERY_RE = token_protocol.QUERY_RE
    QUERY_FORBIDDEN_RE = token_protocol.QUERY_FORBIDDEN_RE
    QUERY_PREFIX_JUNK_RE = token_protocol.QUERY_PREFIX_JUNK_RE
    QUERY_LEAK_STOP_STRINGS = token_protocol.QUERY_LEAK_STOP_STRINGS
    RESULT_BLOCK_RE = token_protocol.RESULT_BLOCK_RE

    PLAIN_FINAL_RE = re.compile(
        r"(?is)(?:^|\n)\s*(?:final answer|answer)\s*:\s*(.+)$"
    )

    GENERATED_RESULT_CLAIM_RE = re.compile(
        r"""(?is)
        (
            \b(?:the\s+)?retrieved\s+result\s+(?:is|says|shows|states)\s*:
            |
            \b(?:the\s+)?retrieval\s+result\s+(?:is|says|shows|states)\s*:
            |
            \b(?:the\s+)?document\s+(?:says|states|shows|mentions)\s*:
            |
            \b(?:the\s+)?source\s+(?:says|states|shows|mentions)\s*:
            |
            \bresult\s*:
            |
            \bretrieved\s+result\s*:
            |
            \bdocument\s*:
        )
        """,
        re.VERBOSE,
    )

    FAKE_RETRIEVAL_CLAIM_RE = re.compile(
        r"""(?is)
        \b(
            retrieved\s+(?:result|text|document|evidence)
            |
            retrieval\s+(?:result|text|document|evidence)
            |
            document
            |
            source
            |
            evidence
        )\s+
        (?:is|says|states|shows|mentions|contains|indicates)\b
        """,
        re.VERBOSE,
    )

    BAD_META_QUERY_RE = re.compile(
        r"""(?is)^\s*
        (
            block\s+to\s+request\s+retrieval
            |query\s+block
            |request\s+retrieval
            |emit\s+(?:a\s+)?query
            |retrieve\s+it
            |look\s+it\s+up
        )
        [\s.]*$
        """,
        re.VERBOSE,
    )

    BAD_META_ANSWER_RE = re.compile(
        r"""(?is)
        <answer>\s*
        (
            .*?\b(?:should|need|must)\s+(?:retrieve|query|search|look\s+up)\b.*
            |
            .*?\bno\s+(?:document|context|result|retrieved\s+evidence)\s+(?:is|was)\s+present\b.*
            |
            .*?\bno\s+retrieved\s+document\s+evidence\s+was\s+found\b.*
            |
            this\s+is\s+a\s+(?:fact|factual|document|external|fictional).{0,240}
        )
        """,
        re.VERBOSE,
    )

    OFF_TOPIC_EVIDENCE_RE = re.compile(
        r"""(?is)
        (
            retrieved\s+(?:result|evidence|document).{0,160}
            (?:not\s+directly\s+applicable|not\s+directly\s+address|does\s+not\s+address|irrelevant|off-topic)
            |
            evidence\s+is\s+incomplete
            |
            retrieved\s+evidence\s+is\s+insufficient
        )
        """,
        re.VERBOSE,
    )

    RETRIEVAL_INTENT_RE = re.compile(
        r"""(?is)
        (
            \b(?:require|requires|required|need|needs|needed|must|should|have\s+to)\b
            .{0,160}?
            \b(?:retrieval|retrieve|query|lookup|search|document|documents|passage|passages|source|sources|evidence)\b
        )
        |
        \bno\s+(?:retrieval\s+)?result\s+is\s+present\b
        |
        \bno\s+direct\s+evidence\s+is\s+present\b
        |
        \bmust\s+create\s+(?:one|a\s+(?:query|retrieval|lookup))\b
        |
        \bneed\s+to\s+check\s+(?:if\s+)?(?:the\s+)?documents?\b
        """,
        re.VERBOSE,
    )

    NO_RETRIEVAL_INTENT_RE = re.compile(
        r"""(?is)
        \b(?:do\s+not|don't|does\s+not|doesn't|no|not)\s+
        (?:need|require|requires|required)\s+
        .{0,80}?
        \b(?:retrieval|retrieve|query|lookup|search)\b
        """,
        re.VERBOSE,
    )

    TOOL_PROMISE_RE = re.compile(
        r"""(?is)
        \b(
            i\s+will
            |i'll
            |i\s+am\s+going\s+to
            |i'm\s+going\s+to
            |let\s+me
            |i\s+should
            |i\s+need\s+to
            |i\s+have\s+to
        )\s+
        (?:now\s+|first\s+|try\s+to\s+)?
        (search|retrieve|query|look\s+up|lookup|check\s+(?:the\s+)?documents?|find\s+(?:evidence|sources?))
        \b
        """,
        re.VERBOSE,
    )

    PLAIN_ANSWER_ATTEMPT_RE = re.compile(
        r"(?is)(?:^|\n)\s*(?:final\s+answer|answer)\s*:\s*[\s\S]*$"
    )

    def __init__(
        self,
        model_name: str = "Qwen/Qwen2.5-3B-Instruct",
        device: Optional[str] = None,
        torch_dtype="auto",
    ):
        self.model_name = model_name

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True,
        )

        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch_dtype,
            trust_remote_code=True,
        ).to(self.device)

        self.model.eval()
        self.prompt = REASONING_SYSTEM_PROMPT

    # ---------------------------------------------------------------------
    # Basic helpers
    # ---------------------------------------------------------------------

    def _stop_strings(self, stop_strings: Optional[List[str]] = None) -> List[str]:
        return list(stop_strings or [])

    def _eos_token_ids(self):
        return self.tokenizer.eos_token_id

    def coerce_plain_final_answer(self, text: str) -> tuple[str, bool]:
        """
        Do not fabricate <answer> in Python.
        The model must emit protocol tags itself.
        """
        return text, False

    @staticmethod
    def has_open_tag(text: str, tag: str) -> bool:
        return text.rfind(f"<{tag}>") > text.rfind(f"</{tag}>")

    @staticmethod
    def first_terminal_tag(text: str) -> tuple[str | None, int]:
        """
        Returns the first completed terminal protocol block.

        If model emits:
            <query>real query</query><query><answer>junk</answer>

        runtime keeps only:
            <query>real query</query>

        and starts retrieval immediately.
        """
        text = text or ""
        candidates: list[tuple[int, str]] = []

        for close_tag, kind in (
            ("</query>", "query"),
            ("</answer>", "answer"),
        ):
            index = text.find(close_tag)
            if index >= 0:
                candidates.append((index + len(close_tag), kind))

        if not candidates:
            return None, -1

        end_index, kind = min(candidates, key=lambda item: item[0])
        return kind, end_index

    def _slice_to_first_terminal(self, text: str) -> tuple[str, str | None]:
        terminal_kind, terminal_end = self.first_terminal_tag(text)
        if terminal_kind is None:
            return text, None
        return text[:terminal_end], terminal_kind

    # ---------------------------------------------------------------------
    # Prompt formatting
    # ---------------------------------------------------------------------

    def _messages(
        self,
        user_prompt: str,
        system_prompt: Optional[str] = None,
    ) -> List[Dict[str, str]]:
        active_prompt = str(system_prompt).strip() if system_prompt is not None else ""
        active_prompt = active_prompt or self.prompt

        return [
            {"role": "system", "content": active_prompt},
            {"role": "user", "content": strip_chat_control_tokens(user_prompt)},
        ]

    def _apply_chat_template(self, messages: List[Dict[str, str]]):
        inputs = self.tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            tokenize=True,
            return_dict=True,
            return_tensors="pt",
        )

        return {k: v.to(self.device) for k, v in inputs.items()}

    def _build_raw_prompt(
        self,
        user_prompt: str,
        system_prompt: Optional[str] = None,
    ) -> str:
        return self.tokenizer.apply_chat_template(
            self._messages(user_prompt, system_prompt=system_prompt),
            add_generation_prompt=True,
            tokenize=False,
        )

    def _tokenize_raw(self, raw_prompt: str):
        inputs = self.tokenizer(
            raw_prompt,
            return_tensors="pt",
            add_special_tokens=False,
        )

        return {k: v.to(self.device) for k, v in inputs.items()}

    # ---------------------------------------------------------------------
    # Basic generation
    # ---------------------------------------------------------------------

    @torch.inference_mode()
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        system_prompt: Optional[str] = None,
    ) -> str:
        messages = self._messages(prompt, system_prompt=system_prompt)
        inputs = self._apply_chat_template(messages)

        do_sample = temperature > 0

        output_ids = self.model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            top_p=top_p if do_sample else None,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self._eos_token_ids(),
        )

        new_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
        return strip_chat_control_tokens(
            self.tokenizer.decode(new_ids, skip_special_tokens=True)
        )

    @torch.inference_mode()
    def _generate_raw(
        self,
        raw_prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        stop_strings: Optional[List[str]] = None,
    ) -> str:
        inputs = self._tokenize_raw(raw_prompt)
        do_sample = temperature > 0

        stop_strings = self._stop_strings(stop_strings)
        stopping_criteria = None

        if stop_strings:
            stopping_criteria = StoppingCriteriaList([
                StopOnSubstrings(
                    self.tokenizer,
                    stop_strings,
                    prompt_length=inputs["input_ids"].shape[-1],
                )
            ])

        output_ids = self.model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            top_p=top_p if do_sample else None,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self._eos_token_ids(),
            stopping_criteria=stopping_criteria,
        )

        new_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
        return strip_chat_control_tokens(
            self.tokenizer.decode(new_ids, skip_special_tokens=True)
        )

    def stream_raw(
        self,
        raw_prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        stop_strings: Optional[List[str]] = None,
        thinking_token_limit: Optional[int] = None,
        thinking_limit_state: Optional[dict] = None,
    ) -> Generator[str, None, None]:
        inputs = self._tokenize_raw(raw_prompt)

        streamer = TextIteratorStreamer(
            self.tokenizer,
            skip_prompt=True,
            skip_special_tokens=True,
            timeout=1.0,
        )

        do_sample = temperature > 0
        stop_strings = self._stop_strings(stop_strings)

        stopping_items = []

        if stop_strings:
            stopping_items.append(
                StopOnSubstrings(
                    self.tokenizer,
                    stop_strings,
                    prompt_length=inputs["input_ids"].shape[-1],
                )
            )

        thinking_stopper = None

        if thinking_token_limit is not None:
            thinking_stopper = StopOnThinkingTokenLimit(
                self.tokenizer,
                thinking_token_limit,
                prompt_length=inputs["input_ids"].shape[-1],
                initial_in_think=self.has_open_tag(raw_prompt, "think"),
            )
            stopping_items.append(thinking_stopper)

        stopping_criteria = StoppingCriteriaList(stopping_items) if stopping_items else None

        kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else None,
            top_p=top_p if do_sample else None,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self._eos_token_ids(),
            stopping_criteria=stopping_criteria,
        )

        generation_error = []

        def run_generation():
            try:
                with torch.inference_mode():
                    self.model.generate(**kwargs)
            except BaseException as exc:
                generation_error.append(exc)
                try:
                    streamer.on_finalized_text("", stream_end=True)
                except Exception:
                    pass

        thread = Thread(target=run_generation, daemon=True)
        thread.start()

        while True:
            try:
                chunk = next(streamer)
            except StopIteration:
                break
            except Empty:
                if not thread.is_alive():
                    break
                continue

            chunk = strip_chat_control_tokens(chunk, trim=False)
            if chunk:
                yield chunk

        thread.join()

        if thinking_limit_state is not None:
            thinking_limit_state["triggered"] = bool(
                thinking_stopper and thinking_stopper.triggered
            )

        if generation_error:
            raise RuntimeError("Model generation failed") from generation_error[0]

    def stream(
        self,
        prompt: str,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        thinking_token_limit: Optional[int] = None,
        system_prompt: Optional[str] = None,
    ) -> Generator[str, None, None]:
        raw_prompt = self._build_raw_prompt(prompt, system_prompt=system_prompt)
        thinking_state: dict = {}
        generated = ""

        for chunk in self.stream_raw(
            raw_prompt=raw_prompt,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            thinking_token_limit=thinking_token_limit,
            thinking_limit_state=thinking_state,
        ):
            generated += chunk
            yield chunk

        raw_prompt += generated

        if thinking_state.get("triggered") and self.has_open_tag(raw_prompt, "think"):
            close_think = "\n</think>\n"
            raw_prompt += close_think
            yield close_think

            forced = (
                "[Runtime instruction: Thinking budget reached. "
                "Produce the final <answer>...</answer> now.]\n"
            )
            raw_prompt += forced
            yield forced

            for chunk in self.stream_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=["</answer>", "<result>", "</result>", "</query>", "<query>", "<think>"],
            ):
                yield chunk

    # ---------------------------------------------------------------------
    # Protocol parsing / sanitizing
    # ---------------------------------------------------------------------

    def extract_query(self, text: str) -> Optional[str]:
        return token_protocol.extract_query(text)

    def extract_open_query(self, text: str) -> Optional[str]:
        """
        Only catches a truly unclosed final <query> block.
        """
        text = text or ""

        lower = text.lower()
        last_open = lower.rfind("<query>")
        last_close = lower.rfind("</query>")

        if last_open < 0:
            return None

        if last_close > last_open:
            return None

        return text[last_open + len("<query>"):].strip()

    def sanitize_query_text(self, query: str) -> str:
        query = str(query or "").strip()

        # Remove repeated/nested query openings.
        query = re.sub(r"(?is)(?:<query>\s*)+", "", query).strip()

        query = re.split(
            r"(?is)</?think>|</?answer>|</?result>|<query>|</query>|final\s+answer\s*:|answer\s*:",
            query,
            maxsplit=1,
        )[0].strip()

        query = re.sub(
            r"(?is)\b(?:retrieved\s+evidence\s+only|evidence\s+only|use\s+retrieved\s+evidence\s+only)\b\.?",
            "",
            query,
        ).strip()

        query = self.QUERY_PREFIX_JUNK_RE.sub("", query).strip()
        query = re.sub(r"(?is)^(query|search|lookup)\s*:\s*", "", query).strip()

        return query

    def has_malformed_query(self, text: str) -> bool:
        open_query = self.extract_open_query(text)
        if open_query is None:
            return False

        clean_query = self.sanitize_query_text(open_query)

        if not clean_query:
            return True

        if self.QUERY_FORBIDDEN_RE.search(open_query):
            return True

        return False

    @staticmethod
    def malformed_query_correction() -> str:
        return token_protocol.malformed_query_correction()

    def strip_generated_result_blocks(self, text: str) -> tuple[str, bool]:
        cleaned, count = self.RESULT_BLOCK_RE.subn("", text)
        return cleaned, count > 0

    def sanitize_result(self, result: str) -> str:
        if result is None:
            result = ""

        result = str(result)
        result = strip_chat_control_tokens(result)

        replacements = {
            "<think>": "[think]",
            "</think>": "[/think]",
            "<query>": "[query]",
            "</query>": "[/query]",
            "<result>": "[result]",
            "</result>": "[/result]",
            "<answer>": "[answer]",
            "</answer>": "[/answer]",
        }

        for old, new in replacements.items():
            result = result.replace(old, new)

        return result.strip()

    def format_query_result(self, query: str, result: str) -> str:
        safe_query = self.sanitize_result(query)
        safe_result = self.sanitize_result(result)

        if not safe_result:
            safe_result = "No retrieval results found."

        return (
            "Query:\n"
            f"{safe_query}\n\n"
            "Retrieved Result:\n"
            f"{safe_result}"
        )

    def retrieval_limit_result(self, max_tool_rounds: int) -> str:
        return self.format_query_result(
            query="retrieval limit",
            result=(
                f"Retrieval limit reached ({max_tool_rounds}). "
                "No more retrieval documents will be provided for this answer. "
                "Use the evidence already shown, say when it is insufficient, "
                "and finish with answer."
            ),
        )

    def is_bad_meta_query(self, query: str) -> bool:
        return self.BAD_META_QUERY_RE.search(query or "") is not None

    def has_bad_meta_answer(self, text: str) -> bool:
        return self.BAD_META_ANSWER_RE.search(text or "") is not None

    def says_evidence_off_topic_or_incomplete(self, text: str) -> bool:
        return self.OFF_TOPIC_EVIDENCE_RE.search(text or "") is not None

    def has_generated_result_claim(self, text: str, tool_rounds: int) -> bool:
        if not text:
            return False

        if tool_rounds > 0:
            return False

        if "<result>" in text or "</result>" in text:
            return False

        return self.GENERATED_RESULT_CLAIM_RE.search(text) is not None

    def strip_generated_result_claim(self, text: str) -> str:
        match = self.GENERATED_RESULT_CLAIM_RE.search(text or "")
        if not match:
            return text

        return text[:match.start()].rstrip()

    def has_fake_retrieval_claim(self, text: str, tool_rounds: int) -> bool:
        if tool_rounds > 0:
            return False

        if "<result>" in text or "</result>" in text:
            return False

        return self.FAKE_RETRIEVAL_CLAIM_RE.search(text or "") is not None

    def strip_fake_retrieval_claim(self, text: str) -> str:
        match = self.FAKE_RETRIEVAL_CLAIM_RE.search(text or "")
        if not match:
            return text

        return text[:match.start()].rstrip()
    
    @staticmethod
    def retrieval_exhausted_answer_instruction(max_tool_rounds: int) -> str:
        return (
            "\n[Runtime instruction: Retrieval budget is exhausted. "
            f"Maximum retrieval results already reached: {max_tool_rounds}. "
            "Do not emit <query> again. Do not say you will search, retrieve, look up, "
            "check documents, or request more evidence. "
            "Produce exactly one final <answer>...</answer> now, using only the existing "
            "<result> evidence. If the evidence is insufficient, say it is insufficient.]\n"
        )

    @staticmethod
    def protocol_correction(force_answer: bool = False) -> str:
        if force_answer:
            return (
                "\n[Runtime instruction: No more retrieval results are available. "
                "Answer from the evidence already shown.]\n"
            )

        return (
            "\n[Runtime instruction: Do not create retrieval results yourself. "
            "Either request one focused lookup or answer directly.]\n"
        )

    @staticmethod
    def fake_retrieval_claim_correction() -> str:
        return (
            "\n[Runtime instruction: You claimed retrieval/document evidence exists, "
            "but no <result> block has been provided by runtime. Do not invent retrieved content. "
            "Emit exactly one focused <query>...</query> next.]\n"
        )

    @staticmethod
    def generated_result_claim_correction() -> str:
        return (
            "\n[Runtime instruction: You cannot create or describe retrieved results yourself. "
            "Only the runtime may provide retrieval results inside <result>...</result>. "
            "If evidence is needed, emit exactly one focused <query>...</query>. "
            "If evidence is not needed, answer directly without mentioning retrieved results.]\n"
        )

    @staticmethod
    def force_query_correction() -> str:
        return (
            "\n[Runtime instruction: You promised or stated that retrieval/search is needed, "
            "but you did not emit a valid <query> block. Natural-language promises like "
            "'I will search' do not execute tools. Emit exactly one concrete "
            "<query>...</query> next.]\n"
        )

    def minimum_retrieval_correction(self, tool_rounds: int, min_tool_rounds: int) -> str:
        remaining = max(0, min_tool_rounds - tool_rounds)

        return (
            "\n[Runtime instruction: Retrieval has started, but not enough retrieval results "
            "have been collected yet. "
            f"Current retrieval results: {tool_rounds}. "
            f"Minimum required retrieval results: {min_tool_rounds}. "
            f"Remaining retrieval result(s): {remaining}. "
            "Request another focused retrieval query next.]\n"
        )

    def declares_retrieval_intent(self, text: str) -> bool:
        if not text:
            return False

        if self.extract_query(text) is not None:
            return False

        if self.NO_RETRIEVAL_INTENT_RE.search(text):
            return False

        return (
            self.RETRIEVAL_INTENT_RE.search(text) is not None
            or self.TOOL_PROMISE_RE.search(text) is not None
        )

    def strip_answer_attempts(self, text: str) -> str:
        text = re.sub(r"(?is)<answer>[\s\S]*?(?:</answer>|$)", "", text)
        text = self.PLAIN_ANSWER_ATTEMPT_RE.sub("", text)
        return text.rstrip()

    # ---------------------------------------------------------------------
    # Non-streaming query loop
    # ---------------------------------------------------------------------

    def generate_with_query(
        self,
        prompt: str,
        query_fn: Callable[[str], str],
        max_tool_rounds: int = 4,
        min_tool_rounds: int = 2,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        system_prompt: Optional[str] = None,
    ) -> str:
        raw_prompt = self._build_raw_prompt(prompt, system_prompt=system_prompt)
        transcript = ""

        tool_rounds = 0
        retrieval_started = False
        force_answer = False
        illegal_result_count = 0

        min_tool_rounds = max(0, int(min_tool_rounds))
        max_tool_rounds = max(min_tool_rounds, int(max_tool_rounds))
        max_iterations = max_tool_rounds + 8

        def inject_query_result(query: str) -> str:
            nonlocal tool_rounds, retrieval_started

            tool_rounds += 1
            retrieval_started = True

            result = self.format_query_result(query, query_fn(query))

            return (
                "\n"
                f"<result>{result}</result>\n"
                "<think>\n"
            )

        for _ in range(max_iterations):
            need_more_retrieval = retrieval_started and tool_rounds < min_tool_rounds

            stop_strings = token_protocol.query_generation_stop_strings(
                need_more_retrieval=need_more_retrieval,
            )

            if self.has_open_tag(raw_prompt, "think"):
                stop_strings = list(dict.fromkeys([
                    "</think>",
                    "<answer>",
                    "<result>",
                    "</result>",
                    *stop_strings,
                ]))

            generated = self._generate_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=stop_strings,
            )

            generated, had_illegal_result = self.strip_generated_result_blocks(generated)
            generated, terminal_kind = self._slice_to_first_terminal(generated)

            # Query wins immediately after </query>.
            if terminal_kind == "query":
                query = self.extract_query(generated) or ""
                clean_query = self.sanitize_query_text(query)

                if not clean_query or self.is_bad_meta_query(clean_query):
                    cleaned = generated.split("<query>", 1)[0].rstrip()
                    raw_prompt += cleaned
                    transcript += cleaned

                    correction = (
                        "\n[Runtime instruction: The query was malformed or not a real lookup. "
                        "Emit exactly one concrete <query>...</query> containing the subject and missing fact.]\n"
                    )
                    raw_prompt += correction
                    transcript += correction
                    continue

                if force_answer or tool_rounds >= max_tool_rounds:
                    cleaned = generated.split("<query>", 1)[0].rstrip()
                    raw_prompt += cleaned
                    transcript += cleaned
                    injected = self.retrieval_exhausted_answer_instruction(max_tool_rounds)
                    raw_prompt += injected
                    transcript += injected
                    force_answer = True
                    continue

                raw_prompt += generated
                transcript += generated

                injected = inject_query_result(clean_query)
                raw_prompt += injected
                transcript += injected
                continue

            # Completed answer.
            if terminal_kind == "answer":
                if self.has_bad_meta_answer(generated) or self.declares_retrieval_intent(generated):
                    cleaned = self.strip_answer_attempts(generated)
                    raw_prompt += cleaned
                    transcript += cleaned

                    correction = self.force_query_correction()
                    raw_prompt += correction
                    transcript += correction
                    continue

                if retrieval_started and tool_rounds < min_tool_rounds:
                    cleaned = self.strip_answer_attempts(generated)
                    raw_prompt += cleaned
                    transcript += cleaned

                    correction = self.minimum_retrieval_correction(
                        tool_rounds=tool_rounds,
                        min_tool_rounds=min_tool_rounds,
                    )
                    raw_prompt += correction
                    transcript += correction
                    continue

                raw_prompt += generated
                transcript += generated
                break

            # Open query / malformed query.
            if self.has_malformed_query(generated):
                raw_query = self.extract_open_query(generated) or ""
                clean_query = self.sanitize_query_text(raw_query)

                raw_prompt += re.sub(r"(?is)<query>[\s\S]*$", "", generated).rstrip()
                transcript += re.sub(r"(?is)<query>[\s\S]*$", "", generated).rstrip()

                if clean_query and not self.is_bad_meta_query(clean_query):
                    normalized = f"<query>{clean_query}</query>"
                    raw_prompt += normalized
                    transcript += normalized

                    if not force_answer and tool_rounds < max_tool_rounds:
                        injected = inject_query_result(clean_query)
                        raw_prompt += injected
                        transcript += injected
                        continue

                correction = self.malformed_query_correction()
                raw_prompt += correction
                transcript += correction
                continue

            if self.has_fake_retrieval_claim(generated, tool_rounds):
                generated = self.strip_fake_retrieval_claim(generated)

                raw_prompt += generated
                transcript += generated

                correction = self.fake_retrieval_claim_correction()
                raw_prompt += correction
                transcript += correction
                continue

            if self.has_generated_result_claim(generated, tool_rounds):
                generated = self.strip_generated_result_claim(generated)

                raw_prompt += generated
                transcript += generated

                correction = self.generated_result_claim_correction()
                raw_prompt += correction
                transcript += correction
                continue

            if self.declares_retrieval_intent(generated):
                generated = self.strip_answer_attempts(generated)

                raw_prompt += generated
                transcript += generated

                correction = self.force_query_correction()
                raw_prompt += correction
                transcript += correction
                continue

            if self.has_open_tag(raw_prompt, "think") and "<answer>" in generated:
                before_answer = generated.split("<answer>", 1)[0].rstrip()
                raw_prompt += before_answer
                transcript += before_answer
                raw_prompt += "\n</think>\n"
                transcript += "\n</think>\n"
                raw_prompt += self.protocol_correction(force_answer=True)
                transcript += self.protocol_correction(force_answer=True)
                continue

            if had_illegal_result:
                illegal_result_count += 1
                force_answer = force_answer or illegal_result_count >= 2
                correction = self.protocol_correction(force_answer=force_answer)
                raw_prompt += correction
                transcript += correction
                continue

            raw_prompt += generated
            transcript += generated

            if generated.strip().endswith("</think>"):
                if retrieval_started and tool_rounds >= min_tool_rounds:
                    instruction = (
                        "\n[Runtime instruction: Retrieved evidence is available. "
                        "Produce the final <answer>...</answer> now.]\n"
                    )
                    raw_prompt += instruction
                    transcript += instruction
                    continue

                continue

            if retrieval_started and tool_rounds >= min_tool_rounds:
                instruction = (
                    "\n[Runtime instruction: Retrieved evidence is available. "
                    "Produce the final <answer>...</answer> now.]\n"
                )
                raw_prompt += instruction
                transcript += instruction
                continue

            if not retrieval_started:
                break

        else:
            correction = (
                "\n[Runtime instruction: The retrieval loop reached its safety limit. "
                "Do not request more retrieval. Emit a final <answer>...</answer> now, "
                "using only the evidence already shown.]\n"
            )
            raw_prompt += correction
            transcript += correction

            generated = self._generate_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=["</answer>", "<result>", "</result>", "</query>", "<query>"],
            )
            generated, _ = self._slice_to_first_terminal(generated)
            generated = generated.split("<query>", 1)[0]
            generated = generated.split("<result>", 1)[0]

            raw_prompt += generated
            transcript += generated

        return transcript

    # ---------------------------------------------------------------------
    # Streaming query loop
    # ---------------------------------------------------------------------

    def stream_with_query(
        self,
        prompt: str,
        query_fn: Callable[[str], str],
        max_tool_rounds: int = 3,
        min_tool_rounds: int = 2,
        max_new_tokens: int = 512,
        temperature: float = 0.2,
        top_p: float = 0.9,
        show_injected_result: bool = True,
        thinking_token_limit: Optional[int] = None,
        system_prompt: Optional[str] = None,
    ) -> Generator[str, None, None]:
        raw_prompt = self._build_raw_prompt(prompt, system_prompt=system_prompt)

        min_tool_rounds = max(0, int(min_tool_rounds))
        max_tool_rounds = max(min_tool_rounds, int(max_tool_rounds))

        tool_rounds = 0
        retrieval_started = False
        force_answer = False
        final_answer_requested = False
        illegal_result_count = 0
        last_query: str | None = None

        max_iterations = max_tool_rounds + 8

        def close_open_think_if_needed(text: str) -> str:
            if self.has_open_tag(text, "think"):
                return "\n</think>\n"

            if not text.endswith("\n"):
                return "\n"

            return ""

        def visible_result_block(result: str) -> str:
            return f"<result>{result}</result>\n"

        def inject_results_for_query(query: str, minimum_fill: bool) -> tuple[str, str]:
            nonlocal tool_rounds

            prompt_parts = []
            visible_parts = []

            needed = max(1, min_tool_rounds - tool_rounds) if minimum_fill else 1
            can_fetch = max(0, max_tool_rounds - tool_rounds)
            fetch_count = min(needed, can_fetch)

            for _ in range(fetch_count):
                tool_rounds += 1
                result = self.format_query_result(query, query_fn(query))
                block = visible_result_block(result)

                prompt_parts.append(block)
                visible_parts.append(block)

            return "".join(prompt_parts), "".join(visible_parts)

        def request_final_answer_instruction() -> str:
            return (
                "\n[Runtime instruction: Retrieved evidence is available. "
                "Do not request more retrieval unless a genuinely different lookup is required. "
                "Continue with a normal reasoning block if needed, then produce the final "
                "<answer>...</answer>.]\n"
            )

        for _ in range(max_iterations):
            generated = ""
            thinking_state: dict = {}
            suppressor = ResultTagSuppressor()
            query_guard = StartTagSuppressor(["<query>"]) if force_answer else None

            block_early_answer = retrieval_started and tool_rounds < min_tool_rounds
            answer_guard = EarlyAnswerSuppressor() if block_early_answer else None

            need_more_retrieval = retrieval_started and tool_rounds < min_tool_rounds

            if force_answer:
                # Answer-only mode. Stop as soon as the model tries to open another query.
                stop_strings = ["</answer>", "<query>", "<result>", "</result>"]
            else:
                stop_strings = token_protocol.query_generation_stop_strings(
                    need_more_retrieval=need_more_retrieval,
                )

            if self.has_open_tag(raw_prompt, "think"):
                stop_strings = list(dict.fromkeys([
                    "</think>",
                    "<answer>",
                    "<result>",
                    "</result>",
                    *stop_strings,
                ]))

            for chunk in self.stream_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=stop_strings,
                thinking_token_limit=thinking_token_limit,
                thinking_limit_state=thinking_state,
            ):
                filtered_chunk = suppressor.feed(chunk)

                if not filtered_chunk:
                    continue

                if query_guard is not None:
                    filtered_chunk = query_guard.feed(filtered_chunk)

                if not filtered_chunk:
                    continue

                if answer_guard is not None:
                    visible_chunk = answer_guard.feed(filtered_chunk)
                else:
                    visible_chunk = filtered_chunk

                if visible_chunk:
                    generated += visible_chunk
                    yield visible_chunk

            filtered_tail = suppressor.flush()

            if filtered_tail:
                if query_guard is not None:
                    filtered_tail = query_guard.feed(filtered_tail)

                if answer_guard is not None:
                    visible_tail = answer_guard.feed(filtered_tail)
                else:
                    visible_tail = filtered_tail

                if visible_tail:
                    generated += visible_tail
                    yield visible_tail

            if query_guard is not None:
                visible_tail = query_guard.flush()

                if visible_tail:
                    generated += visible_tail
                    yield visible_tail

            if answer_guard is not None:
                visible_tail = answer_guard.flush()

                if visible_tail:
                    generated += visible_tail
                    yield visible_tail

            think_limit_hit = bool(thinking_state.get("triggered"))

            if think_limit_hit and self.has_open_tag(raw_prompt + generated, "think"):
                close_think = "\n</think>\n"
                generated += close_think
                yield close_think

            generated, had_illegal_result = self.strip_generated_result_blocks(generated)
            generated, terminal_kind = self._slice_to_first_terminal(generated)
            if force_answer and (("<query>" in generated) or (query_guard is not None and query_guard.blocked)):
                before_query = generated.split("<query>", 1)[0].rstrip()

                if before_query:
                    raw_prompt += before_query

                raw_prompt += close_open_think_if_needed(raw_prompt)
                raw_prompt += self.retrieval_exhausted_answer_instruction(max_tool_rounds)

                final_answer_requested = True
                continue

            # -------------------------------------------------------------
            # CRITICAL FIX:
            # If </query> is the first completed terminal block,
            # start retrieval immediately. Do not inspect later <answer> junk.
            # -------------------------------------------------------------
            if terminal_kind == "query":
                query = self.extract_query(generated) or ""
                clean_query = self.sanitize_query_text(query)

                if force_answer or tool_rounds >= max_tool_rounds:
                    before_query = generated.split("<query>", 1)[0].rstrip()

                    if before_query:
                        raw_prompt += before_query

                    raw_prompt += close_open_think_if_needed(raw_prompt)
                    raw_prompt += self.retrieval_exhausted_answer_instruction(max_tool_rounds)

                    force_answer = True
                    final_answer_requested = True
                    continue

                raw_prompt += generated

                if not clean_query or self.is_bad_meta_query(clean_query):
                    raw_prompt += (
                        "\n[Runtime instruction: The query was malformed or not a real lookup. "
                        "Emit exactly one concrete <query>...</query> containing the subject and missing fact.]\n"
                    )
                    continue

                retrieval_started = True
                last_query = clean_query

                close_think = close_open_think_if_needed(raw_prompt)

                prompt_injected, visible_injected = inject_results_for_query(
                    clean_query,
                    minimum_fill=(tool_rounds < min_tool_rounds),
                )

                prompt_injected = close_think + prompt_injected + "<think>\n"
                visible_injected = close_think + visible_injected + "<think>\n"

                raw_prompt += prompt_injected

                if show_injected_result and visible_injected.strip():
                    yield visible_injected

                continue

            # Completed answer.
            if terminal_kind == "answer":
                if self.has_bad_meta_answer(generated) or self.declares_retrieval_intent(generated):
                    cleaned = self.strip_answer_attempts(generated)
                    raw_prompt += cleaned
                    raw_prompt += self.force_query_correction()
                    continue

                if retrieval_started and tool_rounds < min_tool_rounds:
                    cleaned = self.strip_answer_attempts(generated)
                    raw_prompt += cleaned

                    if last_query and tool_rounds < max_tool_rounds:
                        prompt_injected, visible_injected = inject_results_for_query(
                            last_query,
                            minimum_fill=True,
                        )
                        raw_prompt += prompt_injected

                        if show_injected_result and visible_injected:
                            yield visible_injected

                        continue

                    raw_prompt += self.minimum_retrieval_correction(
                        tool_rounds=tool_rounds,
                        min_tool_rounds=min_tool_rounds,
                    )
                    continue

                raw_prompt += generated
                break

            # Open query / malformed query.
            if self.has_malformed_query(generated):
                raw_query = self.extract_open_query(generated) or ""
                clean_query = self.sanitize_query_text(raw_query)

                generated_without_open_query = re.sub(
                    r"(?is)<query>[\s\S]*$",
                    "",
                    generated,
                ).rstrip()

                raw_prompt += generated_without_open_query

                if clean_query and not self.is_bad_meta_query(clean_query):
                    normalized = f"<query>{clean_query}</query>"
                    raw_prompt += normalized
                    yield normalized

                    if not force_answer and tool_rounds < max_tool_rounds:
                        retrieval_started = True
                        last_query = clean_query

                        close_think = close_open_think_if_needed(raw_prompt)

                        prompt_injected, visible_injected = inject_results_for_query(
                            clean_query,
                            minimum_fill=(tool_rounds < min_tool_rounds),
                        )

                        prompt_injected = close_think + prompt_injected + "<think>\n"
                        visible_injected = close_think + visible_injected + "<think>\n"

                        raw_prompt += prompt_injected

                        if show_injected_result and visible_injected.strip():
                            yield visible_injected

                        continue

                raw_prompt += self.malformed_query_correction()
                continue

            # Fake retrieval claims before runtime provided results.
            if self.has_fake_retrieval_claim(generated, tool_rounds):
                generated = self.strip_fake_retrieval_claim(generated)
                raw_prompt += generated
                raw_prompt += self.fake_retrieval_claim_correction()
                continue

            if self.has_generated_result_claim(generated, tool_rounds):
                generated = self.strip_generated_result_claim(generated)
                raw_prompt += generated
                raw_prompt += self.generated_result_claim_correction()
                continue

            # Retrieval intent without a valid query.
            if self.declares_retrieval_intent(generated):
                cleaned = self.strip_answer_attempts(generated)
                raw_prompt += cleaned
                raw_prompt += self.force_query_correction()
                continue

            # If answer leaked inside a runtime-opened think block, repair it.
            if self.has_open_tag(raw_prompt, "think") and "<answer>" in generated:
                before_answer = generated.split("<answer>", 1)[0].rstrip()

                raw_prompt += before_answer
                raw_prompt += "\n</think>\n"
                raw_prompt += self.retrieval_exhausted_answer_instruction(max_tool_rounds)
                final_answer_requested = True
                continue

            if had_illegal_result or suppressor.suppressed:
                illegal_result_count += 1
                force_answer = force_answer or illegal_result_count >= 2
                raw_prompt += self.protocol_correction(force_answer=force_answer)
                continue

            if think_limit_hit:
                raw_prompt += generated
                raw_prompt += (
                    "\n[Runtime instruction: Thinking budget reached. "
                    "Produce the final <answer>...</answer> now.]\n"
                )
                final_answer_requested = True
                continue

            raw_prompt += generated

            if generated.strip().endswith("</think>"):
                if retrieval_started and tool_rounds >= min_tool_rounds:
                    raw_prompt += request_final_answer_instruction()
                    final_answer_requested = True
                    continue

                # Before retrieval starts, do not force answer.
                continue

            if retrieval_started and tool_rounds >= min_tool_rounds and not final_answer_requested:
                if force_answer or tool_rounds >= max_tool_rounds:
                    raw_prompt += self.retrieval_exhausted_answer_instruction(max_tool_rounds)
                else:
                    raw_prompt += request_final_answer_instruction()

                final_answer_requested = True
                continue

            if not retrieval_started:
                break

            if not final_answer_requested:
                raw_prompt += request_final_answer_instruction()
                final_answer_requested = True
                continue

            # Final fallback pass.
            fallback_guard = StartTagSuppressor(["<query>", "<result>"])
            for chunk in self.stream_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=["</answer>", "<query>", "<result>", "</result>"],
            ):
                visible_chunk = fallback_guard.feed(chunk)

                if visible_chunk:
                    yield visible_chunk

            visible_tail = fallback_guard.flush()
            if visible_tail:
                yield visible_tail

            break

        else:
            raw_prompt += request_final_answer_instruction()

            fallback_guard = StartTagSuppressor(["<query>", "<result>"])
            for chunk in self.stream_raw(
                raw_prompt=raw_prompt,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                stop_strings=["</answer>", "<query>", "<result>", "</result>", "</query>"],
            ):
                visible_chunk = fallback_guard.feed(chunk)

                if visible_chunk:
                    yield visible_chunk

            visible_tail = fallback_guard.flush()
            if visible_tail:
                yield visible_tail