Add pdl__model_input field to model blocks and trace fix

Author: mandel

src/pdl/pdl_ast.py

@@ -410,6 +410,7 @@ class ModelBlock(LeafBlock):
     pdl__usage: Optional[PdlUsage] = None
     """Tokens consumed during model call
     """
+    pdl__model_input: Optional[ModelInput] = None
 
 
 class LitellmModelBlock(ModelBlock):

src/pdl/pdl_ast_utils.py

@@ -43,8 +43,7 @@ def iter_block_children(f: Callable[[BlockType], None], block: BlockType) -> Non
             if block.pdl__trace is not None:
                 f(block.pdl__trace)
         case ModelBlock():
-            if block.input is not None:
-                f(block.input)
+            f(block.input)
         case CodeBlock():
             f(block.code)
         case GetBlock():
@@ -133,12 +132,10 @@ def map_block_children(f: MappedFunctions, block: BlockType) -> BlockType:
                 block.pdl__trace = f.f_block(block.pdl__trace)
         case LitellmModelBlock():
             block.model = f.f_expr(block.model)
-            if block.input is not None:
-                block.input = f.f_block(block.input)
+            block.input = f.f_block(block.input)
         case GraniteioModelBlock():
             block.model = f.f_expr(block.model)
-            if block.input is not None:
-                block.input = f.f_block(block.input)
+            block.input = f.f_block(block.input)
             if block.parameters is not None:
                 block.parameters = f.f_expr(block.parameters)
         case CodeBlock():

src/pdl/pdl_dumper.py

@@ -127,6 +127,8 @@ def block_to_dict(  # noqa: C901
                 d["modelResponse"] = block.modelResponse
             if block.pdl__usage is not None:
                 d["pdl__usage"] = usage_to_dict(block.pdl__usage)
+            if block.pdl__model_input is not None:
+                d["pdl__model_input"] = block.pdl__model_input
         case GraniteioModelBlock():
             d["model"] = expr_to_dict(block.model, json_compatible)
             d["platform"] = str(block.platform)