Add some docs, enhance naming and handle Pipeline inputs too

Author: silvanocerza

src/hayhooks/server/pipelines/models.py

@@ -1,6 +1,7 @@
-from hayhooks.server.utils.create_valid_type import create_valid_type
 from pandas import DataFrame
-from pydantic import BaseModel, create_model, ConfigDict
+from pydantic import BaseModel, ConfigDict, create_model
+
+from hayhooks.server.utils.create_valid_type import handle_unsupported_types
 
 
 class PipelineDefinition(BaseModel):
@@ -23,13 +24,16 @@ def get_request_model(pipeline_name: str, pipeline_inputs):
     config = ConfigDict(arbitrary_types_allowed=True)
 
     for component_name, inputs in pipeline_inputs.items():
-
         component_model = {}
         for name, typedef in inputs.items():
-            component_model[name] = (typedef["type"], typedef.get("default_value", ...))
-        request_model[component_name] = (create_model('ComponentParams', **component_model, __config__=config), ...)
+            input_type = handle_unsupported_types(typedef["type"], {DataFrame: dict})
+            component_model[name] = (
+                input_type,
+                typedef.get("default_value", ...),
+            )
+        request_model[component_name] = (create_model("ComponentParams", **component_model, __config__=config), ...)
 
-    return create_model(f'{pipeline_name.capitalize()}RunRequest', **request_model, __config__=config)
+    return create_model(f"{pipeline_name.capitalize()}RunRequest", **request_model, __config__=config)
 
 
 def get_response_model(pipeline_name: str, pipeline_outputs):
@@ -47,10 +51,10 @@ def get_response_model(pipeline_name: str, pipeline_outputs):
         component_model = {}
         for name, typedef in outputs.items():
             output_type = typedef["type"]
-            component_model[name] = (create_valid_type(output_type, { DataFrame: dict}), ...)
-        response_model[component_name] = (create_model('ComponentParams', **component_model, __config__=config), ...)
+            component_model[name] = (handle_unsupported_types(output_type, {DataFrame: dict}), ...)
+        response_model[component_name] = (create_model("ComponentParams", **component_model, __config__=config), ...)
 
-    return create_model(f'{pipeline_name.capitalize()}RunResponse', **response_model, __config__=config)
+    return create_model(f"{pipeline_name.capitalize()}RunResponse", **response_model, __config__=config)
 
 
 def convert_component_output(component_output):

src/hayhooks/server/utils/create_valid_type.py

@@ -1,32 +1,43 @@
-from typing import get_type_hints, Dict, get_origin, get_args
-from typing_extensions import TypedDict
-from types import GenericAlias
 from inspect import isclass
+from types import GenericAlias
+from typing import Dict, Union, get_args, get_origin, get_type_hints
+
+from typing_extensions import TypedDict
+
 
-def create_valid_type(typed_object:type, invalid_types:Dict[type, type]):
-    """ 
-        Returns a new type specification, replacing invalid_types in typed_object.
-        example: replace_invalid_types(ExtractedAnswer, {DataFrame: List}]) returns 
-        a TypedDict with DataFrame types replaced with List
+def handle_unsupported_types(type_: type, types_mapping: Dict[type, type]) -> Union[GenericAlias, type]:
     """
-    def validate_type(v):
+    Recursively handle types that are not supported by Pydantic by replacing them with the given types mapping.
+
+    :param type_: Type to replace if not supported
+    :param types_mapping: Mapping of types to replace
+    """
+
+    def _handle_generics(t_) -> GenericAlias:
+        """
+        Handle generics recursively
+        """
         child_typing = []
-        for t in get_args(v):
-            if t in invalid_types:
-                result = invalid_types[t]
+        for t in get_args(t_):
+            if t in types_mapping:
+                result = types_mapping[t]
             elif isclass(t):
-                result = create_valid_type(t, invalid_types)
-            else: result = t
+                result = handle_unsupported_types(t, types_mapping)
+            else:
+                result = t
             child_typing.append(result)
-        return GenericAlias(get_origin(v), tuple(child_typing))
-    if isclass(typed_object):
-        new_typing = {}
-        for k, v in get_type_hints(typed_object).items():
-            if(get_args(v) != ()):
-                new_typing[k] = validate_type(v)
-            else: new_typing[k] = v
-        if new_typing == {}:
-            return typed_object
-        else: return TypedDict(typed_object.__name__, new_typing)
-    else:
-        return validate_type(typed_object)
+        return GenericAlias(get_origin(t_), tuple(child_typing))
+
+    if isclass(type_):
+        new_type = {}
+        for arg_name, arg_type in get_type_hints(type_).items():
+            if get_args(arg_type):
+                new_type[arg_name] = _handle_generics(arg_type)
+            else:
+                new_type[arg_name] = arg_type
+        if new_type:
+            return TypedDict(type_.__name__, new_type)
+
+        return type_
+
+    return _handle_generics(type_)