instructor

1
from typing import Any, Callable, cast
2
import pytest
3
import instructor
4

5
from openai import OpenAI
6
from pydantic import BaseModel
7

8
from instructor.distil import (
9
    Instructions,
10
    format_function,
11
    get_signature_from_fn,
12
    is_return_type_base_model_or_instance,
13
)
14

15
client = instructor.patch(OpenAI())
16

17
instructions = Instructions(
18
    name="test_distil",
19
)
20

21

22
class SimpleModel(BaseModel):  # type: ignore[misc]
23
    data: int
24

25

26
def test_must_have_hint() -> None:
27
    with pytest.raises(AssertionError):
28

29
        @instructions.distil
30
        def test_func(x: int):  # type: ignore[no-untyped-def]
31
            return SimpleModel(data=x)
32

33

34
def test_must_be_base_model() -> None:
35
    with pytest.raises(AssertionError):
36

37
        @instructions.distil
38
        def test_func(x: int) -> int:
39
            return SimpleModel(data=x)
40

41

42
def test_is_return_type_base_model_or_instance() -> None:
43
    def valid_function() -> SimpleModel:
44
        return SimpleModel(data=1)
45

46
    def invalid_function() -> int:
47
        return 1
48

49
    assert is_return_type_base_model_or_instance(valid_function)
50
    assert not is_return_type_base_model_or_instance(invalid_function)
51

52

53
def test_get_signature_from_fn() -> None:
54
    def test_function(a: int, b: str) -> float:  # type: ignore[empty-body]
55
        """Sample docstring"""
56
        pass
57

58
    result = get_signature_from_fn(test_function)
59
    expected = "def test_function(a: int, b: str) -> float"
60
    assert expected in result
61
    assert "Sample docstring" in result
62

63

64
def test_format_function() -> None:
65
    def sample_function(x: int) -> SimpleModel:
66
        """This is a docstring."""
67
        return SimpleModel(data=x)
68

69
    formatted = format_function(sample_function)
70
    assert "def sample_function(x: int) -> SimpleModel:" in formatted
71
    assert '"""This is a docstring."""' in formatted
72
    assert "return SimpleModel(data=x)" in formatted
73

74

75
def test_distil_decorator_without_arguments() -> None:
76
    @instructions.distil
77
    def test_func(x: int) -> SimpleModel:
78
        return SimpleModel(data=x)
79

80
    casted_test_func = cast(Callable[[int], SimpleModel], test_func)
81
    result: SimpleModel = casted_test_func(42)
82
    assert result.data == 42
83

84

85
def test_distil_decorator_with_name_argument() -> None:
86
    @instructions.distil(name="custom_name")
87
    def another_test_func(x: int) -> SimpleModel:
88
        return SimpleModel(data=x)
89

90
    casted_another_test_func = cast(Callable[[int], SimpleModel], another_test_func)
91
    result: SimpleModel = casted_another_test_func(55)
92
    assert result.data == 55
93

94

95
# Mock track function for decorator tests
96
def mock_track(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> None:
97
    pass
98

99

100
def fn(a: int, b: int) -> int:
101
    return client.chat.completions.create(
102
        messages=[], model="davinci", response_model=SimpleModel
103
    )
104