app/agent/parser.py

"""Robust JSON action parser for LLM outputs.

The LLM is instructed to respond with structured JSON in one of two forms:

    {"action": "tool_call", "reasoning": "...", "tool_name": "...", "tool_input": {...}}
    {"action": "final_answer", "reasoning": "...", "answer": "..."}

The 'reasoning' field is optional but logged for observability.
It is NEVER returned in the API response.

In practice, LLMs (especially smaller ones) produce:
- Markdown fences around JSON
- Preamble text before the JSON
- Trailing explanations after the JSON
- Slightly malformed JSON (trailing commas, single quotes)

This module handles all of that and returns typed dataclasses.
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


# ── Result types ───────────────────────────────────────────────────────────

@dataclass(frozen=True)
class ToolCallAction:
    """LLM decided to call a tool."""
    tool_name: str
    tool_input: dict[str, Any]
    reasoning: str = ""


@dataclass(frozen=True)
class FinalAnswerAction:
    """LLM decided to give the final answer."""
    answer: str
    reasoning: str = ""


@dataclass(frozen=True)
class ParseError:
    """Failed to parse the LLM output into a valid action."""
    raw_output: str
    error: str


# Union type for convenience
ActionResult = ToolCallAction | FinalAnswerAction | ParseError


# ── Parser ─────────────────────────────────────────────────────────────────

def parse_llm_action(raw: str) -> ActionResult:
    """Parse LLM raw text into a typed action.

    Extraction strategy:
    1. Strip markdown code fences if present
    2. Find the first top-level JSON object via brace matching
    3. Parse and validate required fields

    Args:
        raw: Raw string output from the LLM.

    Returns:
        ToolCallAction, FinalAnswerAction, or ParseError.
    """
    if not raw or not raw.strip():
        return ParseError(raw_output=raw or "", error="Empty LLM output")

    cleaned = _strip_markdown_fences(raw.strip())
    json_str = _extract_json_object(cleaned)

    if json_str is None:
        return ParseError(
            raw_output=raw,
            error="No JSON object found in LLM output",
        )

    try:
        data = json.loads(json_str)
    except json.JSONDecodeError as exc:
        return ParseError(
            raw_output=raw,
            error=f"Invalid JSON: {exc}",
        )

    if not isinstance(data, dict):
        return ParseError(
            raw_output=raw,
            error=f"Expected JSON object, got {type(data).__name__}",
        )

    return _validate_action(data, raw)


# ── Internals ──────────────────────────────────────────────────────────────

def _strip_markdown_fences(text: str) -> str:
    """Remove ```json ... ``` or ``` ... ``` wrappers."""
    # Match ```json\n...\n``` or ```\n...\n```
    pattern = r"^```(?:json)?\s*\n?(.*?)\n?\s*```$"
    match = re.match(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return text


def _extract_json_object(text: str) -> str | None:
    """Find the first balanced top-level { ... } in the text.

    Uses brace-depth counting to handle nested objects correctly,
    unlike naive regex that breaks on nested braces.
    """
    start = text.find("{")
    if start == -1:
        return None

    depth = 0
    in_string = False
    escape_next = False

    for i in range(start, len(text)):
        ch = text[i]

        if escape_next:
            escape_next = False
            continue

        if ch == "\\":
            escape_next = True
            continue

        if ch == '"':
            in_string = not in_string
            continue

        if in_string:
            continue

        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                return text[start : i + 1]

    return None  # Unbalanced braces


def _validate_action(data: dict[str, Any], raw: str) -> ActionResult:
    """Validate the parsed JSON against expected action schemas."""
    action = data.get("action")

    if action is None:
        return ParseError(
            raw_output=raw,
            error="Missing 'action' field in JSON",
        )

    if action == "tool_call":
        return _validate_tool_call(data, raw)

    if action == "final_answer":
        return _validate_final_answer(data, raw)

    return ParseError(
        raw_output=raw,
        error=f"Unknown action type: '{action}'. Must be 'tool_call' or 'final_answer'.",
    )


def _validate_tool_call(data: dict[str, Any], raw: str) -> ToolCallAction | ParseError:
    """Validate a tool_call action."""
    tool_name = data.get("tool_name")
    if not tool_name or not isinstance(tool_name, str):
        return ParseError(
            raw_output=raw,
            error="'tool_call' action requires a non-empty 'tool_name' string",
        )

    tool_input = data.get("tool_input")
    if tool_input is None:
        tool_input = {}
    if not isinstance(tool_input, dict):
        return ParseError(
            raw_output=raw,
            error=f"'tool_input' must be an object, got {type(tool_input).__name__}",
        )

    return ToolCallAction(tool_name=tool_name, tool_input=tool_input, reasoning=data.get("reasoning", ""))


def _stringify_answer(answer: Any) -> str:
    """Convert a non-string answer (list, dict) into clean bullet-point text.

    Handles common LLM patterns:
    - ["fact1", "fact2"]                    → "- fact1\\n- fact2"
    - {"facts": ["fact1", "fact2"]}         → "- fact1\\n- fact2"
    - {"key": "value", "key2": ["a", "b"]} → "- key: value\\n- key2:\\n  - a\\n  - b"
    """
    if isinstance(answer, list):
        items = [str(item).strip() for item in answer if item]
        return "\n".join(f"- {item}" for item in items) if items else str(answer)

    if isinstance(answer, dict):
        # Single-key dict wrapping a list → unwrap (e.g., {"facts": [...]})
        values = list(answer.values())
        if len(values) == 1 and isinstance(values[0], list):
            return _stringify_answer(values[0])
        # Multi-key dict → key-value bullets
        lines: list[str] = []
        for key, value in answer.items():
            if isinstance(value, list):
                lines.append(f"- {key}:")
                for item in value:
                    lines.append(f"  - {str(item).strip()}")
            else:
                lines.append(f"- {key}: {str(value).strip()}")
        return "\n".join(lines) if lines else str(answer)

    return str(answer)


def _validate_final_answer(data: dict[str, Any], raw: str) -> FinalAnswerAction | ParseError:
    """Validate a final_answer action."""
    answer = data.get("answer")
    if answer is None:
        return ParseError(
            raw_output=raw,
            error="'final_answer' action requires an 'answer' field",
        )

    if not isinstance(answer, str):
        # Be lenient — convert structured data to bullet-point text
        answer = _stringify_answer(answer)

    return FinalAnswerAction(answer=answer, reasoning=data.get("reasoning", ""))