mirror of
https://github.com/kennethreitz/instructor.git
synced 2026-06-05 22:50:18 +00:00
Jason
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
from typing import Literal
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from instructor import openai_schema, OpenAISchema, openai_function
|
|
|
|
|
|
def test_openai_schema():
|
|
@openai_schema
|
|
class Dataframe(BaseModel):
|
|
"""
|
|
Class representing a dataframe. This class is used to convert
|
|
data into a frame that can be used by pandas.
|
|
"""
|
|
|
|
data: str
|
|
columns: str
|
|
|
|
def to_pandas(self):
|
|
pass
|
|
|
|
assert hasattr(Dataframe, "openai_schema")
|
|
assert hasattr(Dataframe, "from_response")
|
|
assert hasattr(Dataframe, "to_pandas")
|
|
assert Dataframe.openai_schema["name"] == "Dataframe" # type: ignore
|
|
|
|
|
|
def test_openai_schema_raises_error():
|
|
with pytest.raises(TypeError, match="must be a subclass of pydantic.BaseModel"):
|
|
|
|
@openai_schema
|
|
class Dummy:
|
|
pass
|
|
|
|
|
|
def test_no_docstring():
|
|
class Dummy(OpenAISchema):
|
|
attr: str
|
|
|
|
assert (
|
|
Dummy.openai_schema["description"]
|
|
== "Correctly extracted `Dummy` with all the required parameters with correct types"
|
|
)
|
|
|
|
|
|
def test_openai_function():
|
|
@openai_function
|
|
def get_current_weather(
|
|
location: str, format: Literal["celsius", "fahrenheit"] = "celsius"
|
|
):
|
|
"""
|
|
Gets the current weather in a given location, use this function for any questions related to the weather
|
|
|
|
Parameters
|
|
----------
|
|
location
|
|
The city to get the weather, e.g. San Francisco. Guess the location from user messages
|
|
|
|
format
|
|
A string with the full content of what the given role said
|
|
"""
|
|
|
|
@openai_function
|
|
def get_current_weather_no_format_docstring(
|
|
location: str, format: Literal["celsius", "fahrenheit"] = "celsius"
|
|
):
|
|
"""
|
|
Gets the current weather in a given location, use this function for any questions related to the weather
|
|
|
|
Parameters
|
|
----------
|
|
location
|
|
The city to get the weather, e.g. San Francisco. Guess the location from user messages
|
|
"""
|
|
|
|
scheme_missing_param = get_current_weather_no_format_docstring.openai_schema
|
|
assert (
|
|
scheme_missing_param["parameters"]["properties"]["location"]["description"]
|
|
== "The city to get the weather, e.g. San Francisco. Guess the location from user messages"
|
|
)
|
|
assert scheme_missing_param["parameters"]["properties"]["format"]["enum"] == [
|
|
"celsius",
|
|
"fahrenheit",
|
|
]
|
|
with pytest.raises(KeyError, match="description"):
|
|
scheme_missing_param["parameters"]["properties"]["format"]["description"]
|