agentic_rag/agent/react_parser.py

"""Parser for ReAct outputs — extracts Thought, Action, and Final Answer."""

import json
import re
from dataclasses import dataclass
from typing import Optional


@dataclass
class ReActStep:
    """A single step of the ReAct loop."""
    thought: str = ""
    action: str = ""
    action_input: dict = None
    is_final: bool = False
    final_answer: str = ""
    raw_text: str = ""

    def __post_init__(self):
        if self.action_input is None:
            self.action_input = {}


def parse_react_output(text: str, tool_names: list[str] | None = None) -> ReActStep:
    """Parse the LLM output into a ReAct step.

    Handles these patterns:
    - Thought: ...
    - Action: tool_name
    - Action Input: {...}
    OR
    - Thought: ...
    - Final Answer: ...
    """
    # Truncate repetitive output — if the model loops on the same sentence,
    # cut at the first repetition to keep only useful content.
    text = _trim_repetition(text)

    # Strip any residual reasoning (<think> blocks) — reasoning text may
    # contain rehearsed "Thought:/Action:" lines that must NOT be parsed as
    # real actions. Providers normally strip this already; this is a safety net.
    from agentic_rag.services.llm.base import strip_reasoning
    text = strip_reasoning(text)

    result = ReActStep(raw_text=text)

    # Extract Thought
    thought_match = re.search(r'Thought:\s*(.+?)(?=\n(?:Action|Final Answer)|$)', text, re.DOTALL)
    if thought_match:
        result.thought = thought_match.group(1).strip()

    # Check for Final Answer first
    final_match = re.search(r'Final Answer:\s*(.+)', text, re.DOTALL)
    if final_match:
        result.is_final = True
        result.final_answer = final_match.group(1).strip()
        return result

    # Extract Action — handle Chinese-descriptive Action lines like:
    #   Action: 使用 rag_search 搜索...
    #   Action: 调用 mcp__tavily-mcp__tavily_search 查询...
    action_match = re.search(r'Action:\s*(.+?)(?:\n|$)', text)
    if action_match:
        action_text = action_match.group(1).strip()
        result.action = _resolve_tool_name(action_text, tool_names or [])
        # Re-parse with stricter matching if failed
        if not result.action:
            action_match2 = re.search(r'Action:\s*(\S+)', text)
            if action_match2:
                result.action = _resolve_tool_name(action_match2.group(1).strip(), tool_names or [])

    # Extract Action Input (try JSON first, then key=value)
    action_input_match = re.search(r'Action Input:\s*(\{.+?\}|.+)', text, re.DOTALL)
    if action_input_match:
        input_str = action_input_match.group(1).strip()
        result.action_input = parse_action_input(input_str)

    return result


def _trim_repetition(text: str) -> str:
    """Detect and cut repetitive LLM output at the first repetition point.

    When a model loops — e.g. "I will output the answer... I will output the answer..."
    — truncate everything after the first occurrence of the repeated line.
    """
    lines = text.split("\n")
    seen: set[str] = set()
    clean_lines: list[str] = []

    for line in lines:
        stripped = line.strip()
        # Skip empty lines in repetition check
        if not stripped:
            clean_lines.append(line)
            continue
        # Normalize for comparison
        norm = stripped.lower().rstrip(".。！!?？,，")
        if norm in seen:
            # Found repetition — stop here
            break
        if len(norm) > 15:  # only track meaningful lines
            seen.add(norm)
        clean_lines.append(line)

    return "\n".join(clean_lines)


def _resolve_tool_name(action_text: str, tool_names: list[str]) -> str:
    """Extract the actual tool name from an Action line that may contain Chinese description.

    Example inputs → outputs:
        "rag_search" → "rag_search"
        "使用 rag_search 搜索" → "rag_search"
        "调用 mcp__tavily-mcp__tavily_search 查询" → "mcp__tavily-mcp__tavily_search"
        "搜索文档" → "" (no tool found)
    """
    # Already a clean single identifier
    if re.match(r'^[a-zA-Z_][a-zA-Z0-9_\-/]*$', action_text):
        if tool_names:
            if action_text in tool_names:
                return action_text
            # Fuzzy: unique tool ending with the model's text
            # ("tavily_search" → "mcp__tavily-mcp__tavily_search")
            matches = [t for t in tool_names if t.endswith(action_text)]
            if len(matches) == 1:
                return matches[0]
            return ""
        return action_text

    # Find tool-like patterns in the text: lowercase_with_underscores, possibly with __ or /
    candidates = re.findall(r'[a-zA-Z_][a-zA-Z0-9_\-/]{2,}', action_text)
    for c in candidates:
        # Must contain underscore (real tools look like rag_search, mcp__xxx__yyy)
        if '_' in c:
            if tool_names:
                if c in tool_names:
                    return c
            else:
                return c

    # Last resort: pick the first ascii word
    for c in candidates:
        if tool_names:
            if c in tool_names:
                return c
        else:
            return c

    # Fuzzy match: unique tool whose name contains the model's text.
    # Handles cases like "tavily_search" → "mcp__tavily-mcp__tavily_search".
    # The model's text must be a real suffix/substring, and only ONE tool may
    # match (otherwise it's ambiguous and we refuse to guess).
    if tool_names and len(action_text) >= 4 and '_' in action_text:
        matches = [t for t in tool_names if t.endswith(action_text) or t.split('__')[-1] == action_text]
        if len(matches) == 1:
            return matches[0]

    return ""


def parse_action_input(input_str: str) -> dict:
    """Parse action input string into a dict. Tries JSON first, then key=value."""
    # Try JSON
    try:
        return json.loads(input_str)
    except json.JSONDecodeError:
        pass

    # Try to extract JSON from within the string
    json_match = re.search(r'\{[^{}]*\}', input_str)
    if json_match:
        try:
            return json.loads(json_match.group(0))
        except json.JSONDecodeError:
            pass

    # Fallback: treat as raw string
    if input_str:
        return {"query": input_str}

    return {}


def extract_final_answer(text: str) -> Optional[str]:
    """Extract the final answer from text if present."""
    from agentic_rag.services.llm.base import strip_reasoning
    text = strip_reasoning(text)
    match = re.search(r'Final Answer:\s*(.+)', text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return None


def is_final_answer(text: str) -> bool:
    """Check if the text contains a Final Answer marker."""
    return "Final Answer:" in text


def format_observation(tool_name: str, result: str, error: Optional[str] = None) -> str:
    """Format a tool execution result as an Observation."""
    if error:
        return f"Observation: Error executing '{tool_name}': {error}"
    # Truncate very long results — keeps the ReAct prompt compact so the
    # model has less material to over-analyze on the next turn.
    if len(result) > 1500:
        result = result[:1500] + "... (truncated)"
    return f"Observation: {result}"