fix saving dataset to HF, toolcall sanitizing (#246)

Author: willccbb

tests/test_environment_extra.py

@@ -20,8 +20,8 @@
 from verifiers.envs.environment import Environment
 from verifiers.parsers.parser import Parser
 from verifiers.rubrics.rubric import Rubric
-from verifiers.types import GenerateOutputs
-from verifiers.utils.tool_utils import sanitize_tool_calls
+from verifiers.types import GenerateOutputs, Info, Messages, SamplingArgs
+from verifiers.utils.message_utils import sanitize_tool_calls
 
 
 # Local simple concrete Environment for testing
@@ -30,16 +30,17 @@ async def rollout(
         self,
         client,
         model,
-        prompt,
+        prompt: Messages,
         answer: str = "",
         task: str = "default",
-        info: dict = {},
-        sampling_args: dict = {},
+        info: Info | None = {},
+        sampling_args: SamplingArgs | None = None,
         **kwargs,
     ):
         response = await self.get_model_response(
             prompt=prompt, client=client, model=model, sampling_args=sampling_args
         )
+        assert response is not None
         if self.message_type == "chat":
             completion = [
                 {"role": "assistant", "content": response.choices[0].message.content}

verifiers/envs/environment.py

@@ -27,8 +27,7 @@
     SamplingArgs,
     State,
 )
-from verifiers.utils.message_utils import cleanup_messages
-from verifiers.utils.tool_utils import sanitize_tool_calls
+from verifiers.utils.message_utils import cleanup_messages, sanitize_tool_calls
 
 if TYPE_CHECKING:
     from transformers.tokenization_utils_base import (  # type: ignore
@@ -519,6 +518,7 @@ def make_dataset(
         """
         Make a dataset from the evaluation results.
         """
+        # TODO: enable saving of multimodal datasets
         state_columns = state_columns or []
 
         if push_to_hub and hub_name is None:

verifiers/scripts/eval.py

@@ -12,7 +12,7 @@
 from openai import OpenAI
 
 import verifiers as vf
-from verifiers.utils.message_utils import messages_to_printable
+from verifiers.utils.message_utils import messages_to_printable, sanitize_tool_calls
 
 
 def eval_environment(
@@ -124,8 +124,8 @@ def eval_environment(
         tasks = results.task
         data_dict = {
             "id": ids,
-            "prompt": printable_prompts,
-            "completion": printable_completions,
+            "prompt": [sanitize_tool_calls(p) for p in printable_prompts],
+            "completion": [sanitize_tool_calls(c) for c in printable_completions],
             "task": tasks,
         }
         if results.info[0] != {}:
@@ -170,9 +170,7 @@ def eval_environment(
             print(f"Saved dataset to {results_path}")
         if save_to_hf_hub:
             if hf_hub_dataset_name == "":
-                dataset_name = (
-                    f"{env}_{model}_n={num_examples}_r={rollouts_per_example}"
-                )
+                dataset_name = f"{env}_{model.replace('/', '-')}_n{num_examples}_r{rollouts_per_example}"
             else:
                 dataset_name = hf_hub_dataset_name
             dataset.push_to_hub(dataset_name)

verifiers/utils/logging_utils.py

@@ -61,7 +61,9 @@ def _format_messages(messages) -> Text:
             out.append(content, style=style)
             if "tool_calls" in msg:
                 for tool_call in msg["tool_calls"]:
-                    tool_call_str = json.dumps(dict(tool_call["function"]), indent=2)
+                    name = getattr(tool_call.function, "name", "")
+                    args = getattr(tool_call.function, "arguments", {})
+                    tool_call_str = json.dumps({"name": name, "args": args}, indent=2)
                     out.append(f"\n\n[tool call]\n{tool_call_str}", style=style)
         return out
 

verifiers/utils/message_utils.py

@@ -1,104 +1,17 @@
-from collections.abc import Iterable
+import json
 from typing import cast
 
 from verifiers.types import ChatMessage, Messages
 
 
-def sanitize_object(obj: object):
-    """
-    Recursively convert Pydantic/OpenAI SDK objects to plain Python types
-    (dict/list/str/bool/number). Leaves primitives unchanged.
-    """
-    if isinstance(obj, (str, bytes, bytearray, int, float, bool)) or obj is None:
-        return obj
-    dump = getattr(obj, "model_dump", None)
-    if callable(dump):
-        obj = dump()
-    if isinstance(obj, dict):
-        return {k: sanitize_object(v) for k, v in obj.items()}
-    # check if obj is iterable
-    if isinstance(obj, Iterable):
-        return [sanitize_object(x) for x in obj]
-    return obj
-
-
-def sanitize_chat_message(message: ChatMessage):
-    """
-    input: chat message (dict or object)
-    output: chat message (dict)
-    """
-    # TODO: debug for multimodal messages; content can get consumed as an iterator
-    new_message = {}
-    dump = getattr(message, "model_dump", None)
-    if callable(dump):
-        new_message = dump()
-        return new_message
-    assert isinstance(message, dict)
-    assert isinstance(new_message, dict)
-    new_message["role"] = message["role"]
-    if "content" in message and message["content"]:
-        content = message["content"]
-        if isinstance(content, str):
-            new_message["content"] = content
-        else:
-            new_message["content"] = []
-            parts = list(content) if not isinstance(content, str) else content
-            for c in parts:
-                if isinstance(c, str):
-                    new_message["content"].append(c)
-                else:
-                    new_message["content"].append(sanitize_object(c))
-    if "tool_calls" in message and message["tool_calls"]:
-        tool_calls = list(message["tool_calls"])
-        new_message["tool_calls"] = [
-            sanitize_object(tool_call) for tool_call in tool_calls
-        ]
-    return new_message
-
-
-def sanitize_messages(messages: Messages) -> str | list:
-    """
-    input: list of dicts or Pydantic models, or str
-    output: list of dicts, or str
-    """
-    if isinstance(messages, str):
-        return messages
-    sanitized_list = [sanitize_chat_message(m) for m in list(messages)]
-    return sanitized_list
-
-
-def content_to_printable(content: object) -> str:
+def message_to_printable(message: ChatMessage) -> ChatMessage:
     """
-    Render content to readable text, handling multimodal lists.
-    - Text parts: return their text
-    - Image-like parts: return "[image]"
-    Falls back to str(content).
+    Removes image_url objects from message content.
     """
-    print(str(content)[:100])
-    if isinstance(content, str):
-        return content
-    if isinstance(content, dict):
-        if "type" in content and content["type"] == "text":
-            return content["text"]
-        if "type" in content and content["type"] in {
-            "image_url",
-            "input_image",
-            "image",
-        }:
-            return "[image]"
-    if isinstance(content, (list, tuple)):
-        out = []
-        for x in content:
-            out.append(content_to_printable(x))
-        return "\n\n".join(out)
-    return str(content)
-
-
-def message_to_printable(message: ChatMessage) -> ChatMessage:
     new_message = {}
     new_message["role"] = message["role"]
     new_message["content"] = []
-    if "tool_calls" in message and message["tool_calls"]:
+    if "tool_calls" in message:
         new_message["tool_calls"] = message["tool_calls"]
     content = message.get("content")
     if content is None:
@@ -121,6 +34,9 @@ def message_to_printable(message: ChatMessage) -> ChatMessage:
 
 
 def messages_to_printable(messages: Messages) -> Messages:
+    """
+    Removes image_url objects from messages.
+    """
     if isinstance(messages, str):
         return messages
     return [message_to_printable(m) for m in messages]
@@ -129,6 +45,8 @@ def messages_to_printable(messages: Messages) -> Messages:
 def cleanup_message(message: ChatMessage) -> ChatMessage:
     new_message = {}
     new_message["role"] = message["role"]
+    if "tool_calls" in message:
+        new_message["tool_calls"] = message["tool_calls"]
     new_message["content"] = []
     content = message.get("content")
     if content is None:
@@ -161,3 +79,26 @@ def cleanup_messages(messages: Messages) -> Messages:
     for m in messages:
         new_messages.append(cleanup_message(m))
     return new_messages
+
+
+def sanitize_tool_calls(messages: Messages):
+    """
+    Sanitize tool calls from messages.
+    """
+    if not isinstance(messages, list):
+        return messages
+    sanitized_messages = []
+    for m in messages:
+        if "tool_calls" in m:
+            new_m = {
+                "role": m["role"],
+                "content": m.get("content", ""),
+                "tool_calls": [
+                    json.dumps(tc.model_dump())  # type: ignore
+                    for tc in m.get("tool_calls", [])
+                ],
+            }
+            sanitized_messages.append(new_m)
+        else:
+            sanitized_messages.append(m)
+    return sanitized_messages

verifiers/utils/tool_utils.py

@@ -1,15 +1,13 @@
 from __future__ import annotations
 
 import inspect
-import json
 import re
 from typing import Any, Literal, Union, get_args, get_origin
 
 from verifiers.types import (
     ChatCompletionToolParam,
     FunctionParameters,
     JsonPrimitive,
-    Messages,
 )
 
 _JSON_PRIMITIVE_MAP: dict[type, JsonPrimitive] = {
@@ -178,26 +176,3 @@ def convert_func_to_oai_tool(func: Any) -> ChatCompletionToolParam:
             "parameters": parameters_schema,
         },
     }
-
-
-def sanitize_tool_calls(messages: Messages):
-    """
-    Sanitize tool calls from messages.
-    """
-    if not isinstance(messages, list):
-        return messages
-    sanitized_messages = []
-    for m in messages:
-        if "tool_calls" in m:
-            new_m = {
-                "role": m["role"],
-                "content": m.get("content", ""),
-                "tool_calls": [
-                    json.dumps(tc.model_dump())  # type: ignore
-                    for tc in m.get("tool_calls", [])
-                ],
-            }
-            sanitized_messages.append(new_m)
-        else:
-            sanitized_messages.append(m)
-    return sanitized_messages