Fix in MessageBlock (#1410)

Signed-off-by: Mandana Vaziri <mvaziri@us.ibm.com>

Author: vazirim

examples/test.pdl

@@ -0,0 +1,10 @@
+defs:
+  b: hi
+text:
+- role: assistant
+  content: hello
+  tool_calls: 
+  - data: 
+      a: ${ b }
+      c: d 
+- ${ pdl_context }

src/pdl/pdl-schema.json

@@ -3629,6 +3629,21 @@
           "$ref": "#/$defs/OptionalExpressionStr",
           "default": null,
           "description": "The id of the tool invocation for which this message is the tool response."
+        },
+        "tool_calls": {
+          "anyOf": [
+            {
+              "items": {
+                "$ref": "#/$defs/BlockType"
+              },
+              "type": "array"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Tool Calls"
         }
       },
       "required": [

src/pdl/pdl_ast.py

@@ -867,6 +867,7 @@ class MessageBlock(LeafBlock):
     """For example, the name of the tool that was invoked, for which this message is the tool response."""
     tool_call_id: OptionalExpressionStr = None
     """The id of the tool invocation for which this message is the tool response."""
+    tool_calls: Optional[list["BlockType"]] = None
 
 
 class IfBlock(StructuredBlock):

src/pdl/pdl_dumper.py

@@ -216,6 +216,10 @@ def block_to_dict(  # noqa: C901
                 d["object"] = [block_to_dict(b, json_compatible) for b in block.object]
         case MessageBlock():
             d["content"] = block_to_dict(block.content, json_compatible)
+            if block.tool_calls:
+                d["tool_calls"] = [
+                    block_to_dict(b, json_compatible) for b in block.tool_calls
+                ]
         case ReadBlock():
             d["read"] = expr_to_dict(block.read, json_compatible)
             d["message"] = block.message

src/pdl/pdl_interpreter.py

@@ -838,6 +838,16 @@ def process_block_body(
             if block.tool_call_id is not None:
                 tool_call_id, block = process_expr_of(block, "tool_call_id", scope, loc)
                 message["tool_call_id"] = tool_call_id
+            if block.tool_calls is not None:
+                result, _, _, _ = process_blocks_of(
+                    block,
+                    "tool_calls",
+                    JoinArray(as_="array"),  # pyright: ignore
+                    state,
+                    scope,
+                    loc,
+                )
+                message["tool_calls"] = result
             result = PdlConst(SingletonContext(PdlDict(message)))
             background = SingletonContext(PdlDict(message))
         case IfBlock():

src/pdl/pdl_utils.py

@@ -125,6 +125,8 @@ def message_post_processing(message: dict) -> dict[str, Any]:
     for key, value in message.items():
         if key == "content" and value is not None:
             ret[key] = value
+        if key == "tools_calls" and value is not None:
+            ret[key] = value
         if value is not None:
             if isinstance(value, dict):
                 ret[key] = message_post_processing(value)

tests/test_examples_run.yaml

@@ -39,6 +39,7 @@ skip:
   - examples/skeleton-of-thought/tips.pdl
   - examples/tutorial/sdk/lib.pdl
   - build/lib/pdl/pdl_stdlib.pdl
+  - examples/test.pdl
 
 with_inputs:
   examples/tutorial/programs/chatbot.pdl: