custom code: change type hinting execute_workflow_and_wait to handle pydantic BaseModel

Author: anth2o

src/mistralai/client/workflows.py

@@ -21,6 +21,7 @@
 
 # region imports
 import asyncio
+from pydantic import BaseModel
 import time
 # endregion imports
 
@@ -58,7 +59,7 @@ def _init_sdks(self):
     def execute_workflow_and_wait(
         self,
         workflow_identifier: str,
-        input: OptionalNullable[Dict[str, Any]] = UNSET,
+        input: OptionalNullable[Dict[str, Any] | BaseModel] = UNSET,
         execution_id: OptionalNullable[str] = UNSET,
         deployment_name: OptionalNullable[str] = UNSET,
         custom_tracing_attributes: OptionalNullable[Dict[str, str]] = UNSET,
@@ -71,7 +72,7 @@ def execute_workflow_and_wait(
 
         Args:
             workflow_identifier: The workflow name or ID.
-            input: Input parameters for the workflow
+            input: The input to the workflow. This should be a dictionary or a BaseModel that matches the workflow's input schema.
             execution_id: Optional custom execution ID
             deployment_name: Name of the deployment to route this execution to
             custom_tracing_attributes: Custom tracing attributes
@@ -155,7 +156,7 @@ def wait_for_workflow_completion(
     async def execute_workflow_and_wait_async(
         self,
         workflow_identifier: str,
-        input: OptionalNullable[Dict[str, Any]] = UNSET,
+        input: OptionalNullable[Dict[str, Any] | BaseModel] = UNSET,
         execution_id: OptionalNullable[str] = UNSET,
         deployment_name: OptionalNullable[str] = UNSET,
         custom_tracing_attributes: OptionalNullable[Dict[str, str]] = UNSET,
@@ -168,7 +169,7 @@ async def execute_workflow_and_wait_async(
 
         Args:
             workflow_identifier: The workflow name or ID.
-            input: Input parameters for the workflow
+            input: The input to the workflow. This should be a dictionary or a BaseModel that matches the workflow's input schema.
             execution_id: Optional custom execution ID
             deployment_name: Name of the deployment to route this execution to
             custom_tracing_attributes: Custom tracing attributes