Websocket handle kwargs

Author: kevinmessiaen

giskard/ml_worker/websocket/utils.py

@@ -1,10 +1,9 @@
-from typing import Any, Dict, List, Optional
-
 import logging
 import os
 import posixpath
 import shutil
 from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 from mlflow.store.artifact.artifact_repo import verify_artifact_path
 
@@ -277,6 +276,7 @@ def map_suite_input_ws(i: websocket.SuiteInput):
 
 def function_argument_to_ws(value: Dict[str, Any]):
     args = list()
+    kwargs = dict()
 
     for v in value:
         obj = value[v]
@@ -310,10 +310,18 @@ def function_argument_to_ws(value: Dict[str, Any]):
             funcargs = websocket.FuncArgument(name=v, str=obj, none=False)
         elif isinstance(obj, bool):
             funcargs = websocket.FuncArgument(name=v, bool=obj, none=False)
-        elif isinstance(obj, dict):
-            funcargs = websocket.FuncArgument(name=v, kwargs=str(obj), none=False)
         else:
-            raise IllegalArgumentError("Unknown argument type")
+            kwargs[v] = obj
+            continue
         args.append(funcargs)
 
+    if len(kwargs) > 0:
+        args.append(
+            websocket.FuncArgument(
+                name="kwargs",
+                kwargs="\n".join([f"kwargs[{repr(key)}] = {repr(value)}" for key, value in kwargs.items()]),
+                none=False,
+            )
+        )
+
     return args