Coverage for src/integry/utils/pydantic.py: 14%

1from typing import Optional, Type, Any 

2from pydantic import BaseModel, Field, create_model 

3from pydantic.fields import FieldInfo 

4from pydantic_core import PydanticUndefined 

5 

6 

7def get_pydantic_model_from_json_schema(json_schema: dict[str, Any]) -> Type[BaseModel]: 

8 """ 

9 Generates a Pydantic model class using the provided JSON schema. 

10 

11 Args: 

12 json_schema: The JSON schema to create the Pydantic model from. 

13 

14 Returns: 

15 A subclass of Pydantic's BaseModel class to be used as argument schema in various agentic frameworks. 

16 """ 

17 title = json_schema.get("title", "") 

18 fields = {} 

19 

20 for key, schema in json_schema.get("properties", {}).items(): 

21 updated_key, field_info = get_pydantic_field_from_json_schema( 

22 key, schema, json_schema.get("required", []) 

23 ) 

24 fields[updated_key] = field_info 

25 

26 return create_model(title, **fields) # type: ignore 

27 

28 

29JSON_SCHEMA_TYPE_TO_NATIVE_TYPE_MAP: dict[str, Type[Any]] = { 

30 "boolean": bool, 

31 "number": float, 

32 "string": str, 

33 "array": list, 

34 "object": dict, 

35} 

36 

37 

38def get_pydantic_type_from_json_schema( 

39 json_schema: dict[str, Any], 

40) -> Type[Any]: 

41 _type = json_schema.get("type") 

42 if not isinstance(_type, str): 

43 raise ValueError(f"type must be a string, got: {_type}") 

44 

45 native_type = JSON_SCHEMA_TYPE_TO_NATIVE_TYPE_MAP.get(_type) 

46 if native_type is None: 

47 raise ValueError(f"Unsupported type: {_type}") 

48 

49 if _type == "array" and (items_schema := json_schema.get("items")): 

50 return list[get_pydantic_type_from_json_schema(items_schema)] 

51 

52 if _type == "object": 

53 if not json_schema.get("properties"): 

54 # This occurs for custom fields parameters whose properties are only known 

55 # at runtime. 

56 return dict[str, Any] 

57 

58 return get_pydantic_model_from_json_schema(json_schema) 

59 

60 return native_type 

61 

62 

63def get_pydantic_field_from_json_schema( 

64 name: str, 

65 json_schema: dict[str, Any], 

66 required: list[str], 

67) -> tuple[str, tuple[Type[Any], FieldInfo]]: 

68 keywords = ["_cursor"] 

69 

70 description = json_schema.get("description") 

71 default = json_schema.get("default", PydanticUndefined) 

72 

73 alias = None 

74 if name in keywords: 

75 name = name.lstrip("_") 

76 # alias = name 

77 elif name.startswith("_"): 

78 # Pydantic does not allow fields to start with an underscore 

79 # TODO: Once Langchain fixes/adds support for aliases, use alias instead of raising exception 

80 raise ValueError( 

81 f"Parameters with names starting with an underscore are not supported: {name}" 

82 ) 

83 alias = name 

84 name = name.lstrip("_") 

85 

86 _type = get_pydantic_type_from_json_schema( 

87 json_schema, 

88 ) 

89 

90 if name not in required: 

91 default = None 

92 

93 field_info = ( 

94 _type, 

95 Field( 

96 description=description, 

97 default=default, 

98 alias=alias, 

99 ), 

100 ) 

101 

102 return (name, field_info)