added to_pydantic_model as a method under the toolschema class

Author: twishabansal

packages/toolbox-core/src/toolbox_core/protocol.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 from inspect import Parameter
-from typing import Optional, Type
+from typing import Optional, Type, Any, cast
+
+from pydantic import BaseModel, Field, create_model
 
-from pydantic import BaseModel
 
 
 class ParameterSchema(BaseModel):
@@ -62,6 +63,18 @@ class ToolSchema(BaseModel):
     parameters: list[ParameterSchema]
     authRequired: list[str] = []
 
+    def to_pydantic_model(self) -> Type[BaseModel]:
+        """Converts the given manifest schema to a Pydantic BaseModel class."""
+        field_definitions = {}
+        for field in self.parameters:
+            field_definitions[field.name] = cast(
+                Any,
+                (
+                    field.to_param().annotation,
+                    Field(description=field.description),
+                ),
+            )
+        return create_model("tool_model", **field_definitions)
 
 class ManifestSchema(BaseModel):
     """

packages/toolbox-core/src/toolbox_core/tool.py

@@ -171,7 +171,7 @@ async def __call__(self, *args: Any, **kwargs: Any) -> str:
         payload = all_args.arguments
 
         # Perform argument type checks using pydantic
-        model = _schema_to_model(self.__name__, self.__schema)
+        model = self.__schema.to_pydantic_model()
         model.model_validate(payload)
 
         # apply bounded parameters
@@ -304,18 +304,3 @@ def identify_required_authn_params(
         if required:
             required_params[param] = services
     return required_params
-
-
-def _schema_to_model(model_name: str, tool_schema: ToolSchema) -> Type[BaseModel]:
-    """Converts the given manifest schema to a Pydantic BaseModel class."""
-    field_definitions = {}
-    for field in tool_schema.parameters:
-        field_definitions[field.name] = cast(
-            Any,
-            (
-                field.to_param().annotation,
-                Field(description=field.description),
-            ),
-        )
-
-    return create_model(model_name, **field_definitions)