remove None valued parameters when calling LLMs (#212)

* remove None valued parameters when calling LLMs

Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

src/pdl/pdl_interpreter.py

@@ -1125,6 +1125,7 @@ def get_transformed_inputs(kwargs):
 
         litellm.input_callback = [get_transformed_inputs]
         # append_log(state, "Model Input", messages_to_str(model_input))
+
         msg, raw_result = yield from generate_client_response(
             state, concrete_block, model_input
         )

src/pdl/pdl_llms.py

@@ -16,6 +16,7 @@
     set_default_model_params,
     set_structured_decoding_parameters,
 )
+from .pdl_utils import remove_none_values_from_message
 
 # Load environment variables
 load_dotenv()
@@ -165,7 +166,10 @@ def generate_text(
         msg = response.choices[0].message  # pyright: ignore
         if msg.role is None:
             msg.role = "assistant"
-        return msg.json(), response.json()  # pyright: ignore
+        return (
+            remove_none_values_from_message(msg.json()),
+            response.json(),  # pyright: ignore
+        )
 
     @staticmethod
     def generate_text_stream(
@@ -191,5 +195,5 @@ def generate_text_stream(
             msg = chunk.choices[0].delta  # pyright: ignore
             if msg.role is None:
                 msg.role = "assistant"
-            yield msg.model_dump()
+            yield remove_none_values_from_message(msg.model_dump())
         return result

src/pdl/pdl_utils.py

@@ -1,5 +1,5 @@
 import json
-from typing import Sequence
+from typing import Any, Sequence
 
 from .pdl_ast import ContributeTarget, ContributeValue, FunctionBlock, Message, Messages
 
@@ -70,3 +70,16 @@ def simple_message(message: Message) -> bool:
     if message.keys() == {"role", "content"} and message["content"] is not None:
         return True
     return False
+
+
+def remove_none_values_from_message(message: Any) -> dict[str, Any]:
+    ret = {}
+    for key, value in message.items():
+        if key == "content":
+            ret[key] = value
+        if value is not None:
+            if isinstance(value, dict):
+                ret[key] = remove_none_values_from_message(value)
+            else:
+                ret[key] = value
+    return ret