openai/openai-python
Publicmirrored from https://github.com/openai/openai-pythonAvailable
examples/parsing_tools.py
80lines · modeblame
bf1ca86cRobert Craigie1 years ago | 1 | from enum import Enum |
| 2 | from typing import List, Union | |
| 3 | | |
| 4 | import rich | |
| 5 | from pydantic import BaseModel | |
| 6 | | |
| 7 | import openai | |
| 8 | from openai import OpenAI | |
| 9 | | |
| 10 | | |
| 11 | class Table(str, Enum): | |
| 12 | orders = "orders" | |
| 13 | customers = "customers" | |
| 14 | products = "products" | |
| 15 | | |
| 16 | | |
| 17 | class Column(str, Enum): | |
| 18 | id = "id" | |
| 19 | status = "status" | |
| 20 | expected_delivery_date = "expected_delivery_date" | |
| 21 | delivered_at = "delivered_at" | |
| 22 | shipped_at = "shipped_at" | |
| 23 | ordered_at = "ordered_at" | |
| 24 | canceled_at = "canceled_at" | |
| 25 | | |
| 26 | | |
| 27 | class Operator(str, Enum): | |
| 28 | eq = "=" | |
| 29 | gt = ">" | |
| 30 | lt = "<" | |
| 31 | le = "<=" | |
| 32 | ge = ">=" | |
| 33 | ne = "!=" | |
| 34 | | |
| 35 | | |
| 36 | class OrderBy(str, Enum): | |
| 37 | asc = "asc" | |
| 38 | desc = "desc" | |
| 39 | | |
| 40 | | |
| 41 | class DynamicValue(BaseModel): | |
| 42 | column_name: str | |
| 43 | | |
| 44 | | |
| 45 | class Condition(BaseModel): | |
| 46 | column: str | |
| 47 | operator: Operator | |
| 48 | value: Union[str, int, DynamicValue] | |
| 49 | | |
| 50 | | |
| 51 | class Query(BaseModel): | |
| 52 | table_name: Table | |
| 53 | columns: List[Column] | |
| 54 | conditions: List[Condition] | |
| 55 | order_by: OrderBy | |
| 56 | | |
| 57 | | |
| 58 | client = OpenAI() | |
| 59 | | |
| 60 | completion = client.beta.chat.completions.parse( | |
| 61 | model="gpt-4o-2024-08-06", | |
| 62 | messages=[ | |
| 63 | { | |
| 64 | "role": "system", | |
| 65 | "content": "You are a helpful assistant. The current date is August 6, 2024. You help users query for the data they are looking for by calling the query function.", | |
| 66 | }, | |
| 67 | { | |
| 68 | "role": "user", | |
| 69 | "content": "look up all my orders in november of last year that were fulfilled but not delivered on time", | |
| 70 | }, | |
| 71 | ], | |
| 72 | tools=[ | |
| 73 | openai.pydantic_function_tool(Query), | |
| 74 | ], | |
| 75 | ) | |
| 76 | | |
| 77 | tool_call = (completion.choices[0].message.tool_calls or [])[0] | |
| 78 | rich.print(tool_call.function) | |
| 79 | assert isinstance(tool_call.function.parsed_arguments, Query) | |
| 80 | print(tool_call.function.parsed_arguments.table_name) |