instructor

1
from typing import TypeVar
2

3
import pytest
4
from anthropic.types import Message, Usage
5
from openai.resources.chat.completions import ChatCompletion
6
from pydantic import BaseModel, ValidationError
7

8
import instructor
9
from instructor import OpenAISchema, openai_schema
10
from instructor.exceptions import IncompleteOutputException
11

12
T = TypeVar("T")
13

14

15
@pytest.fixture  # type: ignore[misc]
16
def test_model() -> type[OpenAISchema]:
17
    class TestModel(OpenAISchema):  # type: ignore[misc]
18
        name: str = "TestModel"
19
        data: str
20

21
    return TestModel
22

23

24
@pytest.fixture  # type: ignore[misc]
25
def mock_completion(request: T) -> ChatCompletion:
26
    finish_reason = "stop"
27
    data_content = '{\n"data": "complete data"\n}'
28

29
    if hasattr(request, "param"):
30
        finish_reason = request.param.get("finish_reason", finish_reason)
31
        data_content = request.param.get("data_content", data_content)
32

33
    mock_choices = [
34
        {
35
            "index": 0,
36
            "message": {
37
                "role": "assistant",
38
                "function_call": {"name": "TestModel", "arguments": data_content},
39
                "content": data_content,
40
            },
41
            "finish_reason": finish_reason,
42
        }
43
    ]
44

45
    completion = ChatCompletion(
46
        id="test_id",
47
        choices=mock_choices,
48
        created=1234567890,
49
        model="gpt-3.5-turbo",
50
        object="chat.completion",
51
    )
52

53
    return completion
54

55

56
@pytest.fixture  # type: ignore[misc]
57
def mock_anthropic_message(request: T) -> Message:
58
    data_content = '{\n"data": "Claude says hi"\n}'
59
    if hasattr(request, "param"):
60
        data_content = request.param.get("data_content", data_content)
61
    return Message(
62
        id="test_id",
63
        content=[{"type": "text", "text": data_content}],
64
        model="claude-3-haiku-20240307",
65
        role="assistant",
66
        stop_reason="end_turn",
67
        stop_sequence=None,
68
        type="message",
69
        usage=Usage(
70
            input_tokens=100,
71
            output_tokens=100,
72
        ),
73
    )
74

75

76
def test_openai_schema() -> None:
77
    @openai_schema
78
    class Dataframe(BaseModel):  # type: ignore[misc]
79
        """
80
        Class representing a dataframe. This class is used to convert
81
        data into a frame that can be used by pandas.
82
        """
83

84
        data: str
85
        columns: str
86

87
        def to_pandas(self) -> None:
88
            pass
89

90
    assert hasattr(Dataframe, "openai_schema")
91
    assert hasattr(Dataframe, "from_response")
92
    assert hasattr(Dataframe, "to_pandas")
93
    assert Dataframe.openai_schema["name"] == "Dataframe"
94

95

96
def test_openai_schema_raises_error() -> None:
97
    with pytest.raises(TypeError, match="must be a subclass of pydantic.BaseModel"):
98

99
        @openai_schema
100
        class Dummy:
101
            pass
102

103

104
def test_no_docstring() -> None:
105
    class Dummy(OpenAISchema):  # type: ignore[misc]
106
        attr: str
107

108
    assert (
109
        Dummy.openai_schema["description"]
110
        == "Correctly extracted `Dummy` with all the required parameters with correct types"
111
    )
112

113

114
@pytest.mark.parametrize(
115
    "mock_completion",
116
    [{"finish_reason": "length", "data_content": '{\n"data": "incomplete dat"\n}'}],
117
    indirect=True,
118
)  # type: ignore[misc]
119
def test_incomplete_output_exception(
120
    test_model: type[OpenAISchema], mock_completion: ChatCompletion
121
) -> None:
122
    with pytest.raises(IncompleteOutputException):
123
        test_model.from_response(mock_completion, mode=instructor.Mode.FUNCTIONS)
124

125

126
def test_complete_output_no_exception(
127
    test_model: type[OpenAISchema], mock_completion: ChatCompletion
128
) -> None:
129
    test_model_instance = test_model.from_response(
130
        mock_completion, mode=instructor.Mode.FUNCTIONS
131
    )
132
    assert test_model_instance.data == "complete data"
133

134

135
@pytest.mark.asyncio  # type: ignore[misc]
136
@pytest.mark.parametrize(
137
    "mock_completion",
138
    [{"finish_reason": "length", "data_content": '{\n"data": "incomplete dat"\n}'}],
139
    indirect=True,
140
)  # type: ignore[misc]
141
def test_incomplete_output_exception_raise(
142
    test_model: type[OpenAISchema], mock_completion: ChatCompletion
143
) -> None:
144
    with pytest.raises(IncompleteOutputException):
145
        test_model.from_response(mock_completion, mode=instructor.Mode.FUNCTIONS)
146

147

148
def test_anthropic_no_exception(
149
    test_model: type[OpenAISchema], mock_anthropic_message: Message
150
) -> None:
151
    test_model_instance = test_model.from_response(
152
        mock_anthropic_message, mode=instructor.Mode.ANTHROPIC_JSON
153
    )
154
    assert test_model_instance.data == "Claude says hi"
155

156

157
@pytest.mark.parametrize(
158
    "mock_anthropic_message",
159
    [{"data_content": '{\n"data": "Claude likes\ncontrol\ncharacters"\n}'}],
160
    indirect=True,
161
)  # type: ignore[misc]
162
def test_control_characters_not_allowed_in_anthropic_json_strict_mode(
163
    test_model: type[OpenAISchema], mock_anthropic_message: Message
164
) -> None:
165
    with pytest.raises(ValidationError) as exc_info:
166
        test_model.from_response(
167
            mock_anthropic_message, mode=instructor.Mode.ANTHROPIC_JSON, strict=True
168
        )
169

170
    # https://docs.pydantic.dev/latest/errors/validation_errors/#json_invalid
171
    exc = exc_info.value
172
    assert len(exc.errors()) == 1
173
    assert exc.errors()[0]["type"] == "json_invalid"
174
    assert "control character" in exc.errors()[0]["msg"]
175

176

177
@pytest.mark.parametrize(
178
    "mock_anthropic_message",
179
    [{"data_content": '{\n"data": "Claude likes\ncontrol\ncharacters"\n}'}],
180
    indirect=True,
181
)  # type: ignore[misc]
182
def test_control_characters_allowed_in_anthropic_json_non_strict_mode(
183
    test_model: type[OpenAISchema], mock_anthropic_message: Message
184
) -> None:
185
    test_model_instance = test_model.from_response(
186
        mock_anthropic_message, mode=instructor.Mode.ANTHROPIC_JSON, strict=False
187
    )
188
    assert test_model_instance.data == "Claude likes\ncontrol\ncharacters"
189