Granite 3.1 chat template fix (#433)

* Granite 3.1 chat template fix
Signed-off-by: Mandana Vaziri <[email protected]>

Author: vazirim

src/pdl/pdl_ast.py

@@ -739,13 +739,17 @@ def get_default_model_parameters() -> list[dict[str, Any]]:
                         "post_message": "<|end_of_text|>",
                     },
                     "available_tools": {
-                        "pre_message": "<|start_of_role|>available_tools<|end_of_role|>",
+                        "pre_message": "<|start_of_role|>tools<|end_of_role|>",
                         "post_message": "<|end_of_text|>",
                     },
                     "tool_response": {
                         "pre_message": "<|start_of_role|>tool_response<|end_of_role|>",
                         "post_message": "<|end_of_text|>",
                     },
+                    "documents": {
+                        "pre_message": "<|start_of_role|>documents<|end_of_role|>",
+                        "post_message": "<|end_of_text|>",
+                    },
                 },
                 "final_prompt_value": "<|start_of_role|>assistant<|end_of_role|>",
             }

tests/test_defaults.py

@@ -90,13 +90,17 @@ def test_default_model_params():
                 "post_message": "<|end_of_text|>",
             },
             "available_tools": {
-                "pre_message": "<|start_of_role|>available_tools<|end_of_role|>",
+                "pre_message": "<|start_of_role|>tools<|end_of_role|>",
                 "post_message": "<|end_of_text|>",
             },
             "tool_response": {
                 "pre_message": "<|start_of_role|>tool_response<|end_of_role|>",
                 "post_message": "<|end_of_text|>",
             },
+            "documents": {
+                "pre_message": "<|start_of_role|>documents<|end_of_role|>",
+                "post_message": "<|end_of_text|>",
+            },
         },
         "final_prompt_value": "<|start_of_role|>assistant<|end_of_role|>",
     } == params