Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/pdl/pdl_llms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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


Expand Down
13 changes: 10 additions & 3 deletions src/pdl/pdl_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down