main.py

"""
Parse tool_call JSON blocks embedded in LLM text responses.

Handles common LLM quirks:
- Tool call JSON mixed with prose / markdown
- Stringified JSON inside arguments (with literal newlines, tabs, etc.)
- Multiple tool calls in one response
"""

from __future__ import annotations

import json
from dataclasses import dataclass


@dataclass
class ToolCall:
    """A parsed tool call."""
    name: str
    arguments: dict

    def __repr__(self) -> str:
        return f"ToolCall(name={self.name!r}, arguments={self.arguments!r})"


def parse_tool_call(text: str) -> ToolCall | None:
    """
    Extract the first ``{"tool_call": ...}`` block from *text* and return
    a :class:`ToolCall`, or ``None`` if nothing is found.
    """
    results = parse_all_tool_calls(text)
    return results[0] if results else None


def parse_all_tool_calls(text: str) -> list[ToolCall]:
    """
    Extract every ``{"tool_call": ...}`` block from *text*.
    """
    results: list[ToolCall] = []
    search_start = 0

    while True:
        idx = text.find('{"tool_call"', search_start)
        if idx == -1:
            break

        raw = _extract_balanced_json(text, idx)
        if raw is None:
            search_start = idx + 1
            continue

        obj, end_pos = raw
        tc = _validate_tool_call(obj)
        if tc is not None:
            results.append(tc)

        search_start = end_pos

    return results


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _extract_balanced_json(text: str, start: int) -> tuple[dict, int] | None:
    """
    Starting at *start* (which must point at ``{``), walk the string
    respecting JSON string escaping and return ``(parsed_dict, end_index)``
    or ``None`` if the braces never balance or the result isn't valid JSON.
    """
    depth = 0
    in_string = False
    escape = False

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

        if escape:
            escape = False
            continue

        if in_string:
            if ch == '\\':
                escape = True
            elif ch == '"':
                in_string = False
            continue

        # Outside a string
        if ch == '"':
            in_string = True
        elif ch == '{':
            depth += 1
        elif ch == '}':
            depth -= 1
            if depth == 0:
                candidate = text[start : i + 1]
                try:
                    parsed = json.loads(candidate, strict=False)
                    if isinstance(parsed, dict):
                        return parsed, i + 1
                except json.JSONDecodeError:
                    return None

    return None


def _validate_tool_call(obj: dict) -> ToolCall | None:
    """
    Validate shape and return a :class:`ToolCall` or ``None``.

    Expected shape::

        {
          "tool_call": {
            "name": "<string>",
            "arguments": { ... }
          }
        }
    """
    tc = obj.get("tool_call")
    if not isinstance(tc, dict):
        return None

    name = tc.get("name")
    if not isinstance(name, str) or not name:
        return None

    arguments = tc.get("arguments")
    if not isinstance(arguments, dict):
        return None

    return ToolCall(name=name, arguments=arguments)



"""Tests for tool_call_parser."""

import json
import pytest



# ── The original example from the user ────────────────────────────────────

ORIGINAL_TEXT = r'''stripped:  Em chào anh. Để nhóm người dùng có hoạt động trong vòng 60 ngày gần nhất, mình sẽ dùng điều kiện `latest_active`. Mình sẽ dùng toán tử `>` (lớn hơn) để lọc ra những người dùng có `latest_active` lớn hơn ngày hôm nay trừ đi 60 ngày.
Đầu tiên, hôm nay là 2026-06-22. Vậy mình cần so sánh với ngày 2026-04-22.
Mình sẽ tạo JSON như sau:
```json
{
  "condition": "and",
  "rules": [
    {
      "type": "date",
      "field": "latest_active",
      "label": "latest_active",
      "operator": ">",
      "value": "2026-04-22"
    }
  ],
  "select": ["user_id"]
}
```
Bây giờ em sẽ kiểm tra xem JSON này có hợp lệ không nhé.
{"tool_call": {"name": "validate_query", "arguments": {"query_json": "{\n  \"condition\": \"and\",\n  \"rules\": [\n    {\n      \"type\": \"date\",\n      \"field\": \"latest_active\",\n      \"label\": \"latest_active\",\n      \"operator\": \">\",\n      \"value\": \"2026-04-22\"\n    }\n  ],\n  \"select\": [\"user_id\"]\n}"}}}'''


class TestOriginalExample:
    def test_parses_successfully(self):
        result = parse_tool_call(ORIGINAL_TEXT)
        assert result is not None

    def test_correct_name(self):
        result = parse_tool_call(ORIGINAL_TEXT)
        assert result.name == "validate_query"

    def test_arguments_has_query_json(self):
        result = parse_tool_call(ORIGINAL_TEXT)
        assert "query_json" in result.arguments

    def test_query_json_is_valid_json_string(self):
        result = parse_tool_call(ORIGINAL_TEXT)
        inner = json.loads(result.arguments["query_json"])
        assert inner["condition"] == "and"
        assert inner["rules"][0]["field"] == "latest_active"
        assert inner["rules"][0]["value"] == "2026-04-22"
        assert inner["select"] == ["user_id"]


# ── Basic parsing ─────────────────────────────────────────────────────────

class TestBasicParsing:
    def test_standalone_tool_call(self):
        text = '{"tool_call": {"name": "greet", "arguments": {"msg": "hi"}}}'
        result = parse_tool_call(text)
        assert result == ToolCall(name="greet", arguments={"msg": "hi"})

    def test_tool_call_after_prose(self):
        text = 'Here is my answer.\n{"tool_call": {"name": "run", "arguments": {"x": 1}}}'
        result = parse_tool_call(text)
        assert result.name == "run"
        assert result.arguments == {"x": 1}

    def test_tool_call_with_trailing_text(self):
        text = '{"tool_call": {"name": "a", "arguments": {}}}\nDone!'
        result = parse_tool_call(text)
        assert result.name == "a"

    def test_empty_arguments(self):
        text = '{"tool_call": {"name": "noop", "arguments": {}}}'
        result = parse_tool_call(text)
        assert result.arguments == {}


# ── Edge cases ────────────────────────────────────────────────────────────

class TestEdgeCases:
    def test_no_tool_call_returns_none(self):
        assert parse_tool_call("Just a regular message") is None

    def test_empty_string(self):
        assert parse_tool_call("") is None

    def test_malformed_json_returns_none(self):
        text = '{"tool_call": {"name": "bad", "arguments": {BROKEN}'
        assert parse_tool_call(text) is None

    def test_missing_name_returns_none(self):
        text = '{"tool_call": {"arguments": {"x": 1}}}'
        assert parse_tool_call(text) is None

    def test_empty_name_returns_none(self):
        text = '{"tool_call": {"name": "", "arguments": {"x": 1}}}'
        assert parse_tool_call(text) is None

    def test_missing_arguments_returns_none(self):
        text = '{"tool_call": {"name": "foo"}}'
        assert parse_tool_call(text) is None

    def test_arguments_not_dict_returns_none(self):
        text = '{"tool_call": {"name": "foo", "arguments": "string"}}'
        assert parse_tool_call(text) is None

    def test_tool_call_not_dict_returns_none(self):
        text = '{"tool_call": "not a dict"}'
        assert parse_tool_call(text) is None

    def test_nested_braces_in_prose_before_tool_call(self):
        """Prose containing { } shouldn't confuse the parser."""
        text = 'Use {"key": "val"} format.\n{"tool_call": {"name": "x", "arguments": {"a": 1}}}'
        result = parse_tool_call(text)
        assert result.name == "x"


# ── Stringified / nested JSON in arguments ────────────────────────────────

class TestStringifiedArguments:
    def test_literal_newlines_in_string_value(self):
        """The key scenario: arguments contain a JSON string with \\n chars."""
        inner_json = '{\n  "hello": "world"\n}'
        tool = {
            "tool_call": {
                "name": "check",
                "arguments": {"payload": inner_json},
            }
        }
        text = f"Some text.\n{json.dumps(tool)}"
        result = parse_tool_call(text)
        assert result.name == "check"
        parsed_inner = json.loads(result.arguments["payload"])
        assert parsed_inner["hello"] == "world"

    def test_tabs_in_string_value(self):
        tool = {
            "tool_call": {
                "name": "fmt",
                "arguments": {"code": "def f():\n\treturn 1"},
            }
        }
        text = json.dumps(tool)
        result = parse_tool_call(text)
        assert result.name == "fmt"
        assert "\t" in result.arguments["code"]

    def test_escaped_quotes_in_string_value(self):
        text = r'{"tool_call": {"name": "q", "arguments": {"s": "he said \"hi\""}}}'
        result = parse_tool_call(text)
        assert result.arguments["s"] == 'he said "hi"'

    def test_deeply_nested_json_string(self):
        """JSON string inside JSON string inside arguments."""
        inner = json.dumps({"a": 1})
        outer = json.dumps({"inner": inner})
        tool_text = json.dumps({
            "tool_call": {
                "name": "deep",
                "arguments": {"data": outer},
            }
        })
        result = parse_tool_call(tool_text)
        outer_parsed = json.loads(result.arguments["data"])
        inner_parsed = json.loads(outer_parsed["inner"])
        assert inner_parsed["a"] == 1


# ── Multiple tool calls ──────────────────────────────────────────────────

class TestMultipleToolCalls:
    def test_two_tool_calls(self):
        text = (
            'Step 1:\n'
            '{"tool_call": {"name": "first", "arguments": {"n": 1}}}\n'
            'Step 2:\n'
            '{"tool_call": {"name": "second", "arguments": {"n": 2}}}'
        )
        results = parse_all_tool_calls(text)
        assert len(results) == 2
        assert results[0].name == "first"
        assert results[1].name == "second"

    def test_parse_tool_call_returns_first(self):
        text = (
            '{"tool_call": {"name": "a", "arguments": {}}}\n'
            '{"tool_call": {"name": "b", "arguments": {}}}'
        )
        result = parse_tool_call(text)
        assert result.name == "a"

    def test_mixed_valid_and_invalid(self):
        text = (
            '{"tool_call": {"name": "good", "arguments": {"x": 1}}}\n'
            '{"tool_call": {"name": "", "arguments": {}}}\n'  # invalid (empty name)
            '{"tool_call": {"name": "also_good", "arguments": {"y": 2}}}'
        )
        results = parse_all_tool_calls(text)
        names = [r.name for r in results]
        assert names == ["good", "also_good"]


# ── ToolCall dataclass ────────────────────────────────────────────────────

class TestToolCallDataclass:
    def test_equality(self):
        a = ToolCall(name="f", arguments={"x": 1})
        b = ToolCall(name="f", arguments={"x": 1})
        assert a == b

    def test_inequality(self):
        a = ToolCall(name="f", arguments={"x": 1})
        b = ToolCall(name="g", arguments={"x": 1})
        assert a != b

    def test_repr(self):
        tc = ToolCall(name="f", arguments={})
        assert "f" in repr(tc)


# ── Unicode / i18n ────────────────────────────────────────────────────────

class TestUnicode:
    def test_vietnamese_prose_with_tool_call(self):
        text = 'Xin chào! Đây là câu trả lời.\n{"tool_call": {"name": "vi_tool", "arguments": {"text": "Việt Nam"}}}'
        result = parse_tool_call(text)
        assert result.name == "vi_tool"
        assert result.arguments["text"] == "Việt Nam"

    def test_unicode_in_arguments(self):
        text = '{"tool_call": {"name": "emoji", "arguments": {"mood": "😀🎉"}}}'
        result = parse_tool_call(text)
        assert result.arguments["mood"] == "😀🎉"