openai/openai-python

Public

mirrored from https://github.com/openai/openai-pythonAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
v1.49.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

tests/lib/test_pydantic.py

237lines · modeblame

bf1ca86cRobert Craigie1 years ago1from __future__ import annotations
2
09350cb5Robert Craigie1 years ago3from enum import Enum
4
5from pydantic import Field, BaseModel
bf1ca86cRobert Craigie1 years ago6from inline_snapshot import snapshot
7
8import openai
9from openai._compat import PYDANTIC_V2
10
11from .schema_types.query import Query
12
13
14def test_most_types() -> None:
15if PYDANTIC_V2:
16assert openai.pydantic_function_tool(Query)["function"] == snapshot(
17{
18"name": "Query",
19"strict": True,
20"parameters": {
21"$defs": {
22"Column": {
23"enum": [
24"id",
25"status",
26"expected_delivery_date",
27"delivered_at",
28"shipped_at",
29"ordered_at",
30"canceled_at",
31],
32"title": "Column",
33"type": "string",
34},
35"Condition": {
36"properties": {
37"column": {"title": "Column", "type": "string"},
38"operator": {"$ref": "#/$defs/Operator"},
39"value": {
40"anyOf": [
41{"type": "string"},
42{"type": "integer"},
43{"$ref": "#/$defs/DynamicValue"},
44],
45"title": "Value",
46},
47},
48"required": ["column", "operator", "value"],
49"title": "Condition",
50"type": "object",
51"additionalProperties": False,
52},
53"DynamicValue": {
54"properties": {"column_name": {"title": "Column Name", "type": "string"}},
55"required": ["column_name"],
56"title": "DynamicValue",
57"type": "object",
58"additionalProperties": False,
59},
60"Operator": {"enum": ["=", ">", "<", "<=", ">=", "!="], "title": "Operator", "type": "string"},
61"OrderBy": {"enum": ["asc", "desc"], "title": "OrderBy", "type": "string"},
62"Table": {"enum": ["orders", "customers", "products"], "title": "Table", "type": "string"},
63},
64"properties": {
54156ad1Robert Craigie1 years ago65"name": {"anyOf": [{"type": "string"}, {"type": "null"}], "title": "Name"},
bf1ca86cRobert Craigie1 years ago66"table_name": {"$ref": "#/$defs/Table"},
67"columns": {
68"items": {"$ref": "#/$defs/Column"},
69"title": "Columns",
70"type": "array",
71},
72"conditions": {
73"items": {"$ref": "#/$defs/Condition"},
74"title": "Conditions",
75"type": "array",
76},
77"order_by": {"$ref": "#/$defs/OrderBy"},
78},
54156ad1Robert Craigie1 years ago79"required": ["name", "table_name", "columns", "conditions", "order_by"],
bf1ca86cRobert Craigie1 years ago80"title": "Query",
81"type": "object",
82"additionalProperties": False,
83},
84}
85)
86else:
87assert openai.pydantic_function_tool(Query)["function"] == snapshot(
88{
89"name": "Query",
90"strict": True,
91"parameters": {
92"title": "Query",
93"type": "object",
94"properties": {
54156ad1Robert Craigie1 years ago95"name": {"title": "Name", "type": "string"},
bf1ca86cRobert Craigie1 years ago96"table_name": {"$ref": "#/definitions/Table"},
97"columns": {"type": "array", "items": {"$ref": "#/definitions/Column"}},
98"conditions": {
99"title": "Conditions",
100"type": "array",
101"items": {"$ref": "#/definitions/Condition"},
102},
103"order_by": {"$ref": "#/definitions/OrderBy"},
104},
54156ad1Robert Craigie1 years ago105"required": ["name", "table_name", "columns", "conditions", "order_by"],
bf1ca86cRobert Craigie1 years ago106"definitions": {
107"Table": {
108"title": "Table",
109"description": "An enumeration.",
110"enum": ["orders", "customers", "products"],
111"type": "string",
112},
113"Column": {
114"title": "Column",
115"description": "An enumeration.",
116"enum": [
117"id",
118"status",
119"expected_delivery_date",
120"delivered_at",
121"shipped_at",
122"ordered_at",
123"canceled_at",
124],
125"type": "string",
126},
127"Operator": {
128"title": "Operator",
129"description": "An enumeration.",
130"enum": ["=", ">", "<", "<=", ">=", "!="],
131"type": "string",
132},
133"DynamicValue": {
134"title": "DynamicValue",
135"type": "object",
136"properties": {"column_name": {"title": "Column Name", "type": "string"}},
137"required": ["column_name"],
ef859c88Alex Protsyk1 years ago138"additionalProperties": False,
bf1ca86cRobert Craigie1 years ago139},
140"Condition": {
141"title": "Condition",
142"type": "object",
143"properties": {
144"column": {"title": "Column", "type": "string"},
145"operator": {"$ref": "#/definitions/Operator"},
146"value": {
147"title": "Value",
148"anyOf": [
149{"type": "string"},
150{"type": "integer"},
151{"$ref": "#/definitions/DynamicValue"},
152],
153},
154},
155"required": ["column", "operator", "value"],
ef859c88Alex Protsyk1 years ago156"additionalProperties": False,
bf1ca86cRobert Craigie1 years ago157},
158"OrderBy": {
159"title": "OrderBy",
160"description": "An enumeration.",
161"enum": ["asc", "desc"],
162"type": "string",
163},
164},
165"additionalProperties": False,
166},
167}
168)
09350cb5Robert Craigie1 years ago169
170
171class Color(Enum):
172RED = "red"
173BLUE = "blue"
174GREEN = "green"
175
176
177class ColorDetection(BaseModel):
178color: Color = Field(description="The detected color")
179hex_color_code: str = Field(description="The hex color code of the detected color")
180
181
182def test_enums() -> None:
183if PYDANTIC_V2:
184assert openai.pydantic_function_tool(ColorDetection)["function"] == snapshot(
185{
186"name": "ColorDetection",
187"strict": True,
188"parameters": {
189"$defs": {"Color": {"enum": ["red", "blue", "green"], "title": "Color", "type": "string"}},
190"properties": {
ad5ac9fcRobert Craigie1 years ago191"color": {
192"description": "The detected color",
193"enum": ["red", "blue", "green"],
194"title": "Color",
195"type": "string",
196},
09350cb5Robert Craigie1 years ago197"hex_color_code": {
198"description": "The hex color code of the detected color",
199"title": "Hex Color Code",
200"type": "string",
201},
202},
203"required": ["color", "hex_color_code"],
204"title": "ColorDetection",
205"type": "object",
206"additionalProperties": False,
207},
208}
209)
210else:
211assert openai.pydantic_function_tool(ColorDetection)["function"] == snapshot(
212{
213"name": "ColorDetection",
214"strict": True,
215"parameters": {
216"properties": {
ad5ac9fcRobert Craigie1 years ago217"color": {
218"description": "The detected color",
219"title": "Color",
220"enum": ["red", "blue", "green"],
221},
09350cb5Robert Craigie1 years ago222"hex_color_code": {
223"description": "The hex color code of the detected color",
224"title": "Hex Color Code",
225"type": "string",
226},
227},
228"required": ["color", "hex_color_code"],
229"title": "ColorDetection",
230"definitions": {
231"Color": {"title": "Color", "description": "An enumeration.", "enum": ["red", "blue", "green"]}
232},
233"type": "object",
234"additionalProperties": False,
235},
236}
237)