instructor

1
import json
2
import pytest
3
from instructor.utils import (
4
    classproperty,
5
    extract_json_from_codeblock,
6
    extract_json_from_stream,
7
    extract_json_from_stream_async,
8
    merge_consecutive_messages,
9
)
10

11

12
def test_extract_json_from_codeblock():
13
    example = """
14
    Here is a response
15

16
    ```json
17
    {
18
        "key": "value"
19
    }    
20
    ```
21
    """
22
    result = extract_json_from_codeblock(example)
23
    assert json.loads(result) == {"key": "value"}
24

25

26
def test_extract_json_from_codeblock_no_end():
27
    example = """
28
    Here is a response
29

30
    ```json
31
    {
32
        "key": "value",
33
        "another_key": [{"key": {"key": "value"}}]
34
    }  
35
    """
36
    result = extract_json_from_codeblock(example)
37
    assert json.loads(result) == {
38
        "key": "value",
39
        "another_key": [{"key": {"key": "value"}}],
40
    }
41

42

43
def test_extract_json_from_codeblock_no_start():
44
    example = """
45
    Here is a response
46

47
    {
48
        "key": "value",
49
        "another_key": [{"key": {"key": "value"}}, {"key": "value"}]
50
    }
51
    """
52
    result = extract_json_from_codeblock(example)
53
    assert json.loads(result) == {
54
        "key": "value",
55
        "another_key": [{"key": {"key": "value"}}, {"key": "value"}],
56
    }
57

58

59
def test_stream_json():
60
    text = """here is the json for you! 
61
    
62
    ```json
63
    , here
64
    {
65
        "key": "value",
66
        "another_key": [{"key": {"key": "value"}}]
67
    }
68
    ```
69

70
    What do you think?
71
    """
72

73
    def batch_strings(chunks, n=2):
74
        batch = ""
75
        for chunk in chunks:
76
            for char in chunk:
77
                batch += char
78
                if len(batch) == n:
79
                    yield batch
80
                    batch = ""
81
        if batch:  # Yield any remaining characters in the last batch
82
            yield batch
83

84
    result = json.loads(
85
        "".join(list(extract_json_from_stream(batch_strings(text, n=3))))
86
    )
87
    assert result == {"key": "value", "another_key": [{"key": {"key": "value"}}]}
88

89

90
@pytest.mark.asyncio
91
async def test_stream_json_async():
92
    text = """here is the json for you! 
93
    
94
    ```json
95
    , here
96
    {
97
        "key": "value",
98
        "another_key": [{"key": {"key": "value"}}, {"key": "value"}]
99
    }
100
    ```
101

102
    What do you think?
103
    """
104

105
    async def batch_strings_async(chunks, n=2):
106
        batch = ""
107
        for chunk in chunks:
108
            for char in chunk:
109
                batch += char
110
                if len(batch) == n:
111
                    yield batch
112
                    batch = ""
113
        if batch:  # Yield any remaining characters in the last batch
114
            yield batch
115

116
    result = json.loads(
117
        "".join(
118
            [
119
                chunk
120
                async for chunk in extract_json_from_stream_async(
121
                    batch_strings_async(text, n=3)
122
                )
123
            ]
124
        )
125
    )
126
    assert result == {
127
        "key": "value",
128
        "another_key": [{"key": {"key": "value"}}, {"key": "value"}],
129
    }
130

131

132
def test_merge_consecutive_messages():
133
    messages = [
134
        {"role": "user", "content": "Hello"},
135
        {"role": "user", "content": "How are you"},
136
        {"role": "assistant", "content": "Hello"},
137
        {"role": "assistant", "content": "I am good"},
138
    ]
139
    result = merge_consecutive_messages(messages)
140
    assert result == [
141
        {
142
            "role": "user",
143
            "content": [
144
                {"type": "text", "text": "Hello"},
145
                {"type": "text", "text": "How are you"},
146
            ],
147
        },
148
        {
149
            "role": "assistant",
150
            "content": [
151
                {"type": "text", "text": "Hello"},
152
                {"type": "text", "text": "I am good"},
153
            ],
154
        },
155
    ]
156

157

158
def test_merge_consecutive_messages_empty():
159
    messages = []
160
    result = merge_consecutive_messages(messages)
161
    assert result == []
162

163

164
def test_merge_consecutive_messages_single():
165
    messages = [
166
        {"role": "user", "content": "Hello"},
167
        {"role": "assistant", "content": "Hello"},
168
    ]
169
    result = merge_consecutive_messages(messages)
170
    assert result == [
171
        {"role": "user", "content": [{"type": "text", "text": "Hello"}]},
172
        {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]},
173
    ]
174

175

176
def test_classproperty():
177
    """Test custom `classproperty` descriptor."""
178

179
    class MyClass:
180
        @classproperty
181
        def my_property(cls):
182
            return cls
183

184
    assert MyClass.my_property is MyClass
185

186
    class MyClass:
187
        clvar = 1
188

189
        @classproperty
190
        def my_property(cls):
191
            return cls.clvar
192

193
    assert MyClass.my_property == 1
194