62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
|
from typing import (
|
||
|
|
Dict,
|
||
|
|
List,
|
||
|
|
Type,
|
||
|
|
Union,
|
||
|
|
cast,
|
||
|
|
get_args,
|
||
|
|
get_origin,
|
||
|
|
)
|
||
|
|
|
||
|
|
from pydantic import BaseModel
|
||
|
|
|
||
|
|
from guardrails.utils.safe_get import safe_get
|
||
|
|
|
||
|
|
|
||
|
|
def convert_pydantic_model_to_openai_fn(
|
||
|
|
model: Union[Type[BaseModel], Type[List[Type[BaseModel]]]],
|
||
|
|
) -> Dict:
|
||
|
|
"""Convert a Pydantic BaseModel to an OpenAI function.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model: The Pydantic BaseModel to convert.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
OpenAI function paramters.
|
||
|
|
"""
|
||
|
|
|
||
|
|
schema_model = model
|
||
|
|
|
||
|
|
type_origin = get_origin(model)
|
||
|
|
if type_origin is list:
|
||
|
|
item_types = get_args(model)
|
||
|
|
if len(item_types) > 1:
|
||
|
|
raise ValueError("List data type must have exactly one child.")
|
||
|
|
# No List[List] support; we've already declared that in our types
|
||
|
|
schema_model = safe_get(item_types, 0)
|
||
|
|
|
||
|
|
schema_model = cast(Type[BaseModel], schema_model)
|
||
|
|
|
||
|
|
# Convert Pydantic model to JSON schema
|
||
|
|
json_schema = schema_model.model_json_schema()
|
||
|
|
json_schema["title"] = schema_model.__name__
|
||
|
|
|
||
|
|
if type_origin is list:
|
||
|
|
json_schema = {
|
||
|
|
"title": f"Array<{json_schema.get('title')}>",
|
||
|
|
"type": "array",
|
||
|
|
"items": json_schema,
|
||
|
|
}
|
||
|
|
|
||
|
|
# Create OpenAI function parameters
|
||
|
|
fn_params = {
|
||
|
|
"name": json_schema["title"],
|
||
|
|
"parameters": json_schema,
|
||
|
|
}
|
||
|
|
if "description" in json_schema and json_schema["description"] is not None:
|
||
|
|
fn_params["description"] = json_schema["description"]
|
||
|
|
|
||
|
|
# TODO: Update this to tools
|
||
|
|
# Wrap in { "type": "function", "function": fn_params}
|
||
|
|
return fn_params
|