113 lines
3.1 KiB
113 lines
3.1 KiB
from typing import Any, List
|
|
|
|
from docstring_parser import parse
|
|
from pydantic import BaseModel
|
|
|
|
|
|
def _remove_a_key(d: dict, remove_key: str) -> None:
|
|
"""Remove a key from a dictionary recursively"""
|
|
if isinstance(d, dict):
|
|
for key in list(d.keys()):
|
|
if key == remove_key and "type" in d.keys():
|
|
del d[key]
|
|
else:
|
|
_remove_a_key(d[key], remove_key)
|
|
|
|
|
|
def base_model_to_openai_function(
|
|
pydantic_type: type[BaseModel],
|
|
output_str: bool = False,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Convert a Pydantic model to a dictionary representation of functions.
|
|
|
|
Args:
|
|
pydantic_type (type[BaseModel]): The Pydantic model type to convert.
|
|
|
|
Returns:
|
|
dict[str, Any]: A dictionary representation of the functions.
|
|
|
|
"""
|
|
schema = pydantic_type.model_json_schema()
|
|
|
|
docstring = parse(pydantic_type.__doc__ or "")
|
|
parameters = {
|
|
k: v
|
|
for k, v in schema.items()
|
|
if k not in ("title", "description")
|
|
}
|
|
|
|
for param in docstring.params:
|
|
if (name := param.arg_name) in parameters["properties"] and (
|
|
description := param.description
|
|
):
|
|
if "description" not in parameters["properties"][name]:
|
|
parameters["properties"][name]["description"] = description
|
|
|
|
parameters["type"] = "object"
|
|
|
|
if "description" not in schema:
|
|
if docstring.short_description:
|
|
schema["description"] = docstring.short_description
|
|
else:
|
|
schema["description"] = (
|
|
f"Correctly extracted `{pydantic_type.__name__}` with all "
|
|
f"the required parameters with correct types"
|
|
)
|
|
|
|
_remove_a_key(parameters, "title")
|
|
_remove_a_key(parameters, "additionalProperties")
|
|
|
|
if output_str:
|
|
out = {
|
|
"function_call": {
|
|
"name": pydantic_type.__name__,
|
|
},
|
|
"functions": [
|
|
{
|
|
"name": pydantic_type.__name__,
|
|
"description": schema["description"],
|
|
"parameters": parameters,
|
|
},
|
|
],
|
|
}
|
|
return str(out)
|
|
|
|
else:
|
|
return {
|
|
"function_call": {
|
|
"name": pydantic_type.__name__,
|
|
},
|
|
"functions": [
|
|
{
|
|
"name": pydantic_type.__name__,
|
|
"description": schema["description"],
|
|
"parameters": parameters,
|
|
},
|
|
],
|
|
}
|
|
|
|
|
|
def multi_base_model_to_openai_function(
|
|
pydantic_types: List[BaseModel] = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Converts multiple Pydantic types to a dictionary of functions.
|
|
|
|
Args:
|
|
pydantic_types (List[BaseModel]]): A list of Pydantic types to convert.
|
|
|
|
Returns:
|
|
dict[str, Any]: A dictionary containing the converted functions.
|
|
|
|
"""
|
|
functions: list[dict[str, Any]] = [
|
|
base_model_to_openai_function(pydantic_type)["functions"][0]
|
|
for pydantic_type in pydantic_types
|
|
]
|
|
|
|
return {
|
|
"function_call": "auto",
|
|
"functions": functions,
|
|
}
|