Spaces:
Sleeping
Sleeping
File size: 4,214 Bytes
c96b98a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | from typing import Any, Dict, Union, get_args, get_origin, Optional
from phi.utils.log import logger
def get_json_type_for_py_type(arg: str) -> str:
"""
Get the JSON schema type for a given type.
:param arg: The type to get the JSON schema type for.
:return: The JSON schema type.
"""
# logger.info(f"Getting JSON type for: {arg}")
if arg in ("int", "float", "complex", "Decimal"):
return "number"
elif arg in ("str", "string"):
return "string"
elif arg in ("bool", "boolean"):
return "boolean"
elif arg in ("NoneType", "None"):
return "null"
elif arg in ("list", "tuple", "set", "frozenset"):
return "array"
elif arg in ("dict", "mapping"):
return "object"
# If the type is not recognized, return "object"
return "object"
def get_json_schema_for_arg(t: Any) -> Optional[Dict[str, Any]]:
# logger.info(f"Getting JSON schema for arg: {t}")
type_args = get_args(t)
# logger.info(f"Type args: {type_args}")
type_origin = get_origin(t)
# logger.info(f"Type origin: {type_origin}")
if type_origin is not None:
if type_origin in (list, tuple, set, frozenset):
json_schema_for_items = get_json_schema_for_arg(type_args[0]) if type_args else {"type": "string"}
return {"type": "array", "items": json_schema_for_items}
elif type_origin is dict:
# Handle both key and value types for dictionaries
key_schema = get_json_schema_for_arg(type_args[0]) if type_args else {"type": "string"}
value_schema = get_json_schema_for_arg(type_args[1]) if len(type_args) > 1 else {"type": "string"}
return {"type": "object", "propertyNames": key_schema, "additionalProperties": value_schema}
elif type_origin is Union:
types = []
for arg in type_args:
if arg is not type(None):
try:
schema = get_json_schema_for_arg(arg)
if schema:
types.append(schema)
except Exception:
continue
return {"anyOf": types} if types else None
return {"type": get_json_type_for_py_type(t.__name__)}
def get_json_schema(
type_hints: Dict[str, Any], param_descriptions: Optional[Dict[str, str]] = None, strict: bool = False
) -> Dict[str, Any]:
json_schema: Dict[str, Any] = {
"type": "object",
"properties": {},
}
if strict:
json_schema["additionalProperties"] = False
for k, v in type_hints.items():
# logger.info(f"Parsing arg: {k} | {v}")
if k == "return":
continue
try:
# Check if type is Optional (Union with NoneType)
type_origin = get_origin(v)
type_args = get_args(v)
is_optional = type_origin is Union and len(type_args) == 2 and any(arg is type(None) for arg in type_args)
# Get the actual type if it's Optional
if is_optional:
v = next(arg for arg in type_args if arg is not type(None))
# Handle cases with no type hint
if v:
arg_json_schema = get_json_schema_for_arg(v)
else:
arg_json_schema = {}
if arg_json_schema is not None:
if is_optional:
# Handle null type for optional fields
if isinstance(arg_json_schema["type"], list):
arg_json_schema["type"].append("null")
else:
arg_json_schema["type"] = [arg_json_schema["type"], "null"]
# Add description
if param_descriptions and k in param_descriptions and param_descriptions[k]:
arg_json_schema["description"] = param_descriptions[k]
json_schema["properties"][k] = arg_json_schema
else:
logger.warning(f"Could not parse argument {k} of type {v}")
except Exception as e:
logger.error(f"Error processing argument {k}: {str(e)}")
continue
return json_schema
|