wehub-resource-sync
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
from langchain_core.utils.function_calling import convert_to_openai_function
|
|
from pydantic import BaseModel, Field
|
|
|
|
|
|
def test_convert_pydantic_to_openai_function() -> None:
|
|
class Data(BaseModel):
|
|
"""The data to return."""
|
|
|
|
key: str = Field(..., description="API key")
|
|
days: int = Field(default=0, description="Number of days to forecast")
|
|
|
|
actual = convert_to_openai_function(Data)
|
|
expected = {
|
|
"name": "Data",
|
|
"description": "The data to return.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"key": {"description": "API key", "type": "string"},
|
|
"days": {
|
|
"description": "Number of days to forecast",
|
|
"default": 0,
|
|
"type": "integer",
|
|
},
|
|
},
|
|
"required": ["key"],
|
|
},
|
|
}
|
|
assert actual == expected
|
|
|
|
|
|
def test_convert_pydantic_to_openai_function_nested() -> None:
|
|
class Data(BaseModel):
|
|
"""The data to return."""
|
|
|
|
key: str = Field(..., description="API key")
|
|
days: int = Field(default=0, description="Number of days to forecast")
|
|
|
|
class Model(BaseModel):
|
|
"""The model to return."""
|
|
|
|
data: Data
|
|
|
|
actual = convert_to_openai_function(Model)
|
|
expected = {
|
|
"name": "Model",
|
|
"description": "The model to return.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"data": {
|
|
"description": "The data to return.",
|
|
"type": "object",
|
|
"properties": {
|
|
"key": {
|
|
"description": "API key",
|
|
"type": "string",
|
|
},
|
|
"days": {
|
|
"description": "Number of days to forecast",
|
|
"default": 0,
|
|
"type": "integer",
|
|
},
|
|
},
|
|
"required": ["key"],
|
|
},
|
|
},
|
|
"required": ["data"],
|
|
},
|
|
}
|
|
assert actual == expected
|