From 373c79d5f1a71ede6ff75041e5101c9950f7dea7 Mon Sep 17 00:00:00 2001
From: Vitor Hideyoshi <vitor.h.n.batista@gmail.com>
Date: Wed, 9 Apr 2025 22:14:34 -0300
Subject: [PATCH] Adds TypedDict for JsonSchema for Reference

---
 jambo/schema_converter.py       |  6 ++-
 jambo/types/json_schema_type.py | 80 +++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 jambo/types/json_schema_type.py

diff --git a/jambo/schema_converter.py b/jambo/schema_converter.py
index 196fc8f..676a78b 100644
--- a/jambo/schema_converter.py
+++ b/jambo/schema_converter.py
@@ -7,6 +7,8 @@ from pydantic.fields import Field
 
 from typing import Type
 
+from jambo.types.json_schema_type import JSONSchema
+
 
 class SchemaConverter:
     """
@@ -18,7 +20,7 @@ class SchemaConverter:
     """
 
     @staticmethod
-    def build(schema: dict) -> Type:
+    def build(schema: JSONSchema) -> Type:
         """
         Converts a JSON Schema to a Pydantic model.
         :param schema: The JSON Schema to convert.
@@ -32,7 +34,7 @@ class SchemaConverter:
     @staticmethod
     def build_object(
         name: str,
-        schema: dict,
+        schema: JSONSchema,
     ) -> Type:
         """
         Converts a JSON Schema object to a Pydantic model given a name.
diff --git a/jambo/types/json_schema_type.py b/jambo/types/json_schema_type.py
new file mode 100644
index 0000000..77ec7ff
--- /dev/null
+++ b/jambo/types/json_schema_type.py
@@ -0,0 +1,80 @@
+from typing import List, Dict, Union, TypedDict, Literal
+
+
+JSONSchemaType = Literal[
+    "string", "number", "integer", "boolean", "object", "array", "null"
+]
+
+
+JSONType = Union[str, int, float, bool, None, Dict[str, "JSONType"], List["JSONType"]]
+
+
+class JSONSchema(TypedDict, total=False):
+    # Basic metadata
+    title: str
+    description: str
+    default: JSONType
+    examples: List[JSONType]
+
+    # Type definitions
+    type: Union[JSONSchemaType, List[JSONSchemaType]]
+
+    # Object-specific keywords
+    properties: Dict[str, "JSONSchema"]
+    required: List[str]
+    additionalProperties: Union[bool, "JSONSchema"]
+    minProperties: int
+    maxProperties: int
+    patternProperties: Dict[str, "JSONSchema"]
+    dependencies: Dict[str, Union[List[str], "JSONSchema"]]
+
+    # Array-specific keywords
+    items: Union["JSONSchema", List["JSONSchema"]]
+    additionalItems: Union[bool, "JSONSchema"]
+    minItems: int
+    maxItems: int
+    uniqueItems: bool
+
+    # String-specific keywords
+    minLength: int
+    maxLength: int
+    pattern: str
+    format: str
+
+    # Number-specific keywords
+    minimum: float
+    maximum: float
+    exclusiveMinimum: float
+    exclusiveMaximum: float
+    multipleOf: float
+
+    # Enum and const
+    enum: List[JSONType]
+    const: JSONType
+
+    # Conditionals
+    if_: "JSONSchema"  # 'if' is a reserved word in Python
+    then: "JSONSchema"
+    else_: "JSONSchema"  # 'else' is also a reserved word
+
+    # Combination keywords
+    allOf: List["JSONSchema"]
+    anyOf: List["JSONSchema"]
+    oneOf: List["JSONSchema"]
+    not_: "JSONSchema"  # 'not' is a reserved word
+
+
+# Fix forward references
+JSONSchema.__annotations__["properties"] = Dict[str, JSONSchema]
+JSONSchema.__annotations__["items"] = Union[JSONSchema, List[JSONSchema]]
+JSONSchema.__annotations__["additionalItems"] = Union[bool, JSONSchema]
+JSONSchema.__annotations__["additionalProperties"] = Union[bool, JSONSchema]
+JSONSchema.__annotations__["patternProperties"] = Dict[str, JSONSchema]
+JSONSchema.__annotations__["dependencies"] = Dict[str, Union[List[str], JSONSchema]]
+JSONSchema.__annotations__["if_"] = JSONSchema
+JSONSchema.__annotations__["then"] = JSONSchema
+JSONSchema.__annotations__["else_"] = JSONSchema
+JSONSchema.__annotations__["allOf"] = List[JSONSchema]
+JSONSchema.__annotations__["anyOf"] = List[JSONSchema]
+JSONSchema.__annotations__["oneOf"] = List[JSONSchema]
+JSONSchema.__annotations__["not_"] = JSONSchema