Fix for llm message post-processing (#1106)

* fix for llm message post-processing

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

Author: vazirim

src/pdl/pdl_llms.py

@@ -17,7 +17,7 @@
 )
 from .pdl_lazy import PdlConst, PdlLazy, lazy_apply
 from .pdl_schema_utils import pdltype_to_jsonschema
-from .pdl_utils import remove_none_values_from_message
+from .pdl_utils import message_post_processing
 
 # Load environment variables
 load_dotenv()
@@ -47,7 +47,7 @@ async def async_generate_text(
             if msg.role is None:
                 msg.role = "assistant"
             return (
-                remove_none_values_from_message(msg.json()),
+                message_post_processing(msg.json()),
                 response.json(),  # pyright: ignore
             )
         except httpx.RequestError as exc:
@@ -161,7 +161,7 @@ def generate_text_stream(
             msg = chunk.choices[0].delta  # pyright: ignore
             if msg.role is None:
                 msg.role = "assistant"
-            yield remove_none_values_from_message(msg.model_dump())
+            yield message_post_processing(msg.model_dump())
         return result
 
 

src/pdl/pdl_utils.py

@@ -108,16 +108,23 @@ def get_contribute_value(
     return None
 
 
-def remove_none_values_from_message(message: dict) -> dict[str, Any]:
+def message_post_processing(message: dict) -> dict[str, Any]:
     ret = {}
     for key, value in message.items():
-        if key == "content":
+        if key == "content" and value is not None:
+            ret[key] = value
+        elif (
+            key == "reasoning_content" and value is not None
+        ):  # TODO: replacing reasoning_content with content here
+            key = "content"
             ret[key] = value
         if value is not None:
             if isinstance(value, dict):
-                ret[key] = remove_none_values_from_message(value)
+                ret[key] = message_post_processing(value)
             else:
                 ret[key] = value
+    if "content" not in ret:
+        ret["content"] = ""
     return ret