Skip to content

Commit

Permalink
add check
Browse files Browse the repository at this point in the history
  • Loading branch information
jxnl committed Oct 14, 2023
1 parent 2bb4703 commit bc2c6f9
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
3 changes: 3 additions & 0 deletions instructor/distil.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,9 @@ def is_return_type_base_model_or_instance(func: Callable[..., Any]) -> bool:
:return: True if the return type is a pydantic BaseModel or an instance of it.
"""
return_type = inspect.signature(func).return_annotation
assert (
return_type != inspect.Signature.empty
), "Must have a return type hint that is a pydantic BaseModel"
return inspect.isclass(return_type) and issubclass(return_type, BaseModel)


Expand Down
18 changes: 17 additions & 1 deletion tests/test_distil.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from pyexpat import model
import pytest
import openai
from pydantic import BaseModel
from instructor.distil import (
Expand All @@ -19,6 +19,22 @@ class SimpleModel(BaseModel):
data: int


def test_must_have_hint():
with pytest.raises(AssertionError):

@instructions.distil
def test_func(x: int):
return SimpleModel(data=x)


def test_must_be_base_model():
with pytest.raises(AssertionError):

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


def test_is_return_type_base_model_or_instance():
def valid_function() -> SimpleModel:
return SimpleModel(data=1)
Expand Down

0 comments on commit bc2c6f9

Please sign in to comment.