Handling of localized expressions

Author: mandel

pyproject.toml

@@ -15,7 +15,8 @@ dependencies = [
   "jsonschema~=4.0",
   "litellm~=1.49",
   "termcolor~=2.0",
-  "ipython~=8.0"
+  "ipython~=8.0",
+  "strictyaml~=1.7.3"
 ]
 authors = [
   { name="Mandana Vaziri", email="[email protected]" },

src/pdl/pdl_ast.py

@@ -64,7 +64,7 @@ class LocalizedExpression(BaseModel):
     model_config = ConfigDict(extra="forbid", use_attribute_docstrings=True)
     expr: Any
     location: Optional[LocationType] = None
-    _pdl_yaml_src: Optional[YamlSource] = None
+    pdl_yaml_src: Optional[YamlSource] = None
 
 
 ExpressionType: TypeAlias = Any | LocalizedExpression
@@ -137,7 +137,7 @@ class Block(BaseModel):
     # Fields for internal use
     result: Optional[Any] = None
     location: Optional[LocationType] = None
-    _pdl_yaml_src: Optional[YamlSource] = None
+    pdl_yaml_src: Optional[YamlSource] = None
 
 
 class FunctionBlock(Block):
@@ -275,7 +275,7 @@ class ModelBlock(Block):
 class BamModelBlock(ModelBlock):
     platform: Literal[ModelPlatform.BAM]
     prompt_id: Optional[str] = None
-    parameters: Optional[BamTextGenerationParameters | dict] = None
+    parameters: Optional[BamTextGenerationParameters | ExpressionType] = None
     moderations: Optional[ModerationParameters] = None
     data: Optional[PromptTemplateData] = None
     constraints: Any = None  # TODO
@@ -285,7 +285,7 @@ class LitellmModelBlock(ModelBlock):
     """Call a LLM through the LiteLLM API: https://docs.litellm.ai/."""
 
     platform: Literal[ModelPlatform.LITELLM] = ModelPlatform.LITELLM
-    parameters: Optional[LitellmParameters | dict] = None
+    parameters: Optional[LitellmParameters | ExpressionType] = None
 
 
 class CodeBlock(Block):

src/pdl/pdl_compilers/to_regex.py

@@ -19,6 +19,7 @@
     IncludeBlock,
     LitellmModelBlock,
     LitellmParameters,
+    LocalizedExpression,
     ModelBlock,
     ReadBlock,
     RepeatBlock,
@@ -273,10 +274,14 @@ def compile_block(
                             "include_stop_sequence", False
                         )
                     else:
-                        stop_sequences = block.parameters.stop_sequences or []
+                        if isinstance(block.parameters, LocalizedExpression):
+                            parameters = block.parameters.expr
+                        else:
+                            parameters = block.parameters
+                        stop_sequences = parameters.stop_sequences or []
                         include_stop_sequence = (
-                            block.parameters.include_stop_sequence is None
-                            or block.parameters.include_stop_sequence
+                            parameters.include_stop_sequence is None
+                            or parameters.include_stop_sequence
                         )
                 case LitellmModelBlock():
                     if block.parameters is None:
@@ -285,6 +290,8 @@ def compile_block(
                     else:
                         if isinstance(block.parameters, LitellmParameters):
                             parameters = block.parameters.model_dump()
+                        elif isinstance(block.parameters, LocalizedExpression):
+                            parameters = block.parameters.expr
                         else:
                             parameters = block.parameters
                         stop_sequences = parameters.get("stop", [])

src/pdl/pdl_interpreter.py

@@ -46,6 +46,7 @@
     LastOfBlock,
     LitellmModelBlock,
     LitellmParameters,
+    LocalizedExpression,
     LocationType,
     Message,
     MessageBlock,
@@ -917,8 +918,12 @@ def process_condition_of(
 EXPR_END_STRING = "}"
 
 
-def process_expr(scope: ScopeType, expr: Any, loc: LocationType) -> Any:
+def process_expr(
+    scope: ScopeType, expr: Any, loc: LocationType
+) -> Any:  # pylint: disable=too-many-return-statements
     result: Any
+    if isinstance(expr, LocalizedExpression):
+        return process_expr(scope, expr.expr, loc)
     if isinstance(expr, str):
         try:
             if expr.startswith(EXPR_START_STRING) and expr.endswith(EXPR_END_STRING):
@@ -1001,7 +1006,7 @@ def step_call_model(
     ],
 ]:
     # evaluate model name
-    _, concrete_block = process_expr_of(block, "model", scope, loc)
+    model, concrete_block = process_expr_of(block, "model", scope, loc)
     # evaluate model params
     match concrete_block:
         case BamModelBlock():
@@ -1060,9 +1065,7 @@ def get_transformed_inputs(kwargs):
         if "input" in litellm_params:
             append_log(state, "Model Input", litellm_params["input"])
         else:
-            append_log(
-                state, "Model Input", messages_to_str(concrete_block.model, model_input)
-            )
+            append_log(state, "Model Input", messages_to_str(model, model_input))
         background: Messages = [msg]
         result = msg["content"]
         append_log(state, "Model Output", result)
@@ -1104,6 +1107,8 @@ def generate_client_response_streaming(
     model_input: Messages,
 ) -> Generator[YieldMessage, Any, Message]:
     msg_stream: Generator[Message, Any, None]
+    assert isinstance(block.model, str)  # block is a "concrete block"
+    assert isinstance(block.parameters, dict)  # block is a "concrete block"
     model_input_str = messages_to_str(block.model, model_input)
     match block:
         case BamModelBlock():
@@ -1158,6 +1163,8 @@ def generate_client_response_single(
     block: BamModelBlock | LitellmModelBlock,
     model_input: Messages,
 ) -> Generator[YieldMessage, Any, Message]:
+    assert isinstance(block.model, str)  # block is a "concrete block"
+    assert isinstance(block.parameters, dict)  # block is a "concrete block"
     msg: Message
     model_input_str = messages_to_str(block.model, model_input)
     match block:
@@ -1189,6 +1196,8 @@ def generate_client_response_batching(  # pylint: disable=too-many-arguments
     # model: str,
     model_input: Messages,
 ) -> Generator[YieldMessage, Any, Message]:
+    assert isinstance(block.model, str)  # block is a "concrete block"
+    assert isinstance(block.parameters, dict)  # block is a "concrete block"
     model_input_str = messages_to_str(block.model, model_input)
     match block:
         case BamModelBlock():

src/pdl/pdl_parser.py

@@ -49,8 +49,8 @@ def set_location(
     pdl: Any,
     loc: YamlSource,
 ):
-    if hasattr(pdl, "_pdl_yaml_src"):
-        pdl._pdl_yaml_src = loc
+    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):