custom code: change type hinting execute_workflow_and_wait to handle pydantic BaseModel

Author: anth2o

src/mistralai/client/workflows.py

@@ -20,6 +20,7 @@
 
 # region imports
 import asyncio
+from pydantic import BaseModel
 import time
 # endregion imports
 
@@ -55,7 +56,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,
@@ -68,7 +69,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
@@ -152,7 +153,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,
@@ -165,7 +166,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