Add pdl__model_input field to model blocks and trace fix (#785)

Author: mandel

pdl-live-react/src/pdl_ast.d.ts

@@ -2180,6 +2180,11 @@ export type Input =
 export type Modelresponse = string | null
 export type CompletionTokens = number | null
 export type PromptTokens = number | null
+export type PdlModelInput =
+  | {
+      [k: string]: unknown
+    }[]
+  | null
 /**
  * Optional field to ensure that the block is using granite-io.
  *
@@ -2326,6 +2331,11 @@ export type Input1 =
  *
  */
 export type Modelresponse1 = string | null
+export type PdlModelInput1 =
+  | {
+      [k: string]: unknown
+    }[]
+  | null
 /**
  * Optional field to ensure that the block is using LiteLLM.
  *
@@ -2787,6 +2797,7 @@ export interface LitellmModelBlock {
    *
    */
   pdl__usage?: PdlUsage | null
+  pdl__model_input?: PdlModelInput1
   platform?: Platform1
   parameters?: Parameters1
 }
@@ -2856,6 +2867,7 @@ export interface GraniteioModelBlock {
    *
    */
   pdl__usage?: PdlUsage | null
+  pdl__model_input?: PdlModelInput
   platform?: Platform
   backend: Backend
   processor?: Processor

src/pdl/pdl-schema.json

@@ -4236,6 +4236,21 @@
           "default": null,
           "description": "Tokens consumed during model call\n    "
         },
+        "pdl__model_input": {
+          "anyOf": [
+            {
+              "items": {
+                "type": "object"
+              },
+              "type": "array"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Pdl  Model Input"
+        },
         "platform": {
           "const": "granite-io",
           "default": "granite-io",
@@ -6686,6 +6701,21 @@
           "default": null,
           "description": "Tokens consumed during model call\n    "
         },
+        "pdl__model_input": {
+          "anyOf": [
+            {
+              "items": {
+                "type": "object"
+              },
+              "type": "array"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null,
+          "title": "Pdl  Model Input"
+        },
         "platform": {
           "const": "litellm",
           "default": "litellm",

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)

src/pdl/pdl_interpreter.py

@@ -1277,25 +1277,21 @@ def process_call_model(
             assert False
     # evaluate input
     model_input: ModelInput
-    if concrete_block.input is not None:  # If not implicit, then input must be a block
-        model_input_future, _, _, input_trace = process_block_of(
-            concrete_block,
-            "input",
-            state.with_yield_result(False).with_yield_background(False),
-            scope,
-            loc,
-        )
-        model_input_result = model_input_future.result()
-        if isinstance(model_input_result, str):
-            model_input = [{"role": state.role, "content": model_input_result}]
-        else:
-            model_input = model_input_result
+    model_input_future, _, _, concrete_block = process_block_of(
+        concrete_block,
+        "input",
+        state.with_yield_result(False).with_yield_background(False),
+        scope,
+        loc,
+    )
+    model_input_result = model_input_future.result()
+    if isinstance(model_input_result, str):
+        model_input = [{"role": state.role, "content": model_input_result}]
     else:
-        model_input = scope["pdl_context"]  # pyright: ignore
-        input_trace = None
+        model_input = model_input_result
     concrete_block = concrete_block.model_copy(
         update={
-            "input": input_trace,
+            "pdl__model_input": model_input,
         }
     )
     # Execute model call