diff --git a/pdl-live-react/src/helpers.ts b/pdl-live-react/src/helpers.ts
index 7aeeb7114..110b23ffa 100644
--- a/pdl-live-react/src/helpers.ts
+++ b/pdl-live-react/src/helpers.ts
@@ -12,6 +12,7 @@ import type {
 
 /** Re-export for convenience */
 export type { PdlBlock } from "./pdl_ast"
+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..2edfe2d5f 100644
--- a/pdl-live-react/src/pdl_ast_utils.ts
+++ b/pdl-live-react/src/pdl_ast_utils.ts
@@ -1,10 +1,11 @@
 import { match, P } from "ts-pattern"
 
-import { PdlBlock } from "./pdl_ast"
-import { hasContextInformation, hasResult, isArgs } from "./helpers"
+import { Backend, PdlBlock, Processor } from "./pdl_ast"
+import { ExpressionT, isArgs } from "./helpers"
 
 export function map_block_children(
-  f: (block: PdlBlock) => PdlBlock,
+  f_block: (block: PdlBlock) => PdlBlock,
+  f_expr: (expr: ExpressionT<unknown>) => ExpressionT<unknown>,
   block: PdlBlock,
 ): PdlBlock {
   if (
@@ -21,118 +22,171 @@ 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) {
+    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.array(P._) }, (v) => ({
+                  value: v.value.map(f_expr),
+                }))
+                .otherwise((v) => v),
+            ]),
+          ),
+        )
+        .otherwise((contrib) => contrib),
+    )
+    new_block = { ...new_block, contribute }
+  }
+  // @ts-expect-error: TODO
   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",
+        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-expect-error: f_expr does not preserve the type of the expression
+        const parameters: Parameters = block.parameters
+          ? f_expr(block.parameters)
+          : undefined
+        // @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
+        return {
+          ...block,
+          model,
+          input,
+          parameters,
+          backend,
+          processor,
+        }
+      },
+    )
     .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:
+      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,
-        context: !hasContextInformation(block)
-          ? undefined
-          : JSON.parse(
-              JSON.stringify(block.context, (k, v) =>
-                k === "defsite" ? undefined : v,
-              ),
-            ),
+        platform: "litellm",
+        model,
+        input,
+        parameters,
       }
     })
     .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/pdl_code_cleanup.ts b/pdl-live-react/src/pdl_code_cleanup.ts
new file mode 100644
index 000000000..05d5bf6be
--- /dev/null
+++ b/pdl-live-react/src/pdl_code_cleanup.ts
@@ -0,0 +1,156 @@
+import { match, P } from "ts-pattern"
+
+import {
+  NonScalarPdlBlock,
+  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: PdlBlock): PdlBlock {
+  const block_with_clean_children = map_block_children(
+    block_code_cleanup,
+    expr_code_cleanup,
+    block,
+  )
+  if (
+    block_with_clean_children === null ||
+    typeof block_with_clean_children !== "object"
+  ) {
+    return block_with_clean_children
+  }
+  let block_with_generic_clean = remove_block_default_values(
+    block_with_clean_children,
+  )
+  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,
+          ),
+        ),
+  }
+}
+
+function clean_data_block(block: DataBlock) {
+  return match(block)
+    .with(
+      {
+        kind: "data",
+        data: P.union(P.string, P.number, P.boolean),
+        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),
+        ),
+        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,
+    )
+    .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: NonScalarPdlBlock,
+): NonScalarPdlBlock {
+  // 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: NonScalarPdlBlock) {
+  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)
+    .otherwise((e) => e)
+}
diff --git a/pdl-live-react/src/view/code/Code.tsx b/pdl-live-react/src/view/code/Code.tsx
index b19e3d643..6433fbc33 100644
--- a/pdl-live-react/src/view/code/Code.tsx
+++ b/pdl-live-react/src/view/code/Code.tsx
@@ -3,8 +3,7 @@ import { stringify } from "yaml"
 
 import { tryJsonPrettyPrint } from "../../helpers"
 import { type PdlBlock } from "../../pdl_ast"
-import { map_block_children } from "../../pdl_ast_utils"
-
+import { block_code_cleanup } from "../../pdl_code_cleanup"
 const PreviewLight = lazy(() => import("./PreviewLight"))
 
 export type SupportedLanguage =
@@ -48,34 +47,3 @@ export default function Code({
     </Suspense>
   )
 }
-
-function block_code_cleanup(data: string | PdlBlock): string | PdlBlock {
-  if (data === null || typeof data !== "object") {
-    return data
-  }
-  // remove pdl__result
-  const new_data = {
-    ...data,
-    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_data?.contribute?.includes("result") &&
-    new_data?.contribute?.includes("context")
-  ) {
-    delete new_data.contribute
-  }
-  // remove empty defs list
-  if (Object.keys(data?.defs ?? {}).length === 0) {
-    delete new_data.defs
-  }
-  // recursive cleanup
-  return map_block_children(block_code_cleanup, new_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..63034828f 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 "../../../pdl_code_cleanup"
 
 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" />
     </>
   )
 }