refactor: change granite-io block syntax (#975)

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

examples/demos/granite_io_hallucinations.pdl

@@ -1,4 +1,4 @@
-description: GraniteIO halluncination example
+description: GraniteIO hallucination example
 defs:
   doc:
     data:
@@ -40,13 +40,16 @@ defs:
 
 text:
 - Did Faith Hill take a break from recording after releasing her second album, It Matters to Me?
-- model: "granite3.2:2b"
-  backend: openai
+- processor:
+    model: "granite3.2:2b"
+    backend: openai
   parameters:
     documents:
     - ${ doc }
     controls:
       hallucinations: true
+    generate_inputs:
+      temperature: 0.0
   modelResponse: output
 - "\nHallucinations:\n"
 - for:  

examples/granite-io/granite_io_hallucinations.pdl

@@ -1,3 +1,4 @@
+description: GraniteIO hallucination example
 defs:
   doc:
     data:
@@ -39,14 +40,16 @@ defs:
 
 text:
 - Did Faith Hill take a break from recording after releasing her second album, It Matters to Me?
-- model: "granite3.2:2b"
-  backend: openai
+- processor:
+    model: "granite3.2:2b"
+    backend: openai
   parameters:
     documents:
     - ${ doc }
     controls:
       hallucinations: true
-      citations: true
+    generate_inputs:
+      temperature: 0.0
   modelResponse: output
 - "\nHallucinations:\n"
 - for:  
@@ -57,12 +60,3 @@ text:
     - "\nSentence: ${ hallucination.response_text }"
   join: 
     with: "\n"
-- "\n\nCitations:\n"
-- for:
-    citation: ${ output.results[0].next_message.citations }
-  repeat:
-    text: 
-    - "Citation: ${ citation.context_text }"
-    - "\nSentence: ${ citation.response_text }"
-  join:
-    with: "\n"

examples/granite-io/granite_io_object.pdl

@@ -0,0 +1,14 @@
+defs:
+  io_proc:
+    lang: python
+    code: |
+      from granite_io import make_backend, make_io_processor
+      model_name = "granite3.2:2b"
+      backend = make_backend("openai", { "model_name": model_name })
+      result = make_io_processor(model_name, backend=backend)
+text:
+- "Hello!\n"
+- processor: ${ io_proc }
+  parameters:
+    generate_inputs:
+      temperature: 0.0

examples/granite-io/granite_io_openai.pdl

@@ -1,4 +1,8 @@
 text:
 - "Hello!\n"
-- model: "granite3.2:2b"
-  backend: openai
+- processor:
+    model: "granite3.2:2b"
+    backend: openai
+  parameters:
+    generate_inputs:
+      temperature: 0.0

examples/granite-io/granite_io_thinking.pdl

@@ -1,13 +1,18 @@
 text:
-- |
-  Find the fastest way for a seller to visit all the cities in their region
-  >> Response:
-- model: "granite3.2:2b"
-  backend: openai
+- Find the fastest way for a seller to visit all the cities in their region
+- def: response
+  processor:
+    model: "granite3.2:2b"
+    backend: openai
   parameters: 
     thinking: true
+    generate_inputs:
+      temperature: 0.0
   modelResponse: outputs
-- "\n"
-- | 
-  >> Thoughts:
-  ${ outputs.results[0].next_message.reasoning_content }
+  contribute: []
+- |
+
+    >> Thoughts:
+    ${ outputs.results[0].next_message.reasoning_content }
+    >> Response:
+    ${ response }

examples/granite-io/granite_io_transformers.pdl

@@ -1,6 +1,7 @@
 text:
 - "Hello!\n"
-- model: ibm-granite/granite-3.2-2b-instruct
-  backend:
-    transformers: cpu
-  processor: granite3.2
+- processor:
+    type: granite3.2
+    model: ibm-granite/granite-3.2-2b-instruct
+    backend:
+      transformers: cpu

pdl-live-react/src/helpers.ts

@@ -9,6 +9,7 @@ import type {
   PdlModelInput,
   LocalizedExpression,
 } from "./pdl_ast"
+import { match, P } from "ts-pattern"
 
 /** Re-export for convenience */
 export type { PdlBlock } from "./pdl_ast"
@@ -325,10 +326,27 @@ function isExpr(e: unknown): e is LocalizedExpression {
   )
 }
 
-export function extractModel({ model }: ModelBlock): string {
-  return typeof model === "string"
-    ? model
-    : isExpr(model)
-      ? String(model.pdl__result)
-      : "unknown"
+function getValue<T>(expr: ExpressionT<T>) {
+  if (isExpr(expr)) {
+    if (expr.pdl__result === undefined) {
+      return expr.pdl__expr
+    } else {
+      return expr.pdl__result
+    }
+  } else {
+    return expr
+  }
+}
+
+export function extractModel(block: ModelBlock): string {
+  return match(block)
+    .with({ model: P._ }, (block) => stringify(getValue(block.model)))
+    .with({ processor: { model: P._ } }, (block) =>
+      stringify(getValue(block.processor.model)),
+    )
+    .with({ processor: { type: P._ } }, (block) =>
+      stringify(getValue(block.processor.type)),
+    )
+    .with({ processor: P._ }, (block) => stringify(getValue(block.processor)))
+    .exhaustive()
 }

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

@@ -2354,11 +2354,6 @@ export type PdlContext17 =
 export type PdlId17 = string | null
 export type PdlIsLeaf17 = true
 export type Kind17 = "model"
-/**
- * Model name used by the backend.
- *
- */
-export type Model = LocalizedExpression | string
 /**
  * Messages to send to the model.
  *
@@ -2502,11 +2497,6 @@ export type PdlContext18 =
 export type PdlId18 = string | null
 export type PdlIsLeaf18 = true
 export type Kind18 = "model"
-/**
- * Name of the model following the LiteLLM convention.
- *
- */
-export type Model1 = LocalizedExpression | string
 /**
  * Messages to send to the model.
  *
@@ -2552,6 +2542,11 @@ export type PdlModelInput1 =
  *
  */
 export type Platform1 = "litellm"
+/**
+ * Name of the model following the LiteLLM convention.
+ *
+ */
+export type Model1 = LocalizedExpression | string
 /**
  * Parameters to send to the model.
  *
@@ -3113,7 +3108,6 @@ export interface LitellmModelBlock {
   pdl__timing?: PdlTiming | null
   pdl__is_leaf?: PdlIsLeaf18
   kind?: Kind18
-  model: Model1
   input?: Input1
   modelResponse?: Modelresponse1
   /**
@@ -3123,6 +3117,7 @@ export interface LitellmModelBlock {
   pdl__usage?: PdlUsage | null
   pdl__model_input?: PdlModelInput1
   platform?: Platform1
+  model: Model1
   parameters?: Parameters1
 }
 /**
@@ -3200,7 +3195,6 @@ export interface GraniteioModelBlock {
   pdl__timing?: PdlTiming | null
   pdl__is_leaf?: PdlIsLeaf17
   kind?: Kind17
-  model: Model
   input?: Input
   modelResponse?: Modelresponse
   /**
@@ -3210,8 +3204,7 @@ export interface GraniteioModelBlock {
   pdl__usage?: PdlUsage | null
   pdl__model_input?: PdlModelInput
   platform?: Platform
-  backend: unknown
-  processor?: unknown
+  processor: unknown
   parameters?: Parameters
 }
 /**

pdl-live-react/src/pdl_ast_utils.ts

@@ -62,29 +62,42 @@ export function map_block_children(
       {
         kind: "model",
         platform: "granite-io",
-        backend: P.nonNullable,
         processor: P._,
       },
       (block) => {
-        const model = f_expr(block.model)
+        const processor = match(block.processor)
+          .with(
+            {
+              type: P.union(P._, P.nullish),
+              model: P.union(P._, P.nullish),
+              backend: P.nonNullable,
+            },
+            (proc) => {
+              const type = proc.type ? f_expr(proc.type) : undefined
+              const model = proc.model ? f_expr(proc.model) : undefined
+              const backend = f_expr(proc.backend)
+              return {
+                ...proc,
+                type,
+                model,
+                backend,
+              }
+            },
+          )
+          .otherwise((proc) => f_expr(proc))
         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
+        const parameters = block.parameters
           ? f_expr(block.parameters)
           : undefined
-        const backend = f_expr(block.backend)
-        const processor = block.processor ? f_expr(block.processor) : undefined
         return {
           ...block,
-          model,
+          processor,
           input,
           parameters,
-          backend,
-          processor,
         }
       },
     )
-    .with({ kind: "model" }, (block) => {
+    .with({ kind: "model", model: P._ }, (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
@@ -96,6 +109,15 @@ export function map_block_children(
         parameters,
       }
     })
+    .with({ kind: "model" }, (block) => {
+      const input = block.input ? f_block(block.input) : undefined
+      const parameters = block.parameters ? f_expr(block.parameters) : undefined
+      return {
+        ...block,
+        input,
+        parameters,
+      }
+    })
     .with({ kind: "code" }, (block) => {
       if (isArgs(block)) {
         const args = block.args.map((arg) => f_expr(arg))