llmproxy/tests/test_response_parser.py

"""
Unit tests for the Response Parser module.

Tests cover:
- Parsing text-only responses
- Parsing responses with tool calls
- Parsing native OpenAI-format tool calls
- Parsing streaming chunks
- Error handling and edge cases
"""

import pytest
import json
from app.response_parser import (
    ResponseParser,
    ToolCallParseError,
    parse_response,
    parse_response_with_custom_tags,
    TOOL_CALL_START_TAG,
    TOOL_CALL_END_TAG
)
from app.models import ToolCall, ToolCallFunction


class TestResponseParser:
    """Test suite for ResponseParser class."""

    def test_parse_text_only_response(self):
        """Test parsing a response with no tool calls."""
        parser = ResponseParser()
        text = "Hello, this is a simple response."
        result = parser.parse(text)

        assert result.content == text
        assert result.tool_calls is None

    def test_parse_empty_response(self):
        """Test parsing an empty response."""
        parser = ResponseParser()
        result = parser.parse("")

        assert result.content is None
        assert result.tool_calls is None

    def test_parse_response_with_tool_call(self):
        """Test parsing a response with a single tool call."""
        parser = ResponseParser()
        text = f'''I'll check the weather for you.
{TOOL_CALL_START_TAG}
{{
  "name": "get_weather",
  "arguments": {{
    "location": "San Francisco",
    "units": "celsius"
  }}
}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        assert result.content == "I'll check the weather for you."
        assert result.tool_calls is not None
        assert len(result.tool_calls) == 1

        tool_call = result.tool_calls[0]
        assert tool_call.type == "function"
        assert tool_call.function.name == "get_weather"

        arguments = json.loads(tool_call.function.arguments)
        assert arguments["location"] == "San Francisco"
        assert arguments["units"] == "celsius"

    def test_parse_response_with_tool_call_no_content(self):
        """Test parsing a response with only a tool call."""
        parser = ResponseParser()
        text = f'''{TOOL_CALL_START_TAG}
{{
  "name": "shell",
  "arguments": {{
    "command": ["ls", "-l"]
  }}
}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        assert result.content is None
        assert result.tool_calls is not None
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "shell"

    def test_parse_response_with_malformed_tool_call(self):
        """Test parsing a response with malformed JSON in tool call."""
        parser = ResponseParser()
        text = f'''Here's the result.
{TOOL_CALL_START_TAG}
{{invalid json}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        # Should fall back to treating it as text
        assert result.content == text
        assert result.tool_calls is None

    def test_parse_response_with_missing_tool_name(self):
        """Test parsing a tool call without a name field."""
        parser = ResponseParser()
        text = f'''{TOOL_CALL_START_TAG}
{{
  "arguments": {{
    "command": "echo hello"
  }}
}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        # Should handle gracefully - when name is missing, ToolCallParseError is raised
        # and caught, falling back to treating as text content
        # content will be the text between start and end tags (the JSON object)
        assert result.content is not None

    def test_parse_response_with_complex_arguments(self):
        """Test parsing a tool call with complex nested arguments."""
        parser = ResponseParser()
        text = f'''Executing command.
{TOOL_CALL_START_TAG}
{{
  "name": "shell",
  "arguments": {{
    "command": ["bash", "-lc", "echo 'hello world' && ls -la"],
    "timeout": 5000,
    "env": {{
      "PATH": "/usr/bin"
    }}
  }}
}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        assert result.content == "Executing command."
        assert result.tool_calls is not None

        arguments = json.loads(result.tool_calls[0].function.arguments)
        assert arguments["command"] == ["bash", "-lc", "echo 'hello world' && ls -la"]
        assert arguments["timeout"] == 5000
        assert arguments["env"]["PATH"] == "/usr/bin"

    def test_parse_with_custom_tags(self):
        """Test parsing with custom start and end tags."""
        parser = ResponseParser(
            tool_call_start_tag="<TOOL_CALL>",
            tool_call_end_tag="</TOOL_CALL>"
        )
        text = """I'll help you with that.
<TOOL_CALL>
{
  "name": "search",
  "arguments": {
    "query": "python tutorials"
  }
}
</TOOL_CALL>
"""

        result = parser.parse(text)

        assert "I'll help you with that" in result.content
        assert result.tool_calls is not None
        assert result.tool_calls[0].function.name == "search"

    def test_parse_streaming_chunks(self):
        """Test parsing aggregated streaming chunks."""
        parser = ResponseParser()
        chunks = [
            "I'll run that ",
            "command for you.",
            f'{TOOL_CALL_START_TAG}\n{{"name": "shell", "arguments": {{"command": ["echo", "hello"]}}}}\n{TOOL_CALL_END_TAG}'
        ]

        result = parser.parse_streaming_chunks(chunks)

        assert "I'll run that command for you" in result.content
        assert result.tool_calls is not None
        assert result.tool_calls[0].function.name == "shell"

    def test_parse_native_tool_calls(self):
        """Test parsing a native OpenAI-format response with tool calls."""
        parser = ResponseParser()
        llm_response = {
            "role": "assistant",
            "content": "I'll execute that command.",
            "tool_calls": [
                {
                    "id": "call_abc123",
                    "type": "function",
                    "function": {
                        "name": "shell",
                        "arguments": '{"command": ["ls", "-l"]}'
                    }
                }
            ]
        }

        result = parser.parse_native_tool_calls(llm_response)

        assert result.content == "I'll execute that command."
        assert result.tool_calls is not None
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].id == "call_abc123"
        assert result.tool_calls[0].function.name == "shell"

    def test_parse_native_tool_calls_multiple(self):
        """Test parsing a response with multiple native tool calls."""
        parser = ResponseParser()
        llm_response = {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_1",
                    "type": "function",
                    "function": {
                        "name": "shell",
                        "arguments": '{"command": ["pwd"]}'
                    }
                },
                {
                    "id": "call_2",
                    "type": "function",
                    "function": {
                        "name": "shell",
                        "arguments": '{"command": ["ls", "-la"]}'
                    }
                }
            ]
        }

        result = parser.parse_native_tool_calls(llm_response)

        assert result.tool_calls is not None
        assert len(result.tool_calls) == 2
        assert result.tool_calls[0].id == "call_1"
        assert result.tool_calls[1].id == "call_2"

    def test_parse_native_tool_calls_falls_back_to_text(self):
        """Test that native parsing falls back to text parsing when no tool_calls."""
        parser = ResponseParser()
        llm_response = {
            "role": "assistant",
            "content": "This is a simple text response."
        }

        result = parser.parse_native_tool_calls(llm_response)

        assert result.content == "This is a simple text response."
        assert result.tool_calls is None

    def test_generate_unique_tool_call_ids(self):
        """Test that tool call IDs are unique."""
        parser = ResponseParser()

        text1 = f'{TOOL_CALL_START_TAG}{{"name": "tool1", "arguments": {{}}}}{TOOL_CALL_END_TAG}'
        text2 = f'{TOOL_CALL_START_TAG}{{"name": "tool2", "arguments": {{}}}}{TOOL_CALL_END_TAG}'

        result1 = parser.parse(text1)
        result2 = parser.parse(text2)

        id1 = result1.tool_calls[0].id
        id2 = result2.tool_calls[0].id

        assert id1 != id2
        assert id1.startswith("call_tool1_")
        assert id2.startswith("call_tool2_")


class TestConvenienceFunctions:
    """Test suite for convenience functions."""

    def test_parse_response_default_parser(self):
        """Test the parse_response convenience function."""
        text = f'{TOOL_CALL_START_TAG}{{"name": "search", "arguments": {{"query": "test"}}}}{TOOL_CALL_END_TAG}'
        result = parse_response(text)

        assert result.tool_calls is not None
        assert result.tool_calls[0].function.name == "search"

    def test_parse_response_with_custom_tags_function(self):
        """Test the parse_response_with_custom_tags function."""
        text = """[CALL]
{"name": "test", "arguments": {}}
[/CALL]"""
        result = parse_response_with_custom_tags(
            text,
            start_tag="[CALL]",
            end_tag="[/CALL]"
        )

        assert result.tool_calls is not None
        assert result.tool_calls[0].function.name == "test"


class TestEdgeCases:
    """Test edge cases and error conditions."""

    def test_response_with_whitespace(self):
        """Test parsing responses with various whitespace patterns."""
        parser = ResponseParser()

        # Leading/trailing whitespace
        text = "   Hello world.   "
        result = parser.parse(text)
        assert result.content.strip() == "Hello world."

    def test_response_with_newlines_only(self):
        """Test parsing a response with only newlines."""
        parser = ResponseParser()
        result = parser.parse("\n\n\n")

        assert result.content == ""
        assert result.tool_calls is None

    def test_response_with_special_characters(self):
        """Test parsing responses with special characters in content."""
        parser = ResponseParser()
        special_chars = '@#$%^&*()'
        text = f'''Here's the result with special chars: {special_chars}
{TOOL_CALL_START_TAG}
{{
  "name": "test",
  "arguments": {{
    "special": "!@#$%"
  }}
}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)
        assert "@" in result.content
        assert result.tool_calls is not None

    def test_response_with_escaped_quotes(self):
        """Test parsing tool calls with escaped quotes in arguments."""
        parser = ResponseParser()
        text = f'{TOOL_CALL_START_TAG}{{"name": "echo", "arguments": {{"message": "Hello \\"world\\""}}}}{TOOL_CALL_END_TAG}'

        result = parser.parse(text)
        arguments = json.loads(result.tool_calls[0].function.arguments)
        assert arguments["message"] == 'Hello "world"'

    def test_multiple_tool_calls_in_text_finds_first(self):
        """Test that only the first tool call is extracted."""
        parser = ResponseParser()
        text = f'''First call.
{TOOL_CALL_START_TAG}
{{"name": "tool1", "arguments": {{}}}}
{TOOL_CALL_END_TAG}
Some text in between.
{TOOL_CALL_START_TAG}
{{"name": "tool2", "arguments": {{}}}}
{TOOL_CALL_END_TAG}
'''

        result = parser.parse(text)

        # Should only find the first one
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0].function.name == "tool1"