Add support of PDL code blocks (#185)

Close issue #88

Author: mandel

examples/hello/hello-code-jinja.pdl

@@ -1,4 +1,4 @@
-description: Hello world showing call out to shell command
+description: Hello world showing call out to Jinja
 defs:
   world: "World"
 lang: jinja

examples/hello/hello-code-pdl.pdl

@@ -0,0 +1,8 @@
+description: Hello world showing call out to PDL
+lang: pdl
+code: |
+  description: Hello world
+  text:
+  - "Hello\n"
+  - model: replicate/ibm-granite/granite-3.0-8b-instruct
+  

examples/hello/hello-code.pdl

@@ -7,4 +7,3 @@ text:
     import string
     result = random.choice(string.ascii_lowercase)
 - '!'
-

pdl-live/src/pdl_ast.d.ts

@@ -2322,7 +2322,7 @@ export type Kind15 = "code";
  * Programming language of the code.
  *
  */
-export type Lang = "python" | "command" | "jinja";
+export type Lang = "python" | "command" | "jinja" | "pdl";
 /**
  * Code to execute.
  *

src/pdl/pdl-schema.json

@@ -2692,7 +2692,8 @@
           "enum": [
             "python",
             "command",
-            "jinja"
+            "jinja",
+            "pdl"
           ],
           "title": "Lang",
           "type": "string"

src/pdl/pdl_ast.py

@@ -281,7 +281,7 @@ class CodeBlock(Block):
     """Execute a piece of code."""
 
     kind: Literal[BlockKind.CODE] = BlockKind.CODE
-    lang: Literal["python", "command", "jinja"]
+    lang: Literal["python", "command", "jinja", "pdl"]
     """Programming language of the code.
     """
     code: "BlocksType"

src/pdl/pdl_interpreter.py

@@ -68,7 +68,7 @@
 from .pdl_dumper import blocks_to_dict
 from .pdl_llms import BamModel, LitellmModel
 from .pdl_location_utils import append, get_loc_string
-from .pdl_parser import PDLParseError, parse_file
+from .pdl_parser import PDLParseError, parse_file, parse_str
 from .pdl_scheduler import (
     CodeYieldResultMessage,
     GeneratorWrapper,
@@ -1271,6 +1271,16 @@ def step_call_code(
                     loc=loc,
                     trace=block.model_copy(update={"code": code_s}),
                 ) from exc
+        case "pdl":
+            try:
+                result = call_pdl(code_s, scope)
+                background = [{"role": state.role, "content": result}]
+            except Exception as exc:
+                raise PDLRuntimeError(
+                    f"Code error: {repr(exc)}",
+                    loc=loc,
+                    trace=block.model_copy(update={"code": code_s}),
+                ) from exc
         case _:
             message = f"Unsupported language: {block.lang}"
             raise PDLRuntimeError(
@@ -1318,6 +1328,13 @@ def call_jinja(code: str, scope: dict) -> Any:
     return result
 
 
+def call_pdl(code: str, scope: dict) -> Any:
+    program, loc = parse_str(code)
+    state = InterpreterState()
+    result, _, _, _ = process_prog(state, scope, program, loc)
+    return result
+
+
 def step_call(
     state: InterpreterState, scope: ScopeType, block: CallBlock, loc: LocationType
 ) -> Generator[YieldMessage, Any, tuple[Any, Messages, ScopeType, CallBlock]]:

tests/results/examples/hello/hello-code-pdl.result

@@ -0,0 +1,2 @@
+Hello
+Hello! How can I assist you today?

tests/test_code.py

@@ -123,3 +123,57 @@ def test_jinja3():
     bad
     good"""
     )
+
+
+def test_pdl1():
+    prog_str = """
+lang: pdl
+code: |
+  description: Hello world
+  text:
+  - "Hello World!"
+"""
+    result = exec_str(prog_str)
+    assert result == "Hello World!"
+
+
+def test_pdl2():
+    prog_str = """
+defs:
+  w: World
+lang: pdl
+code: |
+  description: Hello world
+  text:
+  - "Hello ${w}!"
+"""
+    result = exec_str(prog_str)
+    assert result == "Hello World!"
+
+
+def test_pdl3():
+    prog_str = """
+defs:
+  x:
+    code: "result = print"
+    lang: python
+lang: pdl
+code: |
+  data: ${x}
+"""
+    result = exec_str(prog_str)
+    assert result == "<built-in function print>"
+
+
+def test_pdl4():
+    prog_str = """
+defs:
+  x:
+    code: "result = print"
+    lang: python
+lang: pdl
+code: |
+  data: ${ "${" }x ${ "}" }
+"""
+    result = exec_str(prog_str)
+    assert result == print  # pylint: disable=comparison-with-callable