tatlock/tests/integration/test_tatlock_streaming.py

"""
Integration tests for Tatlock agent streaming through full API stack.

These tests verify the complete streaming flow from API endpoint through
StreamingCoordinator to TatlockAgent, ensuring no text duplication and
proper delta calculation.
"""
import json
import pytest
from httpx import AsyncClient
from fastapi.testclient import TestClient


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_streaming_no_duplication(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock streaming produces no text duplication.

    This test catches the bug where accumulated text from PydanticAI was
    being re-streamed multiple times by the StreamingCoordinator.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Say hello"}],
        "stream": True
    }

    collected_deltas = []

    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=30.0,  # Give enough time for Ollama response
    ) as response:
        assert response.status_code == 200
        assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

        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)

                        # Collect output text deltas
                        if chunk.get("event") == "response.output_text.delta":
                            collected_deltas.append(chunk["delta"])

                    except json.JSONDecodeError:
                        pass

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

    # Verify we got some response
    assert len(full_text) > 0, "Should have received some text"

    # Verify no obvious duplication patterns
    # Check that common words don't appear excessively repeated
    words = full_text.lower().split()
    if len(words) > 0:
        # Check for consecutive duplicate words (sign of duplication bug)
        consecutive_dupes = sum(
            1 for i in range(len(words) - 1)
            if words[i] == words[i + 1] and len(words[i]) > 3
        )
        # Allow a few duplicates (natural language), but not excessive
        assert consecutive_dupes < len(words) * 0.1, \
            f"Too many consecutive duplicate words: {consecutive_dupes}/{len(words)}"


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_chat_streaming_no_duplication(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock streaming through Chat Completions API.

    Tests the full stack through the chat completions wrapper to ensure
    streaming works correctly without duplication.
    """
    request_data = {
        "model": "Tatlock",
        "messages": [{"role": "user", "content": "Hello"}],
        "stream": True
    }

    collected_content = []

    async with async_client.stream(
        "POST",
        "/v1/chat/completions",
        json=request_data,
        timeout=30.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]":
                    break

                try:
                    chunk = json.loads(data_str)

                    # Collect content deltas from choices
                    if "choices" in chunk and len(chunk["choices"]) > 0:
                        delta = chunk["choices"][0].get("delta", {})
                        if "content" in delta and delta["content"]:
                            collected_content.append(delta["content"])

                except json.JSONDecodeError:
                    pass

    # Reconstruct full response
    full_response = "".join(collected_content)

    # Verify we got a response
    assert len(full_response) > 0, "Should have received response content"

    # Check for duplication patterns
    words = full_response.lower().split()
    if len(words) > 0:
        consecutive_dupes = sum(
            1 for i in range(len(words) - 1)
            if words[i] == words[i + 1] and len(words[i]) > 3
        )
        assert consecutive_dupes < len(words) * 0.1, \
            f"Too many consecutive duplicate words in chat response: {consecutive_dupes}/{len(words)}"


@pytest.mark.integration
def test_tatlock_non_streaming_responses_api(client: TestClient):
    """
    Integration test: Verify Tatlock non-streaming through Responses API.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Say hello"}],
        "stream": False
    }

    response = client.post("/v1/responses", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Verify response structure
    assert data["status"] == "completed"
    assert "output" in data
    assert len(data["output"]) > 0

    # Get the message content
    message_item = next((item for item in data["output"] if item["type"] == "message"), None)
    assert message_item is not None, "Should have a message output item"
    assert len(message_item["content"]) > 0

    text = message_item["content"][0]["text"]
    assert len(text) > 0, "Should have response text"


@pytest.mark.integration
def test_tatlock_non_streaming_chat_api(client: TestClient):
    """
    Integration test: Verify Tatlock non-streaming through Chat Completions API.
    """
    request_data = {
        "model": "Tatlock",
        "messages": [{"role": "user", "content": "Hello"}],
        "stream": False
    }

    response = client.post("/v1/chat/completions", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Verify OpenAI-compatible structure
    assert "id" in data
    assert data["object"] == "chat.completion"
    assert "choices" in data
    assert len(data["choices"]) > 0

    # Verify content
    choice = data["choices"][0]
    assert choice["message"]["role"] == "assistant"
    assert len(choice["message"]["content"]) > 0


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_streaming_delta_accumulation(async_client: AsyncClient):
    """
    Integration test: Verify deltas accumulate correctly without duplication.

    This test explicitly checks that when we accumulate all deltas,
    we get a coherent response without repeated text.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Count to three"}],
        "stream": True
    }

    collected_deltas = []
    previous_full_text = ""

    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=30.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)

                        if chunk.get("event") == "response.output_text.delta":
                            delta = chunk["delta"]
                            collected_deltas.append(delta)

                            # Verify each delta is new content
                            current_full = "".join(collected_deltas)
                            assert current_full.startswith(previous_full_text), \
                                "Deltas should accumulate progressively"
                            previous_full_text = current_full

                    except json.JSONDecodeError:
                        pass

    full_text = "".join(collected_deltas)
    assert len(full_text) > 0


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_with_reasoning(async_client: AsyncClient):
    """
    Integration test: Verify Tatlock with reasoning enabled.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Hello"}],
        "reasoning": {"effort": "medium", "summary": "auto"},
        "stream": True
    }

    has_reasoning = False
    has_output = False

    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=30.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)

                        if chunk.get("event") == "response.reasoning_summary_text.delta":
                            has_reasoning = True
                        elif chunk.get("event") == "response.output_text.delta":
                            has_output = True

                    except json.JSONDecodeError:
                        pass

    assert has_reasoning, "Should have reasoning summary"
    assert has_output, "Should have output text"


@pytest.mark.integration
@pytest.mark.asyncio
async def test_tatlock_markdown_formatting_preserved(async_client: AsyncClient):
    """
    Integration test: Verify markdown formatting is preserved in responses.

    Tests that code blocks, newlines, and other markdown formatting
    are properly preserved through the streaming pipeline.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Can you give me an HTML5 boilerplate template?"}],
        "stream": True
    }

    collected_deltas = []

    async with async_client.stream(
        "POST",
        "/v1/responses",
        json=request_data,
        timeout=45.0,  # Give extra time for code generation
    ) 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)

                        if chunk.get("event") == "response.output_text.delta":
                            collected_deltas.append(chunk["delta"])

                    except json.JSONDecodeError:
                        pass

    # Reconstruct full response
    full_response = "".join(collected_deltas)

    # Always print the response for debugging
    print("\n" + "="*80)
    print("FULL RESPONSE (repr):")
    print("="*80)
    print(repr(full_response))
    print("\n" + "="*80)
    print("FULL RESPONSE (formatted):")
    print("="*80)
    print(full_response)
    print("="*80 + "\n")

    # Verify we got a response
    assert len(full_response) > 100, "Should have a substantial response"

    # Verify markdown code block is present
    assert "```" in full_response, "Response should contain markdown code blocks"

    # Verify newlines are preserved (not all collapsed to spaces)
    newline_count = full_response.count('\n')
    assert newline_count > 5, f"Should have multiple newlines preserved, got {newline_count}"

    # Verify code block markers are complete
    code_block_starts = full_response.count("```")
    # Should have at least opening and closing markers (even count)
    assert code_block_starts % 2 == 0, "Code blocks should have matching opening/closing markers"
    assert code_block_starts >= 2, "Should have at least one complete code block"

    # Verify HTML tags are present (indicates code block content is preserved)
    assert "<!DOCTYPE html>" in full_response or "<html" in full_response, \
        "Should contain HTML5 boilerplate elements"

    # Verify indentation is preserved (check for multiple spaces in a row)
    # This indicates that code formatting with indentation is maintained
    assert "  " in full_response, "Should preserve indentation (multiple spaces)"

    # Log the response for debugging if test fails
    if "```" not in full_response or newline_count < 5:
        print("\n=== Full Response ===")
        print(repr(full_response))  # Use repr to see escaped characters
        print("\n=== Newline count ===")
        print(f"Found {newline_count} newlines")


@pytest.mark.integration
def test_tatlock_markdown_non_streaming(client: TestClient):
    """
    Integration test: Verify markdown in non-streaming mode.
    """
    request_data = {
        "model": "Tatlock",
        "input": [{"role": "user", "content": "Give me a simple Python hello world code"}],
        "stream": False
    }

    response = client.post("/v1/responses", json=request_data, timeout=30.0)

    assert response.status_code == 200
    data = response.json()

    # Get the message content
    message_item = next((item for item in data["output"] if item["type"] == "message"), None)
    assert message_item is not None

    text = message_item["content"][0]["text"]

    # Verify markdown code block
    assert "```" in text, "Should contain code block markers"
    assert "\n" in text, "Should contain newlines"