tatlock/tests/responses/test_advanced_features.py

"""
Tests for Phase 6 advanced features.

Tests:
- Parameter validation (reasoning effort, max_output_tokens, stop sequences)
- Stop sequence detection and enforcement
- Max tokens enforcement
"""
import json
import pytest
from fastapi.testclient import TestClient
from httpx import AsyncClient
from pydantic import ValidationError

from src.responses.schemas import ResponseRequest
from src.responses.streaming import StreamingCoordinator


# ============================================================================
# Parameter Validation Tests
# ============================================================================

@pytest.mark.unit
def test_temperature_validation():
    """Test temperature parameter validation."""
    # Valid temperatures
    valid_temps = [0.0, 0.5, 1.0, 1.5, 2.0]
    for temp in valid_temps:
        request = ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            temperature=temp
        )
        assert request.temperature == temp

    # Invalid temperatures
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            temperature=-0.1  # Too low
        )
    assert "temperature" in str(exc_info.value).lower()

    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            temperature=2.1  # Too high
        )
    assert "temperature" in str(exc_info.value).lower()


@pytest.mark.unit
def test_reasoning_effort_validation():
    """Test reasoning.effort parameter validation."""
    # Valid effort levels
    valid_efforts = ['none', 'minimal', 'low', 'medium', 'high', 'xhigh']
    for effort in valid_efforts:
        request = ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            reasoning={"effort": effort, "summary": "auto"}
        )
        assert request.reasoning["effort"] == effort

    # Invalid effort level
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            reasoning={"effort": "invalid", "summary": "auto"}
        )
    assert "reasoning.effort" in str(exc_info.value)


@pytest.mark.unit
def test_reasoning_summary_validation():
    """Test reasoning.summary parameter validation."""
    # Valid summary values
    valid_summaries = ['auto', 'off']
    for summary in valid_summaries:
        request = ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            reasoning={"effort": "medium", "summary": summary}
        )
        assert request.reasoning["summary"] == summary

    # Invalid summary value
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            reasoning={"effort": "medium", "summary": "invalid"}
        )
    assert "reasoning.summary" in str(exc_info.value)


@pytest.mark.unit
def test_max_output_tokens_validation():
    """Test max_output_tokens parameter validation."""
    # Valid values
    request = ResponseRequest(
        model="lorem-tester",
        input=[{"role": "user", "content": "Hello"}],
        max_output_tokens=100
    )
    assert request.max_output_tokens == 100

    # None is valid (unlimited)
    request = ResponseRequest(
        model="lorem-tester",
        input=[{"role": "user", "content": "Hello"}],
        max_output_tokens=None
    )
    assert request.max_output_tokens is None

    # Invalid: zero or negative
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            max_output_tokens=0
        )
    assert "max_output_tokens" in str(exc_info.value)

    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            max_output_tokens=-10
        )
    assert "max_output_tokens" in str(exc_info.value)


@pytest.mark.unit
def test_stop_sequences_validation():
    """Test stop sequences parameter validation."""
    # Valid: up to 4 stop sequences
    for num_seqs in range(1, 5):
        stop_seqs = [f"stop{i}" for i in range(num_seqs)]
        request = ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            stop=stop_seqs
        )
        assert request.stop == stop_seqs

    # Invalid: more than 4 stop sequences
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            stop=["stop1", "stop2", "stop3", "stop4", "stop5"]  # 5 sequences
        )
    assert "4 stop sequences" in str(exc_info.value)

    # Invalid: empty string in stop sequences
    with pytest.raises(ValidationError) as exc_info:
        ResponseRequest(
            model="lorem-tester",
            input=[{"role": "user", "content": "Hello"}],
            stop=["stop1", ""]  # Empty string
        )
    assert "non-empty" in str(exc_info.value).lower()


# ============================================================================
# Stop Sequence Enforcement Tests
# ============================================================================

@pytest.mark.unit
def test_stop_sequence_detection_helper():
    """Test stop sequence detection helper method."""
    coordinator = StreamingCoordinator()

    # No stop sequences
    found, text = coordinator._check_stop_sequence("Hello world", None)
    assert found is False
    assert text == "Hello world"

    # Stop sequence not present
    found, text = coordinator._check_stop_sequence(
        "Hello world",
        ["STOP", "END"]
    )
    assert found is False
    assert text == "Hello world"

    # Stop sequence found
    found, text = coordinator._check_stop_sequence(
        "Hello STOP this should not appear",
        ["STOP"]
    )
    assert found is True
    assert text == "Hello "

    # Multiple stop sequences, first one wins
    found, text = coordinator._check_stop_sequence(
        "Hello STOP this END that",
        ["STOP", "END"]
    )
    assert found is True
    assert text == "Hello "


@pytest.mark.unit
@pytest.mark.asyncio
async def test_stop_sequence_in_streaming(async_client: AsyncClient):
    """Test stop sequence enforcement during streaming."""
    # We'll use lorem-tester which generates predictable text
    # The agent generates lorem ipsum text, so we use a stop sequence
    # that's likely to appear
    request_data = {
        "model": "lorem-tester",
        "input": [{"role": "user", "content": "Generate long text"}],
        "stop": ["dolor"],  # Common word in lorem ipsum
        "stream": True
    }

    chunks_received = []
    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=20.0,
    ) as response:
        assert response.status_code == 200

        async for line in response.aiter_lines():
            if not line.strip():
                continue

            if line.startswith("data: "):
                data_str = line[6:].strip()
                if data_str != "[DONE]":
                    try:
                        chunk = json.loads(data_str)
                        chunks_received.append(chunk)
                    except json.JSONDecodeError:
                        pass

    # Should have received chunks and stopped early
    # (Can't verify exact stop behavior with random lorem ipsum,
    # but test ensures no errors occur)
    assert len(chunks_received) > 0


# ============================================================================
# Max Tokens Enforcement Tests
# ============================================================================

@pytest.mark.unit
def test_max_tokens_check_helper():
    """Test max tokens check helper method."""
    coordinator = StreamingCoordinator()

    # No limit
    assert coordinator._check_max_tokens(100, None) is False

    # Under limit
    assert coordinator._check_max_tokens(50, 100) is False

    # At limit
    assert coordinator._check_max_tokens(100, 100) is True

    # Over limit
    assert coordinator._check_max_tokens(150, 100) is True


@pytest.mark.unit
def test_token_counting_helper():
    """Test approximate token counting."""
    coordinator = StreamingCoordinator()

    # Empty string
    assert coordinator._count_tokens_approx("") == 0

    # ~4 chars per token
    text = "Hello world"  # 11 chars
    tokens = coordinator._count_tokens_approx(text)
    assert tokens == 2  # 11 // 4 = 2

    text = "A" * 100  # 100 chars
    tokens = coordinator._count_tokens_approx(text)
    assert tokens == 25  # 100 // 4 = 25


@pytest.mark.unit
@pytest.mark.asyncio
async def test_max_tokens_in_streaming(async_client: AsyncClient):
    """Test max tokens enforcement during streaming."""
    # Set very low max_output_tokens to force early stop
    request_data = {
        "model": "lorem-tester",
        "input": [{"role": "user", "content": "Generate text"}],
        "max_output_tokens": 5,  # Very low limit
        "stream": True
    }

    chunks_received = []
    text_chunks = []

    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=20.0,
    ) as response:
        assert response.status_code == 200

        async for line in response.aiter_lines():
            if not line.strip():
                continue

            if line.startswith("event: "):
                event_type = line[7:].strip()
            elif line.startswith("data: "):
                data_str = line[6:].strip()
                if data_str != "[DONE]":
                    try:
                        chunk = json.loads(data_str)
                        chunks_received.append(chunk)

                        # Collect text deltas
                        if "delta" in chunk:
                            text_chunks.append(chunk["delta"])
                    except json.JSONDecodeError:
                        pass

    # Should have stopped early due to token limit
    total_text = "".join(text_chunks)
    # With max_output_tokens=5, total text should be small
    # (Approximate: 5 tokens * 4 chars ≈ 20 chars)
    assert len(total_text) < 100  # Reasonable upper bound


# ============================================================================
# Combined Features Test
# ============================================================================

@pytest.mark.unit
def test_combined_validation(client: TestClient):
    """Test combined parameter validation in actual request."""
    # Valid request with all advanced features
    request_data = {
        "model": "lorem-tester",
        "input": [{"role": "user", "content": "Hello"}],
        "temperature": 0.7,
        "max_output_tokens": 100,
        "stop": ["STOP", "END"],
        "reasoning": {"effort": "high", "summary": "auto"},
        "stream": False
    }

    response = client.post("/v1/responses", json=request_data)
    assert response.status_code == 200
    data = response.json()
    assert data["status"] == "completed"


@pytest.mark.unit
def test_invalid_combined_parameters(client: TestClient):
    """Test that invalid parameters are rejected."""
    # Invalid temperature
    request_data = {
        "model": "lorem-tester",
        "input": [{"role": "user", "content": "Hello"}],
        "temperature": 3.0,  # Too high
        "stream": False
    }

    response = client.post("/v1/responses", json=request_data)
    assert response.status_code == 422
    data = response.json()
    assert "error" in data
    assert data["error"]["type"] == "invalid_request_error"


# ============================================================================
# Streaming Delta Calculation Tests (No Duplication)
# ============================================================================

@pytest.mark.unit
@pytest.mark.asyncio
async def test_streaming_delta_calculation_no_duplication():
    """
    Test that StreamingCoordinator correctly calculates deltas when agent
    yields accumulated text multiple times (PydanticAI pattern).

    This test prevents the duplication bug where the same text was
    streamed multiple times because we weren't computing deltas correctly.
    """
    from src.agents.base import AgentInterface, OutputItem
    from typing import AsyncGenerator, Any

    # Create a mock agent that simulates PydanticAI's behavior
    # (yielding accumulated text, not deltas)
    class MockStreamingAgent(AgentInterface):
        async def generate_response(
            self,
            messages: list[dict],
            reasoning: dict | None = None,
            tools: list[dict] | None = None,
            temperature: float = 1.0,
            max_tokens: int | None = None,
            stop: list[str] | None = None,
            **kwargs: Any
        ) -> AsyncGenerator[OutputItem, None]:
            """
            Simulate PydanticAI streaming behavior:
            - Yields accumulated text, not deltas
            - Multiple yields with status="in_progress"
            - Final yield with status="completed"
            """
            msg_id = "msg_test_123"

            # Simulate incremental accumulation like PydanticAI does
            accumulated_texts = [
                "Hello",
                "Hello world",
                "Hello world how",
                "Hello world how are",
                "Hello world how are you",
            ]

            for text in accumulated_texts:
                yield OutputItem(
                    type="message",
                    id=msg_id,
                    role="assistant",
                    content=[{
                        "type": "output_text",
                        "text": text,
                        "annotations": []
                    }],
                    status="in_progress"
                )

            # Final message
            yield OutputItem(
                type="message",
                id=msg_id,
                role="assistant",
                content=[{
                    "type": "output_text",
                    "text": "Hello world how are you",
                    "annotations": []
                }],
                status="completed"
            )

        async def supports_tools(self) -> bool:
            return False

        async def supports_reasoning(self) -> bool:
            return False

        async def get_capabilities(self) -> dict:
            return {"streaming": True, "reasoning": False, "tools": False}

    # Register the mock agent
    import time
    from src.agents.registry import ModelRegistry
    ModelRegistry.MODELS["mock-streaming"] = {
        "agent_class": MockStreamingAgent,
        "description": "Mock streaming agent for testing",
        "created": int(time.time()),
        "owned_by": "test",
    }

    try:
        # Create a test request
        request = ResponseRequest(
            model="mock-streaming",
            input=[{"role": "user", "content": "Test"}],
            stream=True
        )

        # Stream the response
        coordinator = StreamingCoordinator()
        collected_deltas = []

        async for event in coordinator.stream_response(request):
            if event.event == "response.output_text.delta":
                collected_deltas.append(event.delta)

        # Reconstruct the full text from deltas
        full_text = "".join(collected_deltas)

        # Verify no duplication - the text should appear exactly once
        assert full_text.count("Hello") == 1, "Text 'Hello' should appear exactly once"
        assert full_text.count("world") == 1, "Text 'world' should appear exactly once"
        assert full_text.count("how") == 1, "Text 'how' should appear exactly once"
        assert full_text.count("are") == 1, "Text 'are' should appear exactly once"
        assert full_text.count("you") == 1, "Text 'you' should appear exactly once"

        # Verify the reconstructed text is correct (no trailing space with chunk streaming)
        expected_text = "Hello world how are you"
        assert full_text == expected_text, f"Expected '{expected_text}', got '{full_text}'"

    finally:
        # Clean up
        del ModelRegistry.MODELS["mock-streaming"]


@pytest.mark.unit
@pytest.mark.asyncio
async def test_streaming_with_multiple_message_items():
    """
    Test that coordinator handles multiple message OutputItems correctly,
    only streaming the delta between each one.
    """
    from src.agents.base import AgentInterface, OutputItem
    from typing import AsyncGenerator, Any

    class MockMultiMessageAgent(AgentInterface):
        async def generate_response(
            self,
            messages: list[dict],
            reasoning: dict | None = None,
            tools: list[dict] | None = None,
            temperature: float = 1.0,
            max_tokens: int | None = None,
            stop: list[str] | None = None,
            **kwargs: Any
        ) -> AsyncGenerator[OutputItem, None]:
            """Yield multiple in_progress messages with accumulated text."""
            # First chunk
            yield OutputItem(
                type="message",
                id="msg_1",
                role="assistant",
                content=[{"type": "output_text", "text": "The answer is", "annotations": []}],
                status="in_progress"
            )

            # Second chunk (more text accumulated)
            yield OutputItem(
                type="message",
                id="msg_1",
                role="assistant",
                content=[{"type": "output_text", "text": "The answer is 42", "annotations": []}],
                status="in_progress"
            )

            # Final chunk
            yield OutputItem(
                type="message",
                id="msg_1",
                role="assistant",
                content=[{"type": "output_text", "text": "The answer is 42", "annotations": []}],
                status="completed"
            )

        async def supports_tools(self) -> bool:
            return False

        async def supports_reasoning(self) -> bool:
            return False

        async def get_capabilities(self) -> dict:
            return {"streaming": True, "reasoning": False, "tools": False}

    # Register mock agent
    import time
    from src.agents.registry import ModelRegistry
    ModelRegistry.MODELS["mock-multi"] = {
        "agent_class": MockMultiMessageAgent,
        "description": "Mock multi-message agent for testing",
        "created": int(time.time()),
        "owned_by": "test",
    }

    try:
        request = ResponseRequest(
            model="mock-multi",
            input=[{"role": "user", "content": "What is the answer?"}],
            stream=True
        )

        coordinator = StreamingCoordinator()
        collected_deltas = []

        async for event in coordinator.stream_response(request):
            if event.event == "response.output_text.delta":
                collected_deltas.append(event.delta)

        full_text = "".join(collected_deltas)

        # Should only see "The answer is 42" once, not repeated
        assert "The answer is 42" in full_text
        # Count occurrences - should only appear once
        assert full_text.count("The") == 1
        assert full_text.count("answer") == 1
        assert full_text.count("42") == 1

    finally:
        # Clean up
        del ModelRegistry.MODELS["mock-multi"]