feat: make ruamel the default yaml handler for better formatting and …

…format preservation; also automatically prettify long strings in yaml; robustness improvement for templates with images in OAI format
MoritzLaurer · Dec 18, 2024 · b267502 · b267502
1 parent 45a8054
commit b267502
Show file tree

Hide file tree

Showing 23 changed files with 518 additions and 391 deletions.
diff --git a/examples/example-usage.ipynb b/examples/example-usage.ipynb
@@ -65,6 +65,60 @@
     "#importlib.reload(prompt_templates.tools)\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "0822c4c6",
+   "metadata": {},
+   "source": [
+    "## Test ruamel outputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a634a66",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from prompt_templates.utils import format_template_content, create_yaml_handler\n",
+    "\n",
+    "# Sample data simulating your prompt template\n",
+    "template = {\n",
+    "    \"prompt\": {\n",
+    "        \"template\": [\n",
+    "            {\n",
+    "                \"role\": \"system\",\n",
+    "                \"content\": \"You are a coding assistant who explains concepts clearly and provides short examples.\\n\\nYou are a coding assistant who explains concepts clearly and provides short examples.\"\n",
+    "            },\n",
+    "            {\n",
+    "                \"role\": \"user\",\n",
+    "                \"content\": \"Explain what {{concept}} is in {{programming_language}}.\"\n",
+    "            }\n",
+    "        ],\n",
+    "        \"template_variables\": [\"concept\", \"programming_language\"],\n",
+    "        \"metadata\": {\n",
+    "            \"name\": \"Code Teacher\",\n",
+    "            \"description\": \"A simple chat prompt for explaining programming concepts with examples\",\n",
+    "            \"tags\": [\"programming\", \"education\"],\n",
+    "            \"version\": \"0.0.1\",\n",
+    "            \"author\": \"Guido van Bossum\"\n",
+    "        },\n",
+    "        \"client_parameters\": {},\n",
+    "        \"custom_data\": {}\n",
+    "    }\n",
+    "}\n",
+    "\n",
+    "# Apply formatting\n",
+    "formatted_template = format_template_content(template)\n",
+    "\n",
+    "# Create YAML handler\n",
+    "yaml = create_yaml_handler()\n",
+    "\n",
+    "# Dump to YAML file\n",
+    "with open(\"code_teacher_test.yaml\", \"w\") as f:\n",
+    "    yaml.dump(formatted_template, f)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "6ded7c48",
@@ -279,7 +333,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 3,
    "id": "f8ad2dde",
    "metadata": {},
    "outputs": [],
@@ -290,7 +344,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 4,
    "id": "0e10ad0f",
    "metadata": {},
    "outputs": [
@@ -305,7 +359,7 @@
    "source": [
     "from prompt_templates import ChatPromptTemplate\n",
     "messages_template = [\n",
-    "    {\"role\": \"system\", \"content\": \"You are a coding assistant who explains concepts clearly and provides short examples.\"},\n",
+    "    {\"role\": \"system\", \"content\": \"You are a coding assistant who explains concepts clearly and provides short examples.\\n\\n\" *10 },\n",
     "    {\"role\": \"user\", \"content\": \"Explain what {{concept}} is in {{programming_language}}.\"}\n",
     "]\n",
     "template_variables = [\"concept\", \"programming_language\"]\n",
@@ -331,24 +385,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "eac98b0b",
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "No files have been modified since last commit. Skipping to prevent empty commit.\n"
-     ]
-    },
     {
      "data": {
       "text/plain": [
-       "CommitInfo(commit_url='https://huggingface.co/datasets/MoritzLaurer/example_prompts_test/commit/abf5c1f51419f69e2d3516faaf6d85a6234cc40e', commit_message='Upload prompt template code_teacher_test.yaml', commit_description='', oid='abf5c1f51419f69e2d3516faaf6d85a6234cc40e', pr_url=None, repo_url=RepoUrl('https://huggingface.co/datasets/MoritzLaurer/example_prompts_test', endpoint='https://huggingface.co', repo_type='dataset', repo_id='MoritzLaurer/example_prompts_test'), pr_revision=None, pr_num=None)"
+       "CommitInfo(commit_url='https://huggingface.co/datasets/MoritzLaurer/example_prompts_test/commit/1749258b19b5551f5d51496c5b2d125a0e45bda2', commit_message='Upload prompt template code_teacher_test.yaml', commit_description='', oid='1749258b19b5551f5d51496c5b2d125a0e45bda2', pr_url=None, repo_url=RepoUrl('https://huggingface.co/datasets/MoritzLaurer/example_prompts_test', endpoint='https://huggingface.co', repo_type='dataset', repo_id='MoritzLaurer/example_prompts_test'), pr_revision=None, pr_num=None)"
       ]
      },
-     "execution_count": 4,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }

diff --git a/poetry.lock b/poetry.lock
diff --git a/prompt_templates/loaders.py b/prompt_templates/loaders.py
@@ -7,13 +7,13 @@
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, Set, Union
 
-import yaml
 from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.utils import validate_repo_id
 
 from .constants import VALID_PROMPT_EXTENSIONS, PopulatorType
 from .prompt_templates import ChatPromptTemplate, TextPromptTemplate
 from .tools import Tool
+from .utils import create_yaml_handler
 
 
 logger = logging.getLogger(__name__)
@@ -56,13 +56,15 @@ def from_local(
         path: Union[str, Path],
         populator: Optional[PopulatorType] = None,
         jinja2_security_level: Literal["strict", "standard", "relaxed"] = "standard",
+        yaml_library: str = "ruamel",
     ) -> Union[TextPromptTemplate, ChatPromptTemplate]:
         """Load a prompt template from a local YAML file.
 
         Args:
             path (Union[str, Path]): Path to the YAML file containing the prompt template
             populator (Optional[PopulatorType], optional): The populator type to use among Literal["double_brace", "single_brace", "jinja2"]. Defaults to None.
             jinja2_security_level (Literal["strict", "standard", "relaxed"], optional): The security level for the Jinja2 populator. Defaults to "standard".
+            yaml_library (str, optional): The YAML library to use ("ruamel" or "pyyaml"). Defaults to "ruamel".
 
         Returns:
             Union[TextPromptTemplate, ChatPromptTemplate]: The loaded template instance
@@ -103,10 +105,14 @@ def from_local(
         if path.suffix not in VALID_PROMPT_EXTENSIONS:
             raise ValueError(f"Template file must be a .yaml or .yml file, got: {path}")
 
+        yaml = create_yaml_handler(yaml_library)
         try:
             with open(path, "r") as file:
-                prompt_file = yaml.safe_load(file)
-        except yaml.YAMLError as e:
+                if yaml_library == "ruamel":
+                    prompt_file = yaml.load(file)
+                else:
+                    prompt_file = yaml.safe_load(file)
+        except Exception as e:
             raise ValueError(
                 f"Failed to parse '{path}' as a valid YAML file. "
                 f"Please ensure the file is properly formatted.\n"
@@ -126,6 +132,7 @@ def from_hub(
         revision: Optional[str] = None,
         populator: Optional[PopulatorType] = None,
         jinja2_security_level: Literal["strict", "standard", "relaxed"] = "standard",
+        yaml_library: str = "ruamel",
     ) -> Union[TextPromptTemplate, ChatPromptTemplate]:
         """Load a prompt template from the Hugging Face Hub.
 
@@ -141,6 +148,7 @@ def from_hub(
                 Can be a branch name, tag, or commit hash. Defaults to None
             populator (Optional[PopulatorType], optional): The populator type to use among Literal["double_brace", "single_brace", "jinja2"]. Defaults to None.
             jinja2_security_level (Literal["strict", "standard", "relaxed"], optional): The security level for the Jinja2 populator. Defaults to "standard".
+            yaml_library (str, optional): The YAML library to use ("ruamel" or "pyyaml"). Defaults to "ruamel".
 
 
         Returns:
@@ -200,10 +208,14 @@ def from_hub(
         except Exception as e:
             raise FileNotFoundError(f"Failed to download template from Hub: {str(e)}") from e
 
+        yaml = create_yaml_handler(yaml_library)
         try:
             with open(file_path, "r") as file:
-                prompt_file = yaml.safe_load(file)
-        except yaml.YAMLError as e:
+                if yaml_library == "ruamel":
+                    prompt_file = yaml.load(file)
+                else:
+                    prompt_file = yaml.safe_load(file)
+        except Exception as e:
             raise ValueError(
                 f"Failed to parse '{filename}' as a valid YAML file. "
                 f"Please ensure the file is properly formatted.\n"
@@ -267,7 +279,7 @@ def _load_template_from_yaml(
         custom_data = {
             k: v
             for k, v in prompt_data.items()
-            if k not in ["template", "template_variables", "metadata", "client_parameters"]
+            if k not in ["template", "template_variables", "metadata", "client_parameters", "custom_data"]
         }
 
         # Determine template type and create appropriate instance