support for structured decoding, removed default watsonx parameters (#168)

Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

src/pdl/pdl_ast.py

@@ -12,6 +12,8 @@
 )
 from pydantic import BaseModel, ConfigDict, Field, RootModel
 
+from .pdl_schema_utils import pdltype_to_jsonschema
+
 ScopeType: TypeAlias = dict[str, Any]
 
 ExpressionType: TypeAlias = Any
@@ -617,44 +619,49 @@ def set_default_model_params(
 
 def set_default_granite_model_parameters(
     model_id: str,
+    spec: Any,
     parameters: Optional[dict[str, Any]],
 ) -> dict[str, Any]:
     if parameters is None:
         parameters = {}
 
-    if "decoding_method" not in parameters:
-        parameters["decoding_method"] = (
-            DECODING_METHOD  # pylint: disable=attribute-defined-outside-init
-        )
-    if "max_tokens" in parameters and parameters["max_tokens"] is None:
-        parameters["max_tokens"] = (
-            MAX_NEW_TOKENS  # pylint: disable=attribute-defined-outside-init
-        )
-    if "min_new_tokens" not in parameters:
-        parameters["min_new_tokens"] = (
-            MIN_NEW_TOKENS  # pylint: disable=attribute-defined-outside-init
-        )
-    if "repetition_penalty" not in parameters:
-        parameters["repetition_penalty"] = (
-            REPETITION_PENATLY  # pylint: disable=attribute-defined-outside-init
-        )
-    if parameters["decoding_method"] == "sample":
-        if "temperature" not in parameters:
-            parameters["temperature"] = (
-                TEMPERATURE_SAMPLING  # pylint: disable=attribute-defined-outside-init
-            )
-        if "top_k" not in parameters:
-            parameters["top_k"] = (
-                TOP_K_SAMPLING  # pylint: disable=attribute-defined-outside-init
-            )
-        if "top_p" not in parameters:
-            parameters["top_p"] = (
-                TOP_P_SAMPLING  # pylint: disable=attribute-defined-outside-init
-            )
+    if spec is not None:
+        schema = pdltype_to_jsonschema(spec, True)
+        parameters["guided_decoding_backend"] = "lm-format-enforcer"
+        parameters["guided_json"] = schema
+
+    # if "decoding_method" not in parameters:
+    #    parameters["decoding_method"] = (
+    #        DECODING_METHOD  # pylint: disable=attribute-defined-outside-init
+    #    )
+    # if "max_tokens" in parameters and parameters["max_tokens"] is None:
+    #    parameters["max_tokens"] = (
+    #        MAX_NEW_TOKENS  # pylint: disable=attribute-defined-outside-init
+    #    )
+    # if "min_new_tokens" not in parameters:
+    #    parameters["min_new_tokens"] = (
+    #        MIN_NEW_TOKENS  # pylint: disable=attribute-defined-outside-init
+    #    )
+    # if "repetition_penalty" not in parameters:
+    #    parameters["repetition_penalty"] = (
+    #        REPETITION_PENATLY  # pylint: disable=attribute-defined-outside-init
+    #    )
+    # if parameters["decoding_method"] == "sample":
+    #    if "temperature" not in parameters:
+    #        parameters["temperature"] = (
+    #            TEMPERATURE_SAMPLING  # pylint: disable=attribute-defined-outside-init
+    #        )
+    #    if "top_k" not in parameters:
+    #        parameters["top_k"] = (
+    #            TOP_K_SAMPLING  # pylint: disable=attribute-defined-outside-init
+    #        )
+    #    if "top_p" not in parameters:
+    #        parameters["top_p"] = (
+    #            TOP_P_SAMPLING  # pylint: disable=attribute-defined-outside-init
+    #        )
     if "granite-3.0" in model_id:
         if "temperature" not in parameters or parameters["temperature"] is None:
             parameters["temperature"] = 0  # setting to decoding greedy
-
         if "roles" not in parameters:
             parameters["roles"] = {
                 "system": {

src/pdl/pdl_interpreter.py

@@ -1127,6 +1127,7 @@ def generate_client_response_streaming(
             msg_stream = LitellmModel.generate_text_stream(
                 model_id=block.model,
                 messages=model_input,
+                spec=block.spec,
                 parameters=litellm_parameters_to_dict(block.parameters),
             )
         case _:
@@ -1186,6 +1187,7 @@ def generate_client_response_single(
             msg, raw_result = LitellmModel.generate_text(
                 model_id=block.model,
                 messages=model_input,
+                spec=block.spec,
                 parameters=litellm_parameters_to_dict(block.parameters),
             )
     if state.yield_result:

src/pdl/pdl_llms.py

@@ -148,10 +148,13 @@ def get_model() -> None:
     def generate_text(
         model_id: str,
         messages: list[Message],
+        spec: Any,
         parameters: dict[str, Any],
     ) -> tuple[Message, Any]:
         if "granite" in model_id and "granite-20b-code-instruct-r1.1" not in model_id:
-            parameters = set_default_granite_model_parameters(model_id, parameters)
+            parameters = set_default_granite_model_parameters(
+                model_id, spec, parameters
+            )
         if parameters.get("mock_response") is not None:
             litellm.suppress_debug_info = True
         response = completion(
@@ -169,10 +172,13 @@ def generate_text(
     def generate_text_stream(
         model_id: str,
         messages: list[Message],
+        spec: Any,
         parameters: dict[str, Any],
     ) -> Generator[Message, Any, Any]:
         if "granite" in model_id and "granite-20b-code-instruct-r1.1" not in model_id:
-            parameters = set_default_granite_model_parameters(model_id, parameters)
+            parameters = set_default_granite_model_parameters(
+                model_id, spec, parameters
+            )
         response = completion(
             model=model_id,
             messages=messages,

src/pdl/pdl_schema_utils.py

@@ -30,7 +30,7 @@ def convert_to_json_type(a_type):
 
 
 def pdltype_to_jsonschema(  # pylint: disable=too-many-return-statements
-    pdl_type: str | dict[str, Any] | list
+    pdl_type: str | dict[str, Any] | list, additional_properties: bool
 ) -> dict[str, Any]:
     match pdl_type:
         case {"enum": choices}:
@@ -46,7 +46,10 @@ def pdltype_to_jsonschema(  # pylint: disable=too-many-return-statements
         case {"int": dict() as details}:
             return {"type": "integer", **details}
         case {"list": str() as type_name}:
-            return {"type": "array", "items": pdltype_to_jsonschema(type_name)}
+            return {
+                "type": "array",
+                "items": pdltype_to_jsonschema(type_name, additional_properties),
+            }
         case {"list": dict() as details}:
             ikws = ["enum", *_PDLTYPE_TO_JSONSCHEMA_NAME.keys()]
             items_details = {k: v for k, v in details.items() if k in ikws}
@@ -55,43 +58,50 @@ def pdltype_to_jsonschema(  # pylint: disable=too-many-return-statements
             other_details = {k: v for k, v in details.items() if k not in ikws}
             return {
                 "type": "array",
-                "items": pdltype_to_jsonschema(items_details),
+                "items": pdltype_to_jsonschema(items_details, additional_properties),
                 **other_details,
             }
         case list() as type_list:
             if len(type_list) != 1:
                 raise ValueError(f"invalid PDL type {pdl_type}")
             return {
                 "type": "array",
-                "items": pdltype_to_jsonschema(type_list[0]),
+                "items": pdltype_to_jsonschema(type_list[0], additional_properties),
             }
         case {"obj": dict() as pdl_props}:
-            return get_json_schema_object(pdl_props)
+            return get_json_schema_object(pdl_props, additional_properties)
         case dict() as pdl_props:
-            return get_json_schema_object(pdl_props)
+            return get_json_schema_object(pdl_props, additional_properties)
     raise ValueError(f"invalid PDL type {pdl_type}")
 
 
-def get_json_schema_object(pdl_props: dict) -> dict[str, Any]:
+def get_json_schema_object(pdl_props: dict, additional_properties) -> dict[str, Any]:
     props = {}
     required = []
     for name, prop_type in pdl_props.items():
         if isinstance(prop_type, dict) and "optional" in prop_type:
-            props[name] = pdltype_to_jsonschema(prop_type["optional"])
+            props[name] = pdltype_to_jsonschema(
+                prop_type["optional"], additional_properties
+            )
         else:
-            props[name] = pdltype_to_jsonschema(prop_type)
+            props[name] = pdltype_to_jsonschema(prop_type, additional_properties)
             required.append(name)
-    return {
-        "type": "object",
-        "properties": props,
-        "required": required,
-        "additionalProperties": False,
-    }
+    if additional_properties is False:
+        return {
+            "type": "object",
+            "properties": props,
+            "required": required,
+            "additionalProperties": False,
+        }
+
+    return {"type": "object", "properties": props, "required": required}
 
 
-def get_json_schema(params: dict[str, Any]) -> Optional[dict[str, Any]]:
+def get_json_schema(
+    params: dict[str, Any], additional_properties
+) -> Optional[dict[str, Any]]:
     try:
-        result = pdltype_to_jsonschema({"obj": params})
+        result = pdltype_to_jsonschema({"obj": params}, additional_properties)
         return result
     except ValueError as e:
         warnings.warn(e.args[0])

src/pdl/pdl_schema_validator.py

@@ -22,14 +22,14 @@ def type_check_args(args: dict[str, Any], params: dict[str, Any], loc) -> list[s
     if "pdl_context" in args_copy:
         # params_copy["pdl_context"] = [{"role": "str?", "content": "str"}]
         params_copy["pdl_context"] = ["obj"]
-    schema = get_json_schema(params_copy)
+    schema = get_json_schema(params_copy, False)
     if schema is None:
         return ["Error obtaining a valid schema from function parameters definition"]
     return type_check(args_copy, schema, loc)
 
 
 def type_check_spec(result: Any, spec: str | dict[str, Any] | list, loc) -> list[str]:
-    schema = pdltype_to_jsonschema(spec)
+    schema = pdltype_to_jsonschema(spec, False)
     if schema is None:
         return ["Error obtaining a valid schema from spec"]
     return type_check(result, schema, loc)

tests/test_type_checking.py

@@ -145,7 +145,7 @@
 def test_pdltype_to_jsonschema():
     for t in _PDLTYPE_TO_JSONSCHEMA_TESTS:
         pdl_type = yaml.safe_load(t["pdl_type"])
-        json_schema = pdltype_to_jsonschema(pdl_type)
+        json_schema = pdltype_to_jsonschema(pdl_type, False)
         assert json_schema == t["json_schema"]