Merge branch 'main' into anonymous

Author: mandel

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

@@ -19,6 +19,7 @@
     BeforeValidator,
     ConfigDict,
     Field,
+    Json,
     RootModel,
     TypeAdapter,
 )
@@ -228,23 +229,32 @@ class OptionalPdlType(PdlType):
     optional: "PdlTypeType"
 
 
+class JsonSchemaTypePdlType(PdlType):
+    """Json Schema type"""
+
+    model_config = ConfigDict(extra="allow")
+    type: str | list[str]
+
+
 class ObjPdlType(PdlType):
-    """Optional type."""
+    """Object type."""
 
     obj: Optional[dict[str, "PdlTypeType"]]
 
 
 PdlTypeType = TypeAliasType(
     "PdlTypeType",
     Annotated[
-        "Union[BasePdlType,"  # pyright: ignore
+        "Union[None,"  # pyright: ignore
+        "      BasePdlType,"
         "      EnumPdlType,"
         "      StrPdlType,"
         "      FloatPdlType,"
         "      IntPdlType,"
         "      ListPdlType,"
         "      list['PdlTypeType'],"
         "      OptionalPdlType,"
+        "      JsonSchemaTypePdlType,"
         "      ObjPdlType,"
         "      dict[str, 'PdlTypeType']]",
         Field(union_mode="left_to_right"),
@@ -261,7 +271,7 @@ class Parser(BaseModel):
     description: Optional[str] = None
     """Documentation associated to the parser.
     """
-    spec: Optional[PdlTypeType] = None
+    spec: PdlTypeType = None
     """Expected type of the parsed value.
     """
 
@@ -340,7 +350,7 @@ class Block(BaseModel):
     description: Optional[str] = None
     """Documentation associated to the block.
     """
-    spec: Optional[PdlTypeType] = None
+    spec: PdlTypeType = None
     """Type specification of the result of the block.
     """
     defs: dict[str, "BlockType"] = {}
@@ -360,6 +370,12 @@ class Block(BaseModel):
     fallback: Optional["BlockType"] = None
     """Block to execute in case of error.
     """
+    retry: Optional[int] = None
+    """The maximum number of times to retry when an error occurs within a block.
+    """
+    trace_error_on_retry: Optional[bool] | str = None
+    """Whether to add the errors while retrying to the trace. Set this to true to use retry feature for multiple LLM trials.
+    """
     role: RoleType = None
     """Role associated to the block and sub-blocks.
     Typical roles are `system`, `user`, and `assistant`,
@@ -402,8 +418,12 @@ class FunctionBlock(LeafBlock):
     """Functions parameters with their types.
     """
     returns: "BlockType" = Field(..., alias="return")
-    """Body of the function
+    """Body of the function.
     """
+    signature: Optional[Json] = None
+    """Function signature computed from the function definition.
+    """
+
     # Field for internal use
     pdl__scope: SkipJsonSchema[Optional[ScopeType]] = Field(default=None, repr=False)
 

src/pdl/pdl-schema.json

@@ -32,6 +32,7 @@
     IntPdlType,
     JoinText,
     JoinType,
+    JsonSchemaTypePdlType,
     LastOfBlock,
     ListPdlType,
     ListPdlTypeConstraints,
@@ -124,6 +125,12 @@ def block_to_dict(  # noqa: C901
         d["defs"] = {
             x: block_to_dict(b, json_compatible) for x, b in block.defs.items()
         }
+    if block.retry is not None:
+        d["retry"] = expr_to_dict(block.retry, json_compatible)
+    if block.trace_error_on_retry is not None:
+        d["trace_error_on_retry"] = expr_to_dict(
+            block.trace_error_on_retry, json_compatible
+        )
     if isinstance(block, StructuredBlock):
         d["context"] = block.context
 
@@ -310,8 +317,10 @@ def expr_to_dict(expr: ExpressionType, json_compatible: bool):
 
 
 def type_to_dict(t: PdlTypeType):
-    d: str | list | dict
+    d: None | str | list | dict
     match t:
+        case None:
+            d = None
         case "null" | "bool" | "str" | "float" | "int" | "list" | "obj":
             d = t
         case EnumPdlType():
@@ -377,6 +386,8 @@ def type_to_dict(t: PdlTypeType):
             assert False, "list must have only one element"
         case OptionalPdlType():
             d = {"optional": type_to_dict(t.optional)}
+        case JsonSchemaTypePdlType():
+            d = t.model_dump()
         case ObjPdlType():
             if t.obj is None:
                 d = "obj"

src/pdl/pdl_ast.py

@@ -10,6 +10,7 @@
 
 # TODO: temporarily disabling warnings to mute a pydantic warning from liteLLM
 import warnings
+from functools import partial
 from os import getenv
 
 warnings.filterwarnings("ignore", "Valid config keys have changed in V2")
@@ -99,6 +100,7 @@
 from .pdl_parser import PDLParseError, parse_file, parse_str  # noqa: E402
 from .pdl_python_repl import PythonREPL  # noqa: E402
 from .pdl_scheduler import yield_background, yield_result  # noqa: E402
+from .pdl_schema_utils import get_json_schema  # noqa: E402
 from .pdl_schema_validator import type_check_args, type_check_spec  # noqa: E402
 from .pdl_utils import (  # noqa: E402
     GeneratorWrapper,
@@ -341,6 +343,34 @@ def identity(result):
     return identity
 
 
+def set_error_to_scope_for_retry(
+    scope: ScopeType, error, block_id: Optional[str] = ""
+) -> ScopeType:
+    repeating_same_error = False
+    pdl_context: Optional[LazyMessages] = scope.get("pdl_context")
+    if pdl_context is None:
+        return scope
+    if pdl_context and isinstance(pdl_context, list):
+        last_msg = pdl_context[-1]
+        last_error = last_msg["content"]
+        if last_error.endswith(error):
+            repeating_same_error = True
+    if repeating_same_error:
+        error = "The previous error occurs multiple times."
+    err_msg = {
+        "role": "assistant",
+        "content": error,
+        "defsite": block_id,
+    }
+    scope = scope | {
+        "pdl_context": lazy_messages_concat(
+            pdl_context,
+            PdlList([err_msg]),
+        )
+    }
+    return scope
+
+
 def process_advanced_block(
     state: InterpreterState,
     scope: ScopeType,
@@ -361,52 +391,85 @@ def process_advanced_block(
     state = state.with_yield_background(
         state.yield_background and context_in_contribute(block)
     )
-    try:
-        result, background, new_scope, trace = process_block_body(
-            state, scope, block, loc
-        )
-        result = lazy_apply(id_with_set_first_use_nanos(block.pdl__timing), result)
-        background = lazy_apply(
-            id_with_set_first_use_nanos(block.pdl__timing), background
-        )
-        trace = trace.model_copy(update={"pdl__result": result})
-        if block.parser is not None:
-            parser = block.parser
-            result = lazy_apply(lambda r: parse_result(parser, r), result)
-            if init_state.yield_result and ContributeTarget.RESULT:
-                yield_result(result, block.kind)
-        if block.spec is not None and not isinstance(block, FunctionBlock):
-            result = lazy_apply(
-                lambda r: result_with_type_checking(
-                    r, block.spec, "Type errors during spec checking:", loc, trace
-                ),
-                result,
+
+    # Bind result variables here with empty values
+    result: PdlLazy[Any] = PdlConst(None)
+    background: LazyMessages = PdlList([{}])
+    new_scope: ScopeType = PdlDict({})
+    trace: AdvancedBlockType = EmptyBlock()
+
+    max_retry = block.retry if block.retry else 0
+    trial_total = max_retry + 1
+    for trial_idx in range(trial_total):
+        try:
+            result, background, new_scope, trace = process_block_body(
+                state, scope, block, loc
             )
-        if block.fallback is not None:
-            result.result()
-    except Exception as exc:
-        if block.fallback is None:
-            raise exc from exc
-        (
-            result,
-            background,
-            new_scope,
-            trace,
-        ) = process_block_of(
-            block,
-            "fallback",
-            state,
-            scope,
-            loc=loc,
-        )
-        if block.spec is not None and not isinstance(block, FunctionBlock):
-            loc = append(loc, "fallback")
-            result = lazy_apply(
-                lambda r: result_with_type_checking(
-                    r, block.spec, "Type errors during spec checking:", loc, trace
-                ),
+            result = lazy_apply(id_with_set_first_use_nanos(block.pdl__timing), result)
+            background = lazy_apply(
+                id_with_set_first_use_nanos(block.pdl__timing), background
+            )
+            trace = trace.model_copy(update={"pdl__result": result})
+            if block.parser is not None:
+                # Use partial to create a function with fixed arguments
+                parser_func = partial(parse_result, block.parser)
+                result = lazy_apply(parser_func, result)
+                if init_state.yield_result and ContributeTarget.RESULT:
+                    yield_result(result, block.kind)
+            if block.spec is not None and not isinstance(block, FunctionBlock):
+                # Use partial to create a function with fixed arguments
+                checker = partial(
+                    result_with_type_checking,
+                    spec=block.spec,
+                    msg="Type errors during spec checking:",
+                    loc=loc,
+                    trace=trace,
+                )
+                result = lazy_apply(checker, result)
+            if block.fallback is not None:
+                result.result()
+            break
+        except Exception as exc:
+            err_msg = exc.args[0]
+            do_retry = (
+                block.retry
+                and trial_idx + 1 < trial_total
+                and "Keyboard Interrupt" not in err_msg
+            )
+            if block.fallback is None and not do_retry:
+                raise exc from exc
+            if do_retry:
+                error = f"An error occurred in a PDL block. Error details: {err_msg}"
+                print(
+                    f"\n\033[0;31m[Retry {trial_idx+1}/{max_retry}] {error}\033[0m\n",
+                    file=sys.stderr,
+                )
+                if block.trace_error_on_retry:
+                    scope = set_error_to_scope_for_retry(scope, error, block.pdl__id)
+                continue
+            (
                 result,
+                background,
+                new_scope,
+                trace,
+            ) = process_block_of(
+                block,
+                "fallback",
+                state,
+                scope,
+                loc=loc,
             )
+            if block.spec is not None and not isinstance(block, FunctionBlock):
+                loc = append(loc, "fallback")
+                # Use partial to create a function with fixed arguments
+                checker = partial(
+                    result_with_type_checking,
+                    spec=block.spec,
+                    msg="Type errors during spec checking:",
+                    loc=loc,
+                    trace=trace,
+                )
+                result = lazy_apply(checker, result)
     if block.def_ is not None:
         var = block.def_
         new_scope = new_scope | PdlDict({var: result})
@@ -832,6 +895,16 @@ def process_block_body(
             if block.def_ is not None:
                 scope = scope | {block.def_: closure}
             closure.pdl__scope = scope
+            signature: dict[str, Any] = {"type": "function"}
+            if block.def_ is not None:
+                signature["name"] = block.def_
+            if block.description is not None:
+                signature["description"] = block.description
+            if block.function is not None:
+                signature["parameters"] = get_json_schema(block.function, False) or {}
+            else:
+                signature["parameters"] = {}
+            closure.signature = signature
             result = PdlConst(closure)
             background = PdlList([])
             trace = closure.model_copy(update={})
@@ -914,6 +987,8 @@ def process_defs(
         state = state.with_iter(idx)
         state = state.with_yield_result(False)
         state = state.with_yield_background(False)
+        if isinstance(block, FunctionBlock) and block.def_ is None:
+            block = block.model_copy(update={"def_": x})
         result, _, _, block_trace = process_block(state, scope, block, newloc)
         scope = scope | PdlDict({x: result})
         defs_trace[x] = block_trace

src/pdl/pdl_dumper.py

@@ -183,7 +183,7 @@ def generate_text_stream(
 
 
 def set_structured_decoding_parameters(
-    spec: Optional[PdlTypeType],
+    spec: PdlTypeType,
     parameters: Optional[dict[str, Any]],
 ) -> dict[str, Any]:
     if parameters is None:

src/pdl/pdl_interpreter.py

@@ -5,6 +5,7 @@
     EnumPdlType,
     FloatPdlType,
     IntPdlType,
+    JsonSchemaTypePdlType,
     ListPdlType,
     ListPdlTypeConstraints,
     ObjPdlType,
@@ -43,7 +44,7 @@ def convert_to_json_type(a_type):
 
 
 def pdltype_to_jsonschema(
-    pdl_type: Optional[PdlTypeType], additional_properties: bool
+    pdl_type: PdlTypeType, additional_properties: bool
 ) -> dict[str, Any]:
     schema: dict[str, Any]
     match pdl_type:
@@ -111,6 +112,12 @@ def pdltype_to_jsonschema(
         case OptionalPdlType(optional=t):
             t_schema = pdltype_to_jsonschema(t, additional_properties)
             schema = {"anyOf": [t_schema, "null"]}
+        case JsonSchemaTypePdlType(type=t):
+            if pdl_type.__pydantic_extra__ is None:
+                extra = {}
+            else:
+                extra = pdl_type.__pydantic_extra__
+            schema = {"type": t, **extra}
         case ObjPdlType(obj=pdl_props):
             if pdl_props is None:
                 schema = {"type": "object"}

src/pdl/pdl_llms.py

@@ -8,7 +8,9 @@
 
 
 def type_check_args(
-    args: Optional[dict[str, Any]], params: Optional[dict[str, Any]], loc
+    args: Optional[dict[str, Any]],
+    params: Optional[dict[str, PdlTypeType]],
+    loc,
 ) -> list[str]:
     if (args == {} or args is None) and (params is None or params == {}):
         return []
@@ -35,7 +37,7 @@ def type_check_args(
     return type_check(args_copy, schema, loc)
 
 
-def type_check_spec(result: Any, spec: Optional[PdlTypeType], loc) -> list[str]:
+def type_check_spec(result: Any, spec: PdlTypeType, loc) -> list[str]:
     schema = pdltype_to_jsonschema(spec, False)
     if schema is None:
         return ["Error obtaining a valid schema from spec"]