Remove pdl__trace from model block (#783)

Author: mandel

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

@@ -2180,32 +2180,6 @@ export type Input =
 export type Modelresponse = string | null
 export type CompletionTokens = number | null
 export type PromptTokens = number | null
-export type PdlTrace3 =
-  | boolean
-  | number
-  | string
-  | FunctionBlock
-  | CallBlock
-  | LitellmModelBlock
-  | GraniteioModelBlock
-  | CodeBlock
-  | ArgsBlock
-  | GetBlock
-  | DataBlock
-  | IfBlock
-  | MatchBlock
-  | RepeatBlock
-  | TextBlock
-  | LastOfBlock
-  | ArrayBlock
-  | ObjectBlock
-  | MessageBlock
-  | ReadBlock
-  | IncludeBlock
-  | ImportBlock
-  | ErrorBlock
-  | EmptyBlock
-  | null
 /**
  * Optional field to ensure that the block is using granite-io.
  *
@@ -2352,32 +2326,6 @@ export type Input1 =
  *
  */
 export type Modelresponse1 = string | null
-export type PdlTrace4 =
-  | boolean
-  | number
-  | string
-  | FunctionBlock
-  | CallBlock
-  | LitellmModelBlock
-  | GraniteioModelBlock
-  | CodeBlock
-  | ArgsBlock
-  | GetBlock
-  | DataBlock
-  | IfBlock
-  | MatchBlock
-  | RepeatBlock
-  | TextBlock
-  | LastOfBlock
-  | ArrayBlock
-  | ObjectBlock
-  | MessageBlock
-  | ReadBlock
-  | IncludeBlock
-  | ImportBlock
-  | ErrorBlock
-  | EmptyBlock
-  | null
 /**
  * Optional field to ensure that the block is using LiteLLM.
  *
@@ -2517,7 +2465,7 @@ export type Context19 =
 export type PdlId19 = string | null
 export type PdlIsLeaf19 = true
 export type Kind19 = "call"
-export type PdlTrace5 =
+export type PdlTrace3 =
   | boolean
   | number
   | string
@@ -2764,7 +2712,7 @@ export interface CallBlock {
   kind?: Kind19
   call: unknown
   args?: unknown
-  pdl__trace?: PdlTrace5
+  pdl__trace?: PdlTrace3
 }
 /**
  * Type specification of the result of the block.
@@ -2839,7 +2787,6 @@ export interface LitellmModelBlock {
    *
    */
   pdl__usage?: PdlUsage | null
-  pdl__trace?: PdlTrace4
   platform?: Platform1
   parameters?: Parameters1
 }
@@ -2909,7 +2856,6 @@ export interface GraniteioModelBlock {
    *
    */
   pdl__usage?: PdlUsage | null
-  pdl__trace?: PdlTrace3
   platform?: Platform
   backend: Backend
   processor?: Processor

src/pdl/pdl-schema.json

@@ -4236,90 +4236,6 @@
           "default": null,
           "description": "Tokens consumed during model call\n    "
         },
-        "pdl__trace": {
-          "anyOf": [
-            {
-              "type": "boolean"
-            },
-            {
-              "type": "integer"
-            },
-            {
-              "type": "number"
-            },
-            {
-              "type": "string"
-            },
-            {
-              "$ref": "#/$defs/FunctionBlock"
-            },
-            {
-              "$ref": "#/$defs/CallBlock"
-            },
-            {
-              "$ref": "#/$defs/LitellmModelBlock"
-            },
-            {
-              "$ref": "#/$defs/GraniteioModelBlock"
-            },
-            {
-              "$ref": "#/$defs/CodeBlock"
-            },
-            {
-              "$ref": "#/$defs/ArgsBlock"
-            },
-            {
-              "$ref": "#/$defs/GetBlock"
-            },
-            {
-              "$ref": "#/$defs/DataBlock"
-            },
-            {
-              "$ref": "#/$defs/IfBlock"
-            },
-            {
-              "$ref": "#/$defs/MatchBlock"
-            },
-            {
-              "$ref": "#/$defs/RepeatBlock"
-            },
-            {
-              "$ref": "#/$defs/TextBlock"
-            },
-            {
-              "$ref": "#/$defs/LastOfBlock"
-            },
-            {
-              "$ref": "#/$defs/ArrayBlock"
-            },
-            {
-              "$ref": "#/$defs/ObjectBlock"
-            },
-            {
-              "$ref": "#/$defs/MessageBlock"
-            },
-            {
-              "$ref": "#/$defs/ReadBlock"
-            },
-            {
-              "$ref": "#/$defs/IncludeBlock"
-            },
-            {
-              "$ref": "#/$defs/ImportBlock"
-            },
-            {
-              "$ref": "#/$defs/ErrorBlock"
-            },
-            {
-              "$ref": "#/$defs/EmptyBlock"
-            },
-            {
-              "type": "null"
-            }
-          ],
-          "default": null,
-          "title": "Pdl  Trace"
-        },
         "platform": {
           "const": "granite-io",
           "default": "granite-io",
@@ -6770,90 +6686,6 @@
           "default": null,
           "description": "Tokens consumed during model call\n    "
         },
-        "pdl__trace": {
-          "anyOf": [
-            {
-              "type": "boolean"
-            },
-            {
-              "type": "integer"
-            },
-            {
-              "type": "number"
-            },
-            {
-              "type": "string"
-            },
-            {
-              "$ref": "#/$defs/FunctionBlock"
-            },
-            {
-              "$ref": "#/$defs/CallBlock"
-            },
-            {
-              "$ref": "#/$defs/LitellmModelBlock"
-            },
-            {
-              "$ref": "#/$defs/GraniteioModelBlock"
-            },
-            {
-              "$ref": "#/$defs/CodeBlock"
-            },
-            {
-              "$ref": "#/$defs/ArgsBlock"
-            },
-            {
-              "$ref": "#/$defs/GetBlock"
-            },
-            {
-              "$ref": "#/$defs/DataBlock"
-            },
-            {
-              "$ref": "#/$defs/IfBlock"
-            },
-            {
-              "$ref": "#/$defs/MatchBlock"
-            },
-            {
-              "$ref": "#/$defs/RepeatBlock"
-            },
-            {
-              "$ref": "#/$defs/TextBlock"
-            },
-            {
-              "$ref": "#/$defs/LastOfBlock"
-            },
-            {
-              "$ref": "#/$defs/ArrayBlock"
-            },
-            {
-              "$ref": "#/$defs/ObjectBlock"
-            },
-            {
-              "$ref": "#/$defs/MessageBlock"
-            },
-            {
-              "$ref": "#/$defs/ReadBlock"
-            },
-            {
-              "$ref": "#/$defs/IncludeBlock"
-            },
-            {
-              "$ref": "#/$defs/ImportBlock"
-            },
-            {
-              "$ref": "#/$defs/ErrorBlock"
-            },
-            {
-              "$ref": "#/$defs/EmptyBlock"
-            },
-            {
-              "type": "null"
-            }
-          ],
-          "default": null,
-          "title": "Pdl  Trace"
-        },
         "platform": {
           "const": "litellm",
           "default": "litellm",

src/pdl/pdl_ast.py

@@ -406,11 +406,10 @@ class ModelBlock(LeafBlock):
     modelResponse: Optional[str] = None
     """Variable where to store the raw response of the model.
     """
+    # Field for internal use
     pdl__usage: Optional[PdlUsage] = None
     """Tokens consumed during model call
     """
-    # Field for internal use
-    pdl__trace: Optional["BlockType"] = None
 
 
 class LitellmModelBlock(ModelBlock):

src/pdl/pdl_ast_utils.py

@@ -45,8 +45,6 @@ def iter_block_children(f: Callable[[BlockType], None], block: BlockType) -> Non
         case ModelBlock():
             if block.input is not None:
                 f(block.input)
-            if block.pdl__trace is not None:
-                f(block.pdl__trace)
         case CodeBlock():
             f(block.code)
         case GetBlock():
@@ -137,16 +135,10 @@ def map_block_children(f: MappedFunctions, block: BlockType) -> BlockType:
             block.model = f.f_expr(block.model)
             if block.input is not None:
                 block.input = f.f_block(block.input)
-            if block.pdl__trace is not None:
-                block.pdl__trace = f.f_block(block.pdl__trace)
-            if block.parameters is not None:
-                block.parameters = f.f_expr(block.parameters)
         case GraniteioModelBlock():
             block.model = f.f_expr(block.model)
             if block.input is not None:
                 block.input = f.f_block(block.input)
-            if block.pdl__trace is not None:
-                block.pdl__trace = f.f_block(block.pdl__trace)
             if block.parameters is not None:
                 block.parameters = f.f_expr(block.parameters)
         case CodeBlock():

src/pdl/pdl_interpreter.py

@@ -1324,13 +1324,11 @@ def get_transformed_inputs(kwargs):
         result = lazy_apply(
             lambda msg: "" if msg["content"] is None else msg["content"], msg
         )
-        trace = block.model_copy(
-            update={"pdl__result": result, "pdl__trace": concrete_block}
-        )
-        if concrete_block.pdl__usage is not None:
-            trace.pdl__usage = concrete_block.pdl__usage
         if block.modelResponse is not None:
             scope = scope | {block.modelResponse: raw_result}
+        trace: BlockTypeTVarProcessCallModel = concrete_block.model_copy(
+            update={"pdl__result": result}
+        )  # pyright: ignore
         return result, background, scope, trace
     except httpx.RequestError as exc:
         message = f"model '{model_id}' encountered {repr(exc)} trying to {exc.request.method} against {exc.request.url}"