Location (WIP)

Author: mandel

src/pdl/pdl-schema.json

@@ -3155,6 +3155,12 @@
           "type": "string"
         },
         "data": {
+          "anyOf": [
+            {},
+            {
+              "$ref": "#/$defs/LocalizedExpression"
+            }
+          ],
           "description": "Value defined.",
           "title": "Data"
         },
@@ -6428,6 +6434,12 @@
           "type": "string"
         },
         "if": {
+          "anyOf": [
+            {},
+            {
+              "$ref": "#/$defs/LocalizedExpression"
+            }
+          ],
           "description": "Condition.\n    ",
           "title": "If"
         },
@@ -9035,6 +9047,31 @@
       "title": "LitellmParameters",
       "type": "object"
     },
+    "LocalizedExpression": {
+      "additionalProperties": false,
+      "description": "Expression with location information\n    ",
+      "properties": {
+        "expr": {
+          "title": "Expr"
+        },
+        "location": {
+          "anyOf": [
+            {
+              "$ref": "#/$defs/LocationType"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "default": null
+        }
+      },
+      "required": [
+        "expr"
+      ],
+      "title": "LocalizedExpression",
+      "type": "object"
+    },
     "LocationType": {
       "additionalProperties": false,
       "properties": {
@@ -12945,6 +12982,12 @@
           "title": "Repeat"
         },
         "until": {
+          "anyOf": [
+            {},
+            {
+              "$ref": "#/$defs/LocalizedExpression"
+            }
+          ],
           "description": "Condition of the loop.\n    ",
           "title": "Until"
         },

src/pdl/pdl_ast.py

@@ -4,6 +4,7 @@
 from enum import StrEnum
 from typing import Any, Literal, Optional, TypeAlias, TypedDict, Union
 
+import strictyaml
 from genai.schema import (
     DecodingMethod,
     ModerationParameters,
@@ -14,17 +15,6 @@
 
 ScopeType: TypeAlias = dict[str, Any]
 
-ExpressionType: TypeAlias = Any
-# (
-#     str
-#     | int
-#     | float
-#     | bool
-#     | None
-#     | list["ExpressionType"]
-#     | dict[str, "ExpressionType"]
-# )
-
 
 class Message(TypedDict):
     role: Optional[str]
@@ -63,9 +53,32 @@ class LocationType(BaseModel):
     table: dict[str, int]
 
 
+YamlSource: TypeAlias = strictyaml.YAML
+
 empty_block_location = LocationType(file="", path=[], table={})
 
 
+class LocalizedExpression(BaseModel):
+    """Expression with location information"""
+
+    model_config = ConfigDict(extra="forbid", use_attribute_docstrings=True)
+    expr: Any
+    location: Optional[LocationType] = None
+    _pdl_yaml_src: Optional[YamlSource] = None
+
+
+ExpressionType: TypeAlias = Any | LocalizedExpression
+# (
+#     str
+#     | int
+#     | float
+#     | bool
+#     | None
+#     | list["ExpressionType"]
+#     | dict[str, "ExpressionType"]
+# )
+
+
 class Parser(BaseModel):
     model_config = ConfigDict(extra="forbid")
     description: Optional[str] = None
@@ -124,6 +137,7 @@ class Block(BaseModel):
     # Fields for internal use
     result: Optional[Any] = None
     location: Optional[LocationType] = None
+    _pdl_yaml_src: Optional[YamlSource] = None
 
 
 class FunctionBlock(Block):

src/pdl/pdl_ast_utils.py

@@ -28,15 +28,18 @@
 )
 
 
+def is_block_list(blocks: BlocksType) -> bool:
+    return not isinstance(blocks, str) and isinstance(blocks, Sequence)
+
+
 def iter_block_children(f: Callable[[BlockType], None], block: BlockType) -> None:
     if not isinstance(block, Block):
         return
     for blocks in block.defs.values():
         iter_blocks(f, blocks)
     match block:
         case FunctionBlock():
-            if block.returns is not None:
-                iter_blocks(f, block.returns)
+            iter_blocks(f, block.returns)
         case CallBlock():
             if block.trace is not None:
                 iter_blocks(f, block.trace)
@@ -105,7 +108,7 @@ def iter_block_children(f: Callable[[BlockType], None], block: BlockType) -> Non
 
 
 def iter_blocks(f: Callable[[BlockType], None], blocks: BlocksType) -> None:
-    if not isinstance(blocks, str) and isinstance(blocks, Sequence):
+    if is_block_list(blocks):
         for block in blocks:
             f(block)
     else:

src/pdl/pdl_dumper.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Sequence, TypeAlias
+from typing import Any, TypeAlias
 
 import yaml
 
@@ -35,6 +35,7 @@
     RepeatUntilBlock,
     TextBlock,
 )
+from .pdl_ast_utils import is_block_list
 
 yaml.SafeDumper.org_represent_str = yaml.SafeDumper.represent_str  # type: ignore
 
@@ -235,7 +236,7 @@ def blocks_to_dict(
     result: (
         int | float | str | dict[str, Any] | list[int | float | str | dict[str, Any]]
     )
-    if not isinstance(blocks, str) and isinstance(blocks, Sequence):
+    if is_block_list(blocks):
         result = [block_to_dict(block, json_compatible) for block in blocks]
     else:
         result = block_to_dict(blocks, json_compatible)

src/pdl/pdl_interpreter.py

@@ -8,7 +8,7 @@
 
 # from itertools import batched
 from pathlib import Path
-from typing import Any, Generator, Optional, Sequence, TypeVar
+from typing import Any, Generator, Optional, TypeVar
 
 import litellm
 import yaml
@@ -65,6 +65,7 @@
     TextBlock,
     empty_block_location,
 )
+from .pdl_ast_utils import is_block_list
 from .pdl_dumper import blocks_to_dict
 from .pdl_llms import BamModel, LitellmModel
 from .pdl_location_utils import append, get_loc_string
@@ -799,7 +800,7 @@ def step_blocks(
     background: Messages
     trace: BlocksType
     results = []
-    if not isinstance(blocks, str) and isinstance(blocks, Sequence):
+    if is_block_list(blocks):
         iteration_state = state.with_yield_result(
             state.yield_result and iteration_type != IterationType.ARRAY
         )

src/pdl/pdl_parser.py

@@ -1,11 +1,26 @@
 import json
 from pathlib import Path
+from typing import Any, Optional
 
+import strictyaml
 import yaml
 from pydantic import ValidationError
 
-from .pdl_ast import LocationType, PDLException, Program
-from .pdl_location_utils import get_line_map
+from .pdl_ast import (
+    Block,
+    BlocksType,
+    CallBlock,
+    DataBlock,
+    FunctionBlock,
+    LastOfBlock,
+    LocationType,
+    ModelBlock,
+    PDLException,
+    Program,
+    YamlSource,
+)
+from .pdl_ast_utils import is_block_list
+from .pdl_location_utils import append, get_line_map
 from .pdl_schema_error_analyzer import analyze_errors
 
 
@@ -25,6 +40,7 @@ def parse_str(pdl_str: str, file_name: str = "") -> tuple[Program, LocationType]
     loc = LocationType(path=[], file=file_name, table=line_table)
     try:
         prog = Program.model_validate(prog_yaml)
+        set_program_location(prog, pdl_str)
     except ValidationError as exc:
         pdl_schema_file = Path(__file__).parent / "pdl-schema.json"
         with open(pdl_schema_file, "r", encoding="utf-8") as schema_fp:
@@ -35,3 +51,119 @@ def parse_str(pdl_str: str, file_name: str = "") -> tuple[Program, LocationType]
             errors = ["The file do not respect the schema."]
         raise PDLParseError(errors) from exc
     return prog, loc
+
+
+def set_program_location(prog: Program, pdl_str: str, file_name: str = ""):
+    loc = strictyaml.dirty_load(pdl_str)
+    set_location(prog.root, loc)
+
+
+def set_location(
+    pdl: Any,
+    loc: YamlSource,
+):
+    if hasattr(pdl, "_pdl_yaml_src"):
+        pdl._pdl_yaml_src = loc
+    if isinstance(loc.data, dict):
+        for x, v in loc.items():
+            if hasattr(pdl, x.data):
+                set_location(getattr(pdl, x.data), v)
+    elif isinstance(pdl, list) and isinstance(loc.data, list):
+        for data_i, loc_i in zip(pdl, loc):
+            set_location(data_i, loc_i)
+
+
+# def set_program_location(prog: Program, pdl_str: str, file_name: str = ""):
+#     line_table = get_line_map(pdl_str)
+#     loc = LocationType(path=[], file=file_name, table=line_table)
+#     return Program(set_blocks_location(prog.root, loc))
+
+# def set_blocks_location(
+#     blocks: BlocksType,
+#     loc: YAML,
+# ):
+#     if is_block_list(blocks):
+#         return [set_block_location(block, append(loc, f"[{i}]")) for i, block in enumerate(blocks)]
+#     return set_block_location(blocks, loc)
+
+
+# def set_block_location(
+#     block: BlocksType,
+#     loc: LocationType,
+# ):
+#     if not isinstance(block, Block):
+#         return DataBlock(data=block, location=loc)
+#     block = block.model_copy(update={"location": loc})
+#     defs_loc = append(loc, "defs")
+#     block.defs = {x: set_block_location(b, append(defs_loc, x)) for x, b in block.defs }
+#     if block.parser is not None:
+#         block.parser = set_parser_location(block.parser)
+#     if block.fallback is not None:
+#         block.fallback = set_block_location(block.fallback, append(loc, "fallback"))
+#     match block:
+#         case FunctionBlock():
+#             block.returns = set_blocks_location(block.returns, append(loc, "return"))
+#         case CallBlock():
+#             block.args = {x: set_expr_location(expr) for x, expr in block.args.items()}
+#         case ModelBlock():
+#             if block.input is not None:
+#                 iter_blocks(f, block.input)
+#         case CodeBlock():
+#             iter_blocks(f, block.code)
+#         case GetBlock():
+#             pass
+#         case DataBlock():
+#             pass
+#         case TextBlock():
+#             iter_blocks(f, block.text)
+#         case LastOfBlock():
+#             iter_blocks(f, block.lastOf)
+#         case ArrayBlock():
+#             iter_blocks(f, block.array)
+#         case ObjectBlock():
+#             if isinstance(block.object, dict):
+#                 body = list(block.object.values())
+#             else:
+#                 body = block.object
+#             iter_blocks(f, body)
+#         case MessageBlock():
+#             iter_blocks(f, block.content)
+#         case IfBlock():
+#             iter_blocks(f, block.then)
+#             if block.elses is not None:
+#                 iter_blocks(f, block.elses)
+#         case RepeatBlock():
+#             iter_blocks(f, block.repeat)
+#             if block.trace is not None:
+#                 for trace in block.trace:
+#                     iter_blocks(f, trace)
+#         case RepeatUntilBlock():
+#             iter_blocks(f, block.repeat)
+#             if block.trace is not None:
+#                 for trace in block.trace:
+#                     iter_blocks(f, trace)
+#         case ForBlock():
+#             iter_blocks(f, block.repeat)
+#             if block.trace is not None:
+#                 for trace in block.trace:
+#                     iter_blocks(f, trace)
+#         case ErrorBlock():
+#             iter_blocks(f, block.program)
+#         case ReadBlock():
+#             pass
+#         case IncludeBlock():
+#             if block.trace is not None:
+#                 iter_blocks(f, block.trace)
+#         case EmptyBlock():
+#             pass
+#         case _:
+#             assert (
+#                 False
+#             ), f"Internal error (missing case iter_block_children({type(block)}))"
+#     match (block.parser):
+#         case "json" | "yaml" | RegexParser():
+#             pass
+#         case PdlParser():
+#             iter_blocks(f, block.parser.pdl)
+#     if block.fallback is not None:
+#         iter_blocks(f, block.fallback)