From db8f396451278756d13d3af2d46aa315fbbd18df Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Wed, 19 Mar 2025 10:18:08 -0400
Subject: [PATCH 01/14] PDL source code cleanup in the UI

---
 pdl-live-react/src/view/code/Code.tsx | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index b19e3d643..1c636b4e7 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -1,3 +1,4 @@
+import { match, P } from "ts-pattern"
 import { lazy, Suspense } from "react"
 import { stringify } from "yaml"
 
@@ -65,6 +66,17 @@ function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
     pdl__location: undefined,
     pdl__model_input: undefined,
   }
+  if (new_data.kind === "model") {
+    let input = match(new_data.input)
+      .with({ kind: "data" }, (i) => i.data)
+      .otherwise((_) => new_data.input)
+    input = match(input)
+      .with({ expr: P._ }, (e) => e.expr)
+      .otherwise((e) => e)
+    if (input) {
+      new_data.input = input
+    }
+  }
   // remove contribute: ["result", context]
   if (
     new_data?.contribute?.includes("result") &&
@@ -77,5 +89,14 @@ function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
     delete new_data.defs
   }
   // recursive cleanup
-  return map_block_children(block_code_cleanup, new_data)
+  const clean_data = map_block_children(block_code_cleanup, new_data)
+  return match(clean_data)
+    .with({ kind: "data" }, (d) => {
+      return match(d.data)
+        .with({ expr: P._ }, (e) => e.expr)
+        .with(P.union(P.string, P.number, P.boolean, {}), (e) => e)
+        .otherwise((_) => d)
+    })
+    .with({ kind: P._ }, (d) => ({ ...d, kind: undefined }))
+    .otherwise((d) => d)
 }

From f703b62a0d9cec900b5a2645afe37b2827872085 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Wed, 19 Mar 2025 21:32:18 -0400
Subject: [PATCH 02/14] Cleanup function body

---
 pdl-live-react/src/view/code/Code.tsx            | 2 +-
 pdl-live-react/src/view/detail/kind/Function.tsx | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index 1c636b4e7..ea2a25a33 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -50,7 +50,7 @@ export default function Code({
   )
 }
 
-function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
+export function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
   if (data === null || typeof data !== "object") {
     return data
   }
diff --git a/pdl-live-react/src/view/detail/kind/Function.tsx b/pdl-live-react/src/view/detail/kind/Function.tsx
index c6fccf1f7..b9ff334ce 100644
--- a/pdl-live-react/src/view/detail/kind/Function.tsx
+++ b/pdl-live-react/src/view/detail/kind/Function.tsx
@@ -2,6 +2,7 @@ import { stringify } from "yaml"
 
 import Group from "../Group"
 import CodeGroup from "../../code/CodeGroup"
+import {block_code_cleanup} from "../../code/Code"
 
 export default function FunctionItems({
   block: { def, function: func, return: retrn },
@@ -12,7 +13,7 @@ export default function FunctionItems({
     <>
       {def && <Group description={def} term="Name" />}
       <CodeGroup code={stringify(func)} term="Parameters" />
-      <CodeGroup code={stringify(retrn)} term="Body" />
+      <CodeGroup code={stringify(block_code_cleanup(retrn))} term="Body" />
     </>
   )
 }

From f30f5846c18a7fd883f71f3ecab9201068e3ba7e Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Wed, 19 Mar 2025 21:39:12 -0400
Subject: [PATCH 03/14] Formatting

---
 pdl-live-react/src/view/detail/kind/Function.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pdl-live-react/src/view/detail/kind/Function.tsx b/pdl-live-react/src/view/detail/kind/Function.tsx
index b9ff334ce..3d12c5ff7 100644
--- a/pdl-live-react/src/view/detail/kind/Function.tsx
+++ b/pdl-live-react/src/view/detail/kind/Function.tsx
@@ -2,7 +2,7 @@ import { stringify } from "yaml"
 
 import Group from "../Group"
 import CodeGroup from "../../code/CodeGroup"
-import {block_code_cleanup} from "../../code/Code"
+import { block_code_cleanup } from "../../code/Code"
 
 export default function FunctionItems({
   block: { def, function: func, return: retrn },

From e6ba89d548822f69c226b83ebcdcccf06dcf0ced Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Mon, 24 Mar 2025 13:54:14 -0400
Subject: [PATCH 04/14] Small fix

---
 pdl-live-react/src/view/code/Code.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index ea2a25a33..6fc88894b 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -93,7 +93,7 @@ export function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
   return match(clean_data)
     .with({ kind: "data" }, (d) => {
       return match(d.data)
-        .with({ expr: P._ }, (e) => e.expr)
+        .with({ pdl__expr: P._ }, (e) => ({ ...d, data: e.pdl__expr }))
         .with(P.union(P.string, P.number, P.boolean, {}), (e) => e)
         .otherwise((_) => d)
     })

From d61c82b2526ad0c91469bc3bb9b758c8eabf2133 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Mon, 24 Mar 2025 14:43:06 -0400
Subject: [PATCH 05/14] Code cleanup

---
 pdl-live-react/src/helpers.ts         |  2 ++
 pdl-live-react/src/view/code/Code.tsx | 38 +++++++++++++++++----------
 2 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/pdl-live-react/src/helpers.ts b/pdl-live-react/src/helpers.ts
index 7aeeb7114..0abdefe13 100644
--- a/pdl-live-react/src/helpers.ts
+++ b/pdl-live-react/src/helpers.ts
@@ -13,6 +13,8 @@ import type {
 /** Re-export for convenience */
 export type { PdlBlock } from "./pdl_ast"
 
+export type BlockType = null | boolean | number | string | PdlBlock
+
 type MakeNonNullable<T> = {
   [K in keyof T]-?: NonNullable<T[K]>
 }
diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index 6fc88894b..6e4833b99 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -4,6 +4,7 @@ import { stringify } from "yaml"
 
 import { tryJsonPrettyPrint } from "../../helpers"
 import { type PdlBlock } from "../../pdl_ast"
+import { type BlockType } from "../../helpers"
 import { map_block_children } from "../../pdl_ast_utils"
 
 const PreviewLight = lazy(() => import("./PreviewLight"))
@@ -50,12 +51,12 @@ export default function Code({
   )
 }
 
-export function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
+export function block_code_cleanup(data: BlockType): BlockType {
   if (data === null || typeof data !== "object") {
     return data
   }
   // remove pdl__result
-  const new_data = {
+  let new_data = {
     ...data,
     pdl__result: undefined,
     pdl__is_leaf: undefined,
@@ -66,17 +67,6 @@ export function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
     pdl__location: undefined,
     pdl__model_input: undefined,
   }
-  if (new_data.kind === "model") {
-    let input = match(new_data.input)
-      .with({ kind: "data" }, (i) => i.data)
-      .otherwise((_) => new_data.input)
-    input = match(input)
-      .with({ expr: P._ }, (e) => e.expr)
-      .otherwise((e) => e)
-    if (input) {
-      new_data.input = input
-    }
-  }
   // remove contribute: ["result", context]
   if (
     new_data?.contribute?.includes("result") &&
@@ -89,7 +79,27 @@ export function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
     delete new_data.defs
   }
   // recursive cleanup
-  const clean_data = map_block_children(block_code_cleanup, new_data)
+  const new_data_rec = map_block_children(block_code_cleanup, new_data)
+  // replace `data: literal` by `literal`
+  const clean_data = match(new_data_rec)
+    .with(
+      {
+        kind: "data",
+        data: P.union(P.string, P.number, P.boolean),
+        raw: false,
+        spec: P.nullish,
+        description: P.nullish,
+        defs: {},
+        def: P.nullish,
+        contribute: P.union(["context", "result"], ["result", "context"]),
+        parser: P.nullish,
+        fallback: P.nullish,
+        role: P.nullish,
+      },
+      (d) => d.data,
+    )
+    .otherwise((d) => d)
+  // remove kind
   return match(clean_data)
     .with({ kind: "data" }, (d) => {
       return match(d.data)

From 3a3bb10d8aae60b3511729ad0b2eba9137c8a3c9 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Mon, 24 Mar 2025 22:46:26 -0400
Subject: [PATCH 06/14] Improve code cleanup

---
 pdl-live-react/src/helpers.ts         |   1 +
 pdl-live-react/src/pdl_ast_utils.ts   | 145 ++++++++++++++++----------
 pdl-live-react/src/view/code/Code.tsx | 102 +++++++++++++-----
 3 files changed, 167 insertions(+), 81 deletions(-)

diff --git a/pdl-live-react/src/helpers.ts b/pdl-live-react/src/helpers.ts
index 0abdefe13..196f77a15 100644
--- a/pdl-live-react/src/helpers.ts
+++ b/pdl-live-react/src/helpers.ts
@@ -14,6 +14,7 @@ import type {
 export type { PdlBlock } from "./pdl_ast"
 
 export type BlockType = null | boolean | number | string | PdlBlock
+export type ExpressionT<T> = T | string | LocalizedExpression
 
 type MakeNonNullable<T> = {
   [K in keyof T]-?: NonNullable<T[K]>
diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index b9fe70d48..a8bacf7a1 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -1,12 +1,13 @@
 import { match, P } from "ts-pattern"
 
 import { PdlBlock } from "./pdl_ast"
-import { hasContextInformation, hasResult, isArgs } from "./helpers"
+import { BlockType, ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
-  f: (block: PdlBlock) => PdlBlock,
-  block: PdlBlock,
-): PdlBlock {
+  f_block: (block: BlockType) => BlockType,
+  f_expr: (expr: ExpressionT<any>) => ExpressionT<any>,
+  block: BlockType,
+): BlockType {
   if (
     block === null ||
     typeof block === "boolean" ||
@@ -21,118 +22,150 @@ export function map_block_children(
   } else {
     const defs: { [k: string]: PdlBlock } = {}
     for (const x in block.defs) {
-      defs[x] = f(block.defs[x])
+      defs[x] = f_block(block.defs[x])
     }
     new_block = { ...block, defs: defs }
   }
+  if (new_block?.contribute !== undefined) {
+    new_block.contribute = new_block.contribute?.map((contrib) =>
+      match(contrib)
+        .with({}, (c) =>
+          Object.fromEntries(
+            Object.entries(c).map(([k, v]) => [
+              k,
+              match(v)
+                .with({ value: P._ }, (v) => ({ value: f_expr(v) }))
+                .otherwise((v) => v),
+            ]),
+          ),
+        )
+        .otherwise((contrib) => contrib),
+    )
+  }
   new_block = match(new_block)
     // .with(P.string, s => s)
     .with({ kind: "empty" }, (block) => block)
     .with({ kind: "function" }, (block) => {
-      const returns = f(block.return)
-      return { ...block, return: returns }
+      const return_ = f_block(block.return)
+      return { ...block, return: return_ }
+    })
+    .with({ kind: "call" }, (block) => {
+      const call = f_expr(block.call)
+      const args = f_expr(block.args)
+      return { ...block, call, args }
     })
-    .with({ kind: "call" }, (block) => block)
     .with({ kind: "model" }, (block) => {
-      if (block.input) {
-        const input = f(block.input)
-        console.error("!!!!!", input)
-        block = { ...block, input }
-      }
-      if (
-        hasResult(block.model) &&
-        typeof block.model.pdl__result === "string"
-      ) {
-        block = { ...block, model: block.model.pdl__result }
-      }
-      // Remove `defsite` from context:
-      return {
-        ...block,
-        context: !hasContextInformation(block)
-          ? undefined
-          : JSON.parse(
-              JSON.stringify(block.context, (k, v) =>
-                k === "defsite" ? undefined : v,
-              ),
-            ),
-      }
+      const model = f_expr(block.model)
+      const input = block.input ? f_block(block.input) : undefined
+      const parameters = block.parameters ? f_expr(block.parameters) : undefined
+      return match(block)
+        .with({ platform: "litellm" }, (_) => ({
+          ...block,
+          model,
+          input,
+          parameters,
+        }))
+        .with(
+          {
+            platform: "granite-io",
+            backend: P.select("b"),
+            processor: P.select("p"),
+          },
+          ({ b, p }) => {
+            const backend = f_expr(b)
+            const processor = f_expr(p)
+            return { ...block, model, input, parameters, backend, processor }
+          },
+        )
+        .exhaustive()
     })
     .with({ kind: "code" }, (block) => {
       if (isArgs(block)) {
-        return block
-      } else {
-        return { ...block, code: f(block.code) }
+        const args = block.args.map((arg) => f_expr(arg))
+        return { ...block, args }
       }
+      return { ...block, code: f_block(block.code) }
     })
     .with({ kind: "get" }, (block) => block)
-    .with(
-      { kind: "data", data: P.union(P.number, P.boolean, P.string) },
-      ({ data }) => data, // { kind: "data", data: "hello" } -> "hello"
-    )
-    .with({ kind: "data" }, (block) => block)
+    .with({ kind: "data" }, (block) => {
+      const data = f_expr(block.data)
+      return { ...block, data }
+    })
     .with({ kind: "text" }, (block) => {
       let text
       if (block.text instanceof Array) {
-        text = block.text.map(f)
+        text = block.text.map(f_block)
       } else {
-        text = f(block.text)
+        text = f_block(block.text)
       }
       return { ...block, text: text }
     })
     .with({ kind: "lastOf" }, (block) => {
-      const lastOf = block.lastOf.map(f)
+      const lastOf = block.lastOf.map(f_block)
       return { ...block, lastOf: lastOf }
     })
     .with({ kind: "array" }, (block) => {
-      const array = block.array.map(f)
+      const array = block.array.map(f_block)
       return { ...block, array: array }
     })
     .with({ kind: "object" }, (block) => {
       let object
       if (block.object instanceof Array) {
-        object = block.object.map(f)
+        object = block.object.map(f_block)
       } else {
         object = Object.fromEntries(
-          Object.entries(block.object).map(([k, v]) => [k, f(v)]),
+          Object.entries(block.object).map(([k, v]) => [k, f_block(v)]),
         )
       }
       return { ...block, object: object }
     })
     .with({ kind: "message" }, (block) => {
-      const content = f(block.content)
+      const content = f_block(block.content)
       return { ...block, content: content }
     })
     .with({ kind: "if" }, (block) => {
-      const then_ = f(block.then)
-      const else_ = block.else ? f(block.else) : undefined
-      return { ...block, then: then_, else: else_ }
+      const if_ = f_expr(block.if)
+      const then_ = f_block(block.then)
+      const else_ = block.else ? f_block(block.else) : undefined
+      return { ...block, if: if_, then: then_, else: else_ }
     })
     .with({ kind: "match" }, (block) => {
+      const match = f_expr(block.match)
       const with_ = block.with.map((match_case) => {
-        return { ...match_case, then: f(match_case.then) }
+        const if_ = f_expr(match_case.if)
+        const then_ = f_block(match_case.then)
+        return { ...match_case, if: if_, then: then_ }
       })
-      return { ...block, with: with_ }
+      return { ...block, match, with: with_ }
     })
     .with({ kind: "repeat" }, (block) => {
-      const repeat = f(block.repeat)
-      return { ...block, repeat: repeat }
+      const for_ = block?.for ? f_expr(block.for) : undefined
+      const until = block?.until ? f_expr(block.until) : undefined
+      const max_iterations = block?.max_iterations
+        ? f_expr(block.max_iterations)
+        : undefined
+      const repeat = f_block(block.repeat)
+      return { ...block, for: for_, repeat, until, max_iterations }
     })
     .with({ kind: "error" }, (block) => {
-      const doc = f(block.program)
+      const doc = f_block(block.program)
       return { ...block, program: doc }
     })
-    .with({ kind: "read" }, (block) => block)
+    .with({ kind: "read" }, (block) => {
+      const read = f_expr(block.read)
+      return { ...block, read }
+    })
     .with({ kind: "include" }, (block) => block)
     .with({ kind: "import" }, (block) => block)
     .with({ kind: undefined }, (block) => block)
     .exhaustive()
   match(new_block)
     .with({ parser: { pdl: P._ } }, (block) => {
-      block.parser.pdl = f(block.parser.pdl)
+      block.parser.pdl = f_block(block.parser.pdl)
     })
     .otherwise(() => {})
   if (block.fallback) {
-    block.fallback = f(block.fallback)
+    block.fallback = f_block(block.fallback)
   }
   return new_block
 }
diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index 6e4833b99..0a45ac517 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -4,7 +4,11 @@ import { stringify } from "yaml"
 
 import { tryJsonPrettyPrint } from "../../helpers"
 import { type PdlBlock } from "../../pdl_ast"
-import { type BlockType } from "../../helpers"
+import {
+  type BlockType,
+  type ExpressionT,
+  hasContextInformation,
+} from "../../helpers"
 import { map_block_children } from "../../pdl_ast_utils"
 
 const PreviewLight = lazy(() => import("./PreviewLight"))
@@ -51,13 +55,13 @@ export default function Code({
   )
 }
 
-export function block_code_cleanup(data: BlockType): BlockType {
-  if (data === null || typeof data !== "object") {
-    return data
+export function block_code_cleanup(block: BlockType): BlockType {
+  if (block === null || typeof block !== "object") {
+    return block
   }
   // remove pdl__result
-  let new_data = {
-    ...data,
+  let new_block = {
+    ...block,
     pdl__result: undefined,
     pdl__is_leaf: undefined,
     pdl__usage: undefined,
@@ -69,19 +73,34 @@ export function block_code_cleanup(data: BlockType): BlockType {
   }
   // remove contribute: ["result", context]
   if (
-    new_data?.contribute?.includes("result") &&
-    new_data?.contribute?.includes("context")
+    new_block?.contribute?.includes("result") &&
+    new_block?.contribute?.includes("context")
   ) {
-    delete new_data.contribute
+    delete new_block.contribute
   }
   // remove empty defs list
-  if (Object.keys(data?.defs ?? {}).length === 0) {
-    delete new_data.defs
+  if (Object.keys(block?.defs ?? {}).length === 0) {
+    delete new_block.defs
   }
   // recursive cleanup
-  const new_data_rec = map_block_children(block_code_cleanup, new_data)
-  // replace `data: literal` by `literal`
-  const clean_data = match(new_data_rec)
+  const new_block_rec = map_block_children(
+    block_code_cleanup,
+    expr_code_cleanup,
+    new_block,
+  )
+  const clean_block = match(new_block_rec)
+    // Remove `defsite` from context:
+    .with({ kind: "model" }, (block) => ({
+      ...block,
+      context: !hasContextInformation(block)
+        ? undefined
+        : JSON.parse(
+            JSON.stringify(block.context, (k, v) =>
+              k === "defsite" ? undefined : v,
+            ),
+          ),
+    }))
+    // replace `data: literal` by `literal`
     .with(
       {
         kind: "data",
@@ -96,17 +115,50 @@ export function block_code_cleanup(data: BlockType): BlockType {
         fallback: P.nullish,
         role: P.nullish,
       },
-      (d) => d.data,
+      (block) => block.data,
     )
-    .otherwise((d) => d)
-  // remove kind
-  return match(clean_data)
-    .with({ kind: "data" }, (d) => {
-      return match(d.data)
-        .with({ pdl__expr: P._ }, (e) => ({ ...d, data: e.pdl__expr }))
-        .with(P.union(P.string, P.number, P.boolean, {}), (e) => e)
-        .otherwise((_) => d)
+    .with({ kind: "match", with: P._ }, (block) => {
+      const with_ = block.with.map((case_) => {
+        const clean_case = {
+          ...case_,
+          pdl__case_result: undefined,
+          pdl__if_result: undefined,
+          pdl__matched: undefined,
+        }
+        if (clean_case.case === null) {
+          delete clean_case.case
+        }
+        if (clean_case.if === null) {
+          delete clean_case.if
+        }
+        return clean_case
+      })
+      return { ...block, with: with_ }
+    })
+    .with({ kind: "repeat" }, (block) => {
+      if (block.for === null) {
+        delete block.for
+      }
+      if (block.while === true) {
+        delete block.while
+      }
+      if (block.until === false) {
+        delete block.until
+      }
+      if (block.max_iterations === null) {
+        delete block.max_iterations
+      }
+      return block
     })
-    .with({ kind: P._ }, (d) => ({ ...d, kind: undefined }))
-    .otherwise((d) => d)
+    .otherwise((block) => block)
+  // remove kind
+  return match(clean_block)
+    .with({ kind: P._ }, (block) => ({ ...block, kind: undefined }))
+    .otherwise((block) => block)
+}
+
+export function expr_code_cleanup(expr: ExpressionT<any>): ExpressionT<any> {
+  return match(expr)
+    .with({ pdl__expr: P._ }, (e) => e.pdl__expr)
+    .otherwise((e) => e)
 }

From 5fe592b3479983d7cdc8e7738176f68fe4f3de42 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Mon, 24 Mar 2025 23:26:09 -0400
Subject: [PATCH 07/14] pre-commit fixes

---
 pdl-live-react/src/pdl_ast_utils.ts    |   9 +-
 pdl-live-react/src/pdl_code_cleanup.ts | 116 ++++++++++++++++++++++++
 pdl-live-react/src/view/code/Code.tsx  | 117 +------------------------
 3 files changed, 123 insertions(+), 119 deletions(-)
 create mode 100644 pdl-live-react/src/pdl_code_cleanup.ts

diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index a8bacf7a1..69647fd38 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -5,7 +5,7 @@ import { BlockType, ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
   f_block: (block: BlockType) => BlockType,
-  f_expr: (expr: ExpressionT<any>) => ExpressionT<any>,
+  f_expr: (expr: ExpressionT<unknown>) => ExpressionT<unknown>,
   block: BlockType,
 ): BlockType {
   if (
@@ -27,20 +27,23 @@ export function map_block_children(
     new_block = { ...block, defs: defs }
   }
   if (new_block?.contribute !== undefined) {
-    new_block.contribute = new_block.contribute?.map((contrib) =>
+    const contribute = new_block.contribute?.map((contrib) =>
       match(contrib)
         .with({}, (c) =>
           Object.fromEntries(
             Object.entries(c).map(([k, v]) => [
               k,
               match(v)
-                .with({ value: P._ }, (v) => ({ value: f_expr(v) }))
+                .with({ value: P.array(P._) }, (v) => ({
+                  value: v.value.map(f_expr),
+                }))
                 .otherwise((v) => v),
             ]),
           ),
         )
         .otherwise((contrib) => contrib),
     )
+    new_block = { ...new_block, contribute }
   }
   new_block = match(new_block)
     // .with(P.string, s => s)
diff --git a/pdl-live-react/src/pdl_code_cleanup.ts b/pdl-live-react/src/pdl_code_cleanup.ts
new file mode 100644
index 000000000..832f8e973
--- /dev/null
+++ b/pdl-live-react/src/pdl_code_cleanup.ts
@@ -0,0 +1,116 @@
+import { match, P } from "ts-pattern"
+
+import {
+  type BlockType,
+  type ExpressionT,
+  hasContextInformation,
+} from "./helpers"
+import { map_block_children } from "./pdl_ast_utils"
+
+export function block_code_cleanup(block: BlockType): BlockType {
+  if (block === null || typeof block !== "object") {
+    return block
+  }
+  // remove pdl__result
+  const new_block = {
+    ...block,
+    pdl__result: undefined,
+    pdl__is_leaf: undefined,
+    pdl__usage: undefined,
+    pdl__trace: undefined,
+    pdl__id: undefined,
+    pdl__timing: undefined,
+    pdl__location: undefined,
+    pdl__model_input: undefined,
+  }
+  // remove contribute: ["result", context]
+  if (
+    new_block?.contribute?.includes("result") &&
+    new_block?.contribute?.includes("context")
+  ) {
+    delete new_block.contribute
+  }
+  // remove empty defs list
+  if (Object.keys(block?.defs ?? {}).length === 0) {
+    delete new_block.defs
+  }
+  // recursive cleanup
+  const new_block_rec = map_block_children(
+    block_code_cleanup,
+    expr_code_cleanup,
+    new_block,
+  )
+  const clean_block = match(new_block_rec)
+    // Remove `defsite` from context:
+    .with({ kind: "model" }, (block) => ({
+      ...block,
+      context: !hasContextInformation(block)
+        ? undefined
+        : JSON.parse(
+            JSON.stringify(block.context, (k, v) =>
+              k === "defsite" ? undefined : v,
+            ),
+          ),
+    }))
+    // replace `data: literal` by `literal`
+    .with(
+      {
+        kind: "data",
+        data: P.union(P.string, P.number, P.boolean),
+        raw: false,
+        spec: P.nullish,
+        description: P.nullish,
+        defs: {},
+        def: P.nullish,
+        contribute: P.union(["context", "result"], ["result", "context"]),
+        parser: P.nullish,
+        fallback: P.nullish,
+        role: P.nullish,
+      },
+      (block) => block.data,
+    )
+    .with({ kind: "match", with: P._ }, (block) => {
+      const with_ = block.with.map((case_) => {
+        const clean_case = {
+          ...case_,
+          pdl__case_result: undefined,
+          pdl__if_result: undefined,
+          pdl__matched: undefined,
+        }
+        if (clean_case.case === null) {
+          delete clean_case.case
+        }
+        if (clean_case.if === null) {
+          delete clean_case.if
+        }
+        return clean_case
+      })
+      return { ...block, with: with_ }
+    })
+    .with({ kind: "repeat" }, (block) => {
+      if (block.for === null) {
+        delete block.for
+      }
+      if (block.while === true) {
+        delete block.while
+      }
+      if (block.until === false) {
+        delete block.until
+      }
+      if (block.max_iterations === null) {
+        delete block.max_iterations
+      }
+      return block
+    })
+    .otherwise((block) => block)
+  // remove kind
+  return match(clean_block)
+    .with({ kind: P._ }, (block) => ({ ...block, kind: undefined }))
+    .otherwise((block) => block)
+}
+
+function expr_code_cleanup(expr: ExpressionT<unknown>): ExpressionT<unknown> {
+  return match(expr)
+    .with({ pdl__expr: P._ }, (e) => e.pdl__expr)
+    .otherwise((e) => e)
+}
diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index 0a45ac517..6433fbc33 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -1,16 +1,9 @@
-import { match, P } from "ts-pattern"
 import { lazy, Suspense } from "react"
 import { stringify } from "yaml"
 
 import { tryJsonPrettyPrint } from "../../helpers"
 import { type PdlBlock } from "../../pdl_ast"
-import {
-  type BlockType,
-  type ExpressionT,
-  hasContextInformation,
-} from "../../helpers"
-import { map_block_children } from "../../pdl_ast_utils"
-
+import { block_code_cleanup } from "../../pdl_code_cleanup"
 const PreviewLight = lazy(() => import("./PreviewLight"))
 
 export type SupportedLanguage =
@@ -54,111 +47,3 @@ export default function Code({
     </Suspense>
   )
 }
-
-export function block_code_cleanup(block: BlockType): BlockType {
-  if (block === null || typeof block !== "object") {
-    return block
-  }
-  // remove pdl__result
-  let new_block = {
-    ...block,
-    pdl__result: undefined,
-    pdl__is_leaf: undefined,
-    pdl__usage: undefined,
-    pdl__trace: undefined,
-    pdl__id: undefined,
-    pdl__timing: undefined,
-    pdl__location: undefined,
-    pdl__model_input: undefined,
-  }
-  // remove contribute: ["result", context]
-  if (
-    new_block?.contribute?.includes("result") &&
-    new_block?.contribute?.includes("context")
-  ) {
-    delete new_block.contribute
-  }
-  // remove empty defs list
-  if (Object.keys(block?.defs ?? {}).length === 0) {
-    delete new_block.defs
-  }
-  // recursive cleanup
-  const new_block_rec = map_block_children(
-    block_code_cleanup,
-    expr_code_cleanup,
-    new_block,
-  )
-  const clean_block = match(new_block_rec)
-    // Remove `defsite` from context:
-    .with({ kind: "model" }, (block) => ({
-      ...block,
-      context: !hasContextInformation(block)
-        ? undefined
-        : JSON.parse(
-            JSON.stringify(block.context, (k, v) =>
-              k === "defsite" ? undefined : v,
-            ),
-          ),
-    }))
-    // replace `data: literal` by `literal`
-    .with(
-      {
-        kind: "data",
-        data: P.union(P.string, P.number, P.boolean),
-        raw: false,
-        spec: P.nullish,
-        description: P.nullish,
-        defs: {},
-        def: P.nullish,
-        contribute: P.union(["context", "result"], ["result", "context"]),
-        parser: P.nullish,
-        fallback: P.nullish,
-        role: P.nullish,
-      },
-      (block) => block.data,
-    )
-    .with({ kind: "match", with: P._ }, (block) => {
-      const with_ = block.with.map((case_) => {
-        const clean_case = {
-          ...case_,
-          pdl__case_result: undefined,
-          pdl__if_result: undefined,
-          pdl__matched: undefined,
-        }
-        if (clean_case.case === null) {
-          delete clean_case.case
-        }
-        if (clean_case.if === null) {
-          delete clean_case.if
-        }
-        return clean_case
-      })
-      return { ...block, with: with_ }
-    })
-    .with({ kind: "repeat" }, (block) => {
-      if (block.for === null) {
-        delete block.for
-      }
-      if (block.while === true) {
-        delete block.while
-      }
-      if (block.until === false) {
-        delete block.until
-      }
-      if (block.max_iterations === null) {
-        delete block.max_iterations
-      }
-      return block
-    })
-    .otherwise((block) => block)
-  // remove kind
-  return match(clean_block)
-    .with({ kind: P._ }, (block) => ({ ...block, kind: undefined }))
-    .otherwise((block) => block)
-}
-
-export function expr_code_cleanup(expr: ExpressionT<any>): ExpressionT<any> {
-  return match(expr)
-    .with({ pdl__expr: P._ }, (e) => e.pdl__expr)
-    .otherwise((e) => e)
-}

From 5a56f5b683369179eb1e55b8299bd33b6e625fbb Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Mon, 24 Mar 2025 23:29:50 -0400
Subject: [PATCH 08/14] Fix build

---
 pdl-live-react/src/view/detail/kind/Function.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pdl-live-react/src/view/detail/kind/Function.tsx b/pdl-live-react/src/view/detail/kind/Function.tsx
index 3d12c5ff7..63034828f 100644
--- a/pdl-live-react/src/view/detail/kind/Function.tsx
+++ b/pdl-live-react/src/view/detail/kind/Function.tsx
@@ -2,7 +2,7 @@ import { stringify } from "yaml"
 
 import Group from "../Group"
 import CodeGroup from "../../code/CodeGroup"
-import { block_code_cleanup } from "../../code/Code"
+import { block_code_cleanup } from "../../../pdl_code_cleanup"
 
 export default function FunctionItems({
   block: { def, function: func, return: retrn },

From ed5d7352da22a91bb9b7ee73c94c35ae5b0de11a Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Tue, 25 Mar 2025 10:34:16 -0400
Subject: [PATCH 09/14] Fix / ignore typing issues

---
 pdl-live-react/src/pdl_ast_utils.ts | 53 +++++++++++++++++------------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index 69647fd38..1db88ea88 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -1,6 +1,6 @@
 import { match, P } from "ts-pattern"
 
-import { PdlBlock } from "./pdl_ast"
+import { Backend, PdlBlock, Processor } from "./pdl_ast"
 import { BlockType, ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
@@ -45,6 +45,7 @@ export function map_block_children(
     )
     new_block = { ...new_block, contribute }
   }
+  // @ts-ignore
   new_block = match(new_block)
     // .with(P.string, s => s)
     .with({ kind: "empty" }, (block) => block)
@@ -57,30 +58,40 @@ export function map_block_children(
       const args = f_expr(block.args)
       return { ...block, call, args }
     })
-    .with({ kind: "model" }, (block) => {
-      const model = f_expr(block.model)
-      const input = block.input ? f_block(block.input) : undefined
-      const parameters = block.parameters ? f_expr(block.parameters) : undefined
-      return match(block)
-        .with({ platform: "litellm" }, (_) => ({
+    .with(
+      { kind: "model", platform: "granite-io", backend: P.nonNullable, processor: P._ },
+      (block) => {
+        const model = f_expr(block.model)
+        const input = block.input ? f_block(block.input) : undefined
+        // @ts-ignore
+        const parameters: Parameters = block.parameters
+          ? f_expr(block.parameters)
+          : undefined
+        // @ts-ignore
+        const backend: Backend = f_expr(block.backend)
+        // @ts-ignore
+        const processor: Processor = block.processor ? f_expr(block.processor) : undefined
+        return {
           ...block,
           model,
           input,
           parameters,
-        }))
-        .with(
-          {
-            platform: "granite-io",
-            backend: P.select("b"),
-            processor: P.select("p"),
-          },
-          ({ b, p }) => {
-            const backend = f_expr(b)
-            const processor = f_expr(p)
-            return { ...block, model, input, parameters, backend, processor }
-          },
-        )
-        .exhaustive()
+          backend,
+          processor,
+        }
+      },
+    )
+    .with({ kind: "model" }, (block) => {
+      const model = f_expr(block.model)
+      const input = block.input ? f_block(block.input) : undefined
+      const parameters = block.parameters ? f_expr(block.parameters) : undefined
+      return {
+        ...block,
+        platform: "litellm",
+        model,
+        input,
+        parameters,
+      }
     })
     .with({ kind: "code" }, (block) => {
       if (isArgs(block)) {

From df9f3a23496243f3583fe5b52ba8fd2a9f4c68c4 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Tue, 25 Mar 2025 12:00:15 -0400
Subject: [PATCH 10/14] Use @ts-expect-error

---
 pdl-live-react/src/pdl_ast_utils.ts | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index 1db88ea88..03dff4bdb 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -45,7 +45,7 @@ export function map_block_children(
     )
     new_block = { ...new_block, contribute }
   }
-  // @ts-ignore
+  // @ts-expect-error: TODO
   new_block = match(new_block)
     // .with(P.string, s => s)
     .with({ kind: "empty" }, (block) => block)
@@ -63,13 +63,13 @@ export function map_block_children(
       (block) => {
         const model = f_expr(block.model)
         const input = block.input ? f_block(block.input) : undefined
-        // @ts-ignore
+        // @ts-expect-error: f_expr does not preserve the type of the expression
         const parameters: Parameters = block.parameters
           ? f_expr(block.parameters)
           : undefined
-        // @ts-ignore
+        // @ts-expect-error: f_expr does not preserve the type of the expression
         const backend: Backend = f_expr(block.backend)
-        // @ts-ignore
+        // @ts-expect-error: f_expr does not preserve the type of the expression
         const processor: Processor = block.processor ? f_expr(block.processor) : undefined
         return {
           ...block,

From 6e3a26072bfc120bc6d71f77366ca9d5d7409a3c Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Tue, 25 Mar 2025 14:52:30 -0400
Subject: [PATCH 11/14] Formatting

---
 pdl-live-react/src/pdl_ast_utils.ts | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index 03dff4bdb..4cb492405 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -59,7 +59,12 @@ export function map_block_children(
       return { ...block, call, args }
     })
     .with(
-      { kind: "model", platform: "granite-io", backend: P.nonNullable, processor: P._ },
+      {
+        kind: "model",
+        platform: "granite-io",
+        backend: P.nonNullable,
+        processor: P._,
+      },
       (block) => {
         const model = f_expr(block.model)
         const input = block.input ? f_block(block.input) : undefined
@@ -70,7 +75,9 @@ export function map_block_children(
         // @ts-expect-error: f_expr does not preserve the type of the expression
         const backend: Backend = f_expr(block.backend)
         // @ts-expect-error: f_expr does not preserve the type of the expression
-        const processor: Processor = block.processor ? f_expr(block.processor) : undefined
+        const processor: Processor = block.processor
+          ? f_expr(block.processor)
+          : undefined
         return {
           ...block,
           model,

From 595443bc4ee47ea1b588e6aca6d2aad4663d6ba7 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Tue, 25 Mar 2025 22:40:06 -0400
Subject: [PATCH 12/14] Refactoring

---
 pdl-live-react/src/helpers.ts          |   6 +-
 pdl-live-react/src/pdl_ast_utils.ts    |   8 +-
 pdl-live-react/src/pdl_code_cleanup.ts | 204 +++++++++++++++----------
 3 files changed, 130 insertions(+), 88 deletions(-)

diff --git a/pdl-live-react/src/helpers.ts b/pdl-live-react/src/helpers.ts
index 196f77a15..3edec1d48 100644
--- a/pdl-live-react/src/helpers.ts
+++ b/pdl-live-react/src/helpers.ts
@@ -12,8 +12,10 @@ import type {
 
 /** Re-export for convenience */
 export type { PdlBlock } from "./pdl_ast"
-
-export type BlockType = null | boolean | number | string | PdlBlock
+export type AdvancedBlockType = Exclude<
+  PdlBlock,
+  null | string | boolean | number
+>
 export type ExpressionT<T> = T | string | LocalizedExpression
 
 type MakeNonNullable<T> = {
diff --git a/pdl-live-react/src/pdl_ast_utils.ts b/pdl-live-react/src/pdl_ast_utils.ts
index 4cb492405..2edfe2d5f 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -1,13 +1,13 @@
 import { match, P } from "ts-pattern"
 
 import { Backend, PdlBlock, Processor } from "./pdl_ast"
-import { BlockType, ExpressionT, isArgs } from "./helpers"
+import { ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
-  f_block: (block: BlockType) => BlockType,
+  f_block: (block: PdlBlock) => PdlBlock,
   f_expr: (expr: ExpressionT<unknown>) => ExpressionT<unknown>,
-  block: BlockType,
-): BlockType {
+  block: PdlBlock,
+): PdlBlock {
   if (
     block === null ||
     typeof block === "boolean" ||
diff --git a/pdl-live-react/src/pdl_code_cleanup.ts b/pdl-live-react/src/pdl_code_cleanup.ts
index 832f8e973..33da8a3a0 100644
--- a/pdl-live-react/src/pdl_code_cleanup.ts
+++ b/pdl-live-react/src/pdl_code_cleanup.ts
@@ -1,58 +1,65 @@
 import { match, P } from "ts-pattern"
 
 import {
-  type BlockType,
-  type ExpressionT,
+  AdvancedBlockType,
+  ExpressionT,
   hasContextInformation,
 } from "./helpers"
 import { map_block_children } from "./pdl_ast_utils"
+import {
+  DataBlock,
+  GraniteioModelBlock,
+  LitellmModelBlock,
+  MatchBlock,
+  PdlBlock,
+  RepeatBlock,
+} from "./pdl_ast"
 
-export function block_code_cleanup(block: BlockType): BlockType {
-  if (block === null || typeof block !== "object") {
-    return block
-  }
-  // remove pdl__result
-  const new_block = {
-    ...block,
-    pdl__result: undefined,
-    pdl__is_leaf: undefined,
-    pdl__usage: undefined,
-    pdl__trace: undefined,
-    pdl__id: undefined,
-    pdl__timing: undefined,
-    pdl__location: undefined,
-    pdl__model_input: undefined,
-  }
-  // remove contribute: ["result", context]
+export function block_code_cleanup(block: PdlBlock): PdlBlock {
+  const block_with_clean_children = map_block_children(
+    block_code_cleanup,
+    expr_code_cleanup,
+    block,
+  )
   if (
-    new_block?.contribute?.includes("result") &&
-    new_block?.contribute?.includes("context")
+    block_with_clean_children === null ||
+    typeof block_with_clean_children !== "object"
   ) {
-    delete new_block.contribute
+    return block_with_clean_children
   }
-  // remove empty defs list
-  if (Object.keys(block?.defs ?? {}).length === 0) {
-    delete new_block.defs
-  }
-  // recursive cleanup
-  const new_block_rec = map_block_children(
-    block_code_cleanup,
-    expr_code_cleanup,
-    new_block,
+  let block_with_generic_clean = remove_block_default_values(
+    block_with_clean_children,
   )
-  const clean_block = match(new_block_rec)
-    // Remove `defsite` from context:
-    .with({ kind: "model" }, (block) => ({
-      ...block,
-      context: !hasContextInformation(block)
-        ? undefined
-        : JSON.parse(
-            JSON.stringify(block.context, (k, v) =>
-              k === "defsite" ? undefined : v,
-            ),
+  block_with_generic_clean = remove_internal_block_fields(
+    block_with_generic_clean,
+  )
+  const clean_block = match(block_with_generic_clean)
+    .with({ kind: "model" }, clean_model_block)
+    .with({ kind: "data" }, clean_data_block)
+    .with({ kind: "match", with: P._ }, clean_match_block)
+    .with({ kind: "repeat" }, clean_repeat_block)
+    .otherwise((block) => block)
+  // remove kind
+  return match(clean_block)
+    .with({ kind: P._ }, (block) => ({ ...block, kind: undefined }))
+    .otherwise((block) => block)
+}
+
+function clean_model_block(block: LitellmModelBlock | GraniteioModelBlock) {
+  return {
+    ...block,
+    context: !hasContextInformation(block)
+      ? undefined
+      : JSON.parse(
+          JSON.stringify(block.context, (k, v) =>
+            k === "defsite" ? undefined : v,
           ),
-    }))
-    // replace `data: literal` by `literal`
+        ),
+  }
+}
+
+function clean_data_block(block: DataBlock) {
+  return match(block)
     .with(
       {
         kind: "data",
@@ -60,55 +67,88 @@ export function block_code_cleanup(block: BlockType): BlockType {
         raw: false,
         spec: P.nullish,
         description: P.nullish,
-        defs: {},
+        defs: P.union(undefined, {}),
         def: P.nullish,
-        contribute: P.union(["context", "result"], ["result", "context"]),
+        contribute: P.union(
+          ["context", "result"],
+          ["result", "context"],
+          undefined,
+        ),
         parser: P.nullish,
         fallback: P.nullish,
         role: P.nullish,
       },
       (block) => block.data,
     )
-    .with({ kind: "match", with: P._ }, (block) => {
-      const with_ = block.with.map((case_) => {
-        const clean_case = {
-          ...case_,
-          pdl__case_result: undefined,
-          pdl__if_result: undefined,
-          pdl__matched: undefined,
-        }
-        if (clean_case.case === null) {
-          delete clean_case.case
-        }
-        if (clean_case.if === null) {
-          delete clean_case.if
-        }
-        return clean_case
-      })
-      return { ...block, with: with_ }
-    })
-    .with({ kind: "repeat" }, (block) => {
-      if (block.for === null) {
-        delete block.for
-      }
-      if (block.while === true) {
-        delete block.while
-      }
-      if (block.until === false) {
-        delete block.until
-      }
-      if (block.max_iterations === null) {
-        delete block.max_iterations
-      }
-      return block
-    })
-    .otherwise((block) => block)
-  // remove kind
-  return match(clean_block)
-    .with({ kind: P._ }, (block) => ({ ...block, kind: undefined }))
     .otherwise((block) => block)
 }
 
+function clean_match_block(block: MatchBlock) {
+  const with_ = block.with.map((case_) => {
+    const clean_case = {
+      ...case_,
+      pdl__case_result: undefined,
+      pdl__if_result: undefined,
+      pdl__matched: undefined,
+    }
+    if (clean_case.case === null) {
+      delete clean_case.case
+    }
+    if (clean_case.if === null) {
+      delete clean_case.if
+    }
+    return clean_case
+  })
+  return { ...block, with: with_ }
+}
+
+function remove_block_default_values(
+  block: AdvancedBlockType,
+): AdvancedBlockType {
+  // remove contribute: ["result", context]
+  if (
+    block?.contribute?.includes("result") &&
+    block?.contribute?.includes("context")
+  ) {
+    block = { ...block, contribute: undefined }
+  }
+  // remove empty defs list
+  if (Object.keys(block?.defs ?? {}).length === 0) {
+    block = { ...block, defs: undefined }
+  }
+  return block
+}
+
+function clean_repeat_block(block: RepeatBlock) {
+  if (block.for === null) {
+    delete block.for
+  }
+  if (block.while === true) {
+    delete block.while
+  }
+  if (block.until === false) {
+    delete block.until
+  }
+  if (block.max_iterations === null) {
+    delete block.max_iterations
+  }
+  return block
+}
+
+function remove_internal_block_fields(block: AdvancedBlockType) {
+  return {
+    ...block,
+    pdl__result: undefined,
+    pdl__is_leaf: undefined,
+    pdl__usage: undefined,
+    pdl__trace: undefined,
+    pdl__id: undefined,
+    pdl__timing: undefined,
+    pdl__location: undefined,
+    pdl__model_input: undefined,
+  }
+}
+
 function expr_code_cleanup(expr: ExpressionT<unknown>): ExpressionT<unknown> {
   return match(expr)
     .with({ pdl__expr: P._ }, (e) => e.pdl__expr)

From c75fb607564f4c651fd545c0433edafdd0aa7139 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Tue, 25 Mar 2025 23:09:01 -0400
Subject: [PATCH 13/14] Fix data cleanup

---
 pdl-live-react/src/pdl_code_cleanup.ts | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/pdl-live-react/src/pdl_code_cleanup.ts b/pdl-live-react/src/pdl_code_cleanup.ts
index 33da8a3a0..ac11bb450 100644
--- a/pdl-live-react/src/pdl_code_cleanup.ts
+++ b/pdl-live-react/src/pdl_code_cleanup.ts
@@ -64,19 +64,19 @@ function clean_data_block(block: DataBlock) {
       {
         kind: "data",
         data: P.union(P.string, P.number, P.boolean),
-        raw: false,
-        spec: P.nullish,
-        description: P.nullish,
-        defs: P.union(undefined, {}),
-        def: P.nullish,
-        contribute: P.union(
-          ["context", "result"],
-          ["result", "context"],
-          undefined,
+        raw: P.optional(false),
+        spec: P.optional(P.nullish),
+        description: P.optional(P.union(P.nullish, "")),
+        defs: P.optional(
+          P.when((defs) => Object.keys(defs ?? {}).length === 0),
         ),
-        parser: P.nullish,
-        fallback: P.nullish,
-        role: P.nullish,
+        def: P.optional(P.nullish),
+        contribute: P.optional(
+          P.union(["context", "result"], ["result", "context"]),
+        ),
+        parser: P.optional(P.nullish),
+        fallback: P.optional(P.nullish),
+        role: P.optional(P.nullish),
       },
       (block) => block.data,
     )

From b879700a354df59b41291145b1b5245940da1812 Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Wed, 26 Mar 2025 09:10:59 -0400
Subject: [PATCH 14/14] Consolidate types

---
 pdl-live-react/src/helpers.ts          | 4 ----
 pdl-live-react/src/pdl_code_cleanup.ts | 8 ++++----
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/pdl-live-react/src/helpers.ts b/pdl-live-react/src/helpers.ts
index 3edec1d48..110b23ffa 100644
--- a/pdl-live-react/src/helpers.ts
+++ b/pdl-live-react/src/helpers.ts
@@ -12,10 +12,6 @@ import type {
 
 /** Re-export for convenience */
 export type { PdlBlock } from "./pdl_ast"
-export type AdvancedBlockType = Exclude<
-  PdlBlock,
-  null | string | boolean | number
->
 export type ExpressionT<T> = T | string | LocalizedExpression
 
 type MakeNonNullable<T> = {
diff --git a/pdl-live-react/src/pdl_code_cleanup.ts b/pdl-live-react/src/pdl_code_cleanup.ts
index ac11bb450..05d5bf6be 100644
--- a/pdl-live-react/src/pdl_code_cleanup.ts
+++ b/pdl-live-react/src/pdl_code_cleanup.ts
@@ -1,7 +1,7 @@
 import { match, P } from "ts-pattern"
 
 import {
-  AdvancedBlockType,
+  NonScalarPdlBlock,
   ExpressionT,
   hasContextInformation,
 } from "./helpers"
@@ -103,8 +103,8 @@ function clean_match_block(block: MatchBlock) {
 }
 
 function remove_block_default_values(
-  block: AdvancedBlockType,
-): AdvancedBlockType {
+  block: NonScalarPdlBlock,
+): NonScalarPdlBlock {
   // remove contribute: ["result", context]
   if (
     block?.contribute?.includes("result") &&
@@ -135,7 +135,7 @@ function clean_repeat_block(block: RepeatBlock) {
   return block
 }
 
-function remove_internal_block_fields(block: AdvancedBlockType) {
+function remove_internal_block_fields(block: NonScalarPdlBlock) {
   return {
     ...block,
     pdl__result: undefined,