openai/openai-python

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
v1.41.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

tests/lib/test_pydantic.py

235lines · modecode

1from __future__ import annotations
2
3from enum import Enum
4
5from pydantic import Field, BaseModel
6from 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:
15 if PYDANTIC_V2:
16 assert 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": {
65 "table_name": {"$ref": "#/$defs/Table"},
66 "columns": {
67 "items": {"$ref": "#/$defs/Column"},
68 "title": "Columns",
69 "type": "array",
70 },
71 "conditions": {
72 "items": {"$ref": "#/$defs/Condition"},
73 "title": "Conditions",
74 "type": "array",
75 },
76 "order_by": {"$ref": "#/$defs/OrderBy"},
77 },
78 "required": ["table_name", "columns", "conditions", "order_by"],
79 "title": "Query",
80 "type": "object",
81 "additionalProperties": False,
82 },
83 }
84 )
85 else:
86 assert openai.pydantic_function_tool(Query)["function"] == snapshot(
87 {
88 "name": "Query",
89 "strict": True,
90 "parameters": {
91 "title": "Query",
92 "type": "object",
93 "properties": {
94 "table_name": {"$ref": "#/definitions/Table"},
95 "columns": {"type": "array", "items": {"$ref": "#/definitions/Column"}},
96 "conditions": {
97 "title": "Conditions",
98 "type": "array",
99 "items": {"$ref": "#/definitions/Condition"},
100 },
101 "order_by": {"$ref": "#/definitions/OrderBy"},
102 },
103 "required": ["table_name", "columns", "conditions", "order_by"],
104 "definitions": {
105 "Table": {
106 "title": "Table",
107 "description": "An enumeration.",
108 "enum": ["orders", "customers", "products"],
109 "type": "string",
110 },
111 "Column": {
112 "title": "Column",
113 "description": "An enumeration.",
114 "enum": [
115 "id",
116 "status",
117 "expected_delivery_date",
118 "delivered_at",
119 "shipped_at",
120 "ordered_at",
121 "canceled_at",
122 ],
123 "type": "string",
124 },
125 "Operator": {
126 "title": "Operator",
127 "description": "An enumeration.",
128 "enum": ["=", ">", "<", "<=", ">=", "!="],
129 "type": "string",
130 },
131 "DynamicValue": {
132 "title": "DynamicValue",
133 "type": "object",
134 "properties": {"column_name": {"title": "Column Name", "type": "string"}},
135 "required": ["column_name"],
136 "additionalProperties": False,
137 },
138 "Condition": {
139 "title": "Condition",
140 "type": "object",
141 "properties": {
142 "column": {"title": "Column", "type": "string"},
143 "operator": {"$ref": "#/definitions/Operator"},
144 "value": {
145 "title": "Value",
146 "anyOf": [
147 {"type": "string"},
148 {"type": "integer"},
149 {"$ref": "#/definitions/DynamicValue"},
150 ],
151 },
152 },
153 "required": ["column", "operator", "value"],
154 "additionalProperties": False,
155 },
156 "OrderBy": {
157 "title": "OrderBy",
158 "description": "An enumeration.",
159 "enum": ["asc", "desc"],
160 "type": "string",
161 },
162 },
163 "additionalProperties": False,
164 },
165 }
166 )
167
168
169class Color(Enum):
170 RED = "red"
171 BLUE = "blue"
172 GREEN = "green"
173
174
175class ColorDetection(BaseModel):
176 color: Color = Field(description="The detected color")
177 hex_color_code: str = Field(description="The hex color code of the detected color")
178
179
180def test_enums() -> None:
181 if PYDANTIC_V2:
182 assert openai.pydantic_function_tool(ColorDetection)["function"] == snapshot(
183 {
184 "name": "ColorDetection",
185 "strict": True,
186 "parameters": {
187 "$defs": {"Color": {"enum": ["red", "blue", "green"], "title": "Color", "type": "string"}},
188 "properties": {
189 "color": {
190 "description": "The detected color",
191 "enum": ["red", "blue", "green"],
192 "title": "Color",
193 "type": "string",
194 },
195 "hex_color_code": {
196 "description": "The hex color code of the detected color",
197 "title": "Hex Color Code",
198 "type": "string",
199 },
200 },
201 "required": ["color", "hex_color_code"],
202 "title": "ColorDetection",
203 "type": "object",
204 "additionalProperties": False,
205 },
206 }
207 )
208 else:
209 assert openai.pydantic_function_tool(ColorDetection)["function"] == snapshot(
210 {
211 "name": "ColorDetection",
212 "strict": True,
213 "parameters": {
214 "properties": {
215 "color": {
216 "description": "The detected color",
217 "title": "Color",
218 "enum": ["red", "blue", "green"],
219 },
220 "hex_color_code": {
221 "description": "The hex color code of the detected color",
222 "title": "Hex Color Code",
223 "type": "string",
224 },
225 },
226 "required": ["color", "hex_color_code"],
227 "title": "ColorDetection",
228 "definitions": {
229 "Color": {"title": "Color", "description": "An enumeration.", "enum": ["red", "blue", "green"]}
230 },
231 "type": "object",
232 "additionalProperties": False,
233 },
234 }
235 )
236