File size: 5,965 Bytes
9f2df60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a7bed5e
9f2df60
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
"""
Offline unit tests for the streaming response-field parser.
No API key, no network. Run: pytest tests/test_stream_parser.py -v
"""
import json
import os
import random
import sys

import pytest

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from src.rag.stream_parser import ResponseFieldStreamParser


def stream_through(parser: ResponseFieldStreamParser, text: str, chunk_size: int) -> str:
    """Feed `text` in chunks of `chunk_size`, return concatenated deltas."""
    out = []
    for i in range(0, len(text), chunk_size):
        out.append(parser.feed(text[i:i + chunk_size]))
    return "".join(out)


SAMPLE_JSON = json.dumps({
    "response": "Der **EMBA HSG** kostet CHF 77'500.\nFrühere Bewerbung = reduzierte Gebühr.",
    "additional_details": "Mehr Infos zur Finanzierung auf Anfrage.",
    "appointment_requested": False,
    "show_booking_widget": False,
    "relevant_programs": [],
}, ensure_ascii=False)

EXPECTED_RESPONSE = (
    "Der **EMBA HSG** kostet CHF 77'500.\nFrühere Bewerbung = reduzierte Gebühr."
)


class TestJsonMode:
    @pytest.mark.parametrize("chunk_size", [1, 2, 3, 5, 7, 16, 64, 10_000])
    def test_extracts_only_response_field_at_any_chunking(self, chunk_size):
        parser = ResponseFieldStreamParser()
        result = stream_through(parser, SAMPLE_JSON, chunk_size)
        assert result == EXPECTED_RESPONSE
        assert parser.field_complete

    def test_never_leaks_other_fields(self):
        parser = ResponseFieldStreamParser()
        result = stream_through(parser, SAMPLE_JSON, 3)
        assert "additional_details" not in result
        assert "appointment_requested" not in result
        assert "{" not in result

    def test_escaped_quotes_inside_value(self):
        payload = json.dumps({"response": 'Das Programm "emba X" ist ein Joint Degree.'})
        parser = ResponseFieldStreamParser()
        assert stream_through(parser, payload, 2) == 'Das Programm "emba X" ist ein Joint Degree.'

    def test_unicode_escapes(self):
        payload = '{"response": "Studiengeb\\u00fchr: CHF 85\'000"}'
        parser = ResponseFieldStreamParser()
        assert stream_through(parser, payload, 3) == "Studiengebühr: CHF 85'000"

    def test_response_field_not_first(self):
        payload = json.dumps({"additional_details": "Interne Details.", "response": "Antwort hier."})
        parser = ResponseFieldStreamParser()
        assert stream_through(parser, payload, 4) == "Antwort hier."

    def test_ignores_response_text_inside_earlier_string_values(self):
        payload = json.dumps({
            "query": 'Use "response": "internal query text" for the database lookup',
            "response": "User-visible answer.",
        })
        parser = ResponseFieldStreamParser()
        result = stream_through(parser, payload, 2)
        assert result == "User-visible answer."
        assert "internal query text" not in result

    def test_ignores_nested_response_fields(self):
        payload = json.dumps({
            "tool_call": {
                "name": "retrieve_context",
                "args": {"response": "internal nested text"},
            },
            "response": "Final answer.",
        })
        parser = ResponseFieldStreamParser()
        result = stream_through(parser, payload, 3)
        assert result == "Final answer."
        assert "internal nested text" not in result

    def test_random_chunk_boundaries_fuzz(self):
        rng = random.Random(42)
        for _ in range(25):
            parser = ResponseFieldStreamParser()
            out, i = [], 0
            while i < len(SAMPLE_JSON):
                step = rng.randint(1, 9)
                out.append(parser.feed(SAMPLE_JSON[i:i + step]))
                i += step
            assert "".join(out) == EXPECTED_RESPONSE


class TestPlainMode:
    def test_plain_text_passes_through(self):
        parser = ResponseFieldStreamParser()
        text = "Eine ganz normale Antwort ohne JSON."
        assert stream_through(parser, text, 5) == text

    def test_leading_whitespace_then_plain(self):
        parser = ResponseFieldStreamParser()
        assert stream_through(parser, "   Hallo Welt", 4) == "   Hallo Welt"

    def test_strict_mode_suppresses_plain_preamble_and_streams_response(self):
        parser = ResponseFieldStreamParser(allow_plain_text=False)
        text = (
            "Thinking about database query internals...\n"
            '{"query": "retrieve hidden text", "program": "emba"}'
            '{"response": "Only this answer is visible."}'
        )
        result = stream_through(parser, text, 4)
        assert result == "Only this answer is visible."
        assert "Thinking" not in result
        assert "retrieve hidden text" not in result

    def test_strict_mode_recovers_after_noisy_retrieval_stream(self):
        parser = ResponseFieldStreamParser(allow_plain_text=False)
        text = (
            'Need retrieve_context(query="cost {EMBA", program="emba"). '
            '[programme: emba | source: db]\n'
            'Snippet with unmatched brace { and quoted "response" text. '
            '{"query": "database lookup only", "program": "emba"}'
            '{"response": "Streaming resumes after retrieval."}'
        )
        result = stream_through(parser, text, 3)
        assert result == "Streaming resumes after retrieval."
        assert "retrieve_context" not in result
        assert "database lookup only" not in result


class TestEdgeCases:
    def test_empty_feeds_are_safe(self):
        parser = ResponseFieldStreamParser()
        assert parser.feed("") == ""
        assert parser.feed('{"response": "ok"}') == "ok"

    def test_nothing_emitted_before_field_starts(self):
        parser = ResponseFieldStreamParser()
        assert parser.feed('{"resp') == ""
        assert parser.feed('onse": ') == ""
        assert parser.feed('"Hal') == "Hal"
        assert parser.feed('lo"}') == "lo"