Creating a custom packager - round-trip (pack + unpack) tutorial

Creating a custom packager - round-trip (pack + unpack) tutorial#

This tutorial walks through creating a round-trip custom packager for langchain_core.prompts.ChatPromptTemplate. Round-trip means the packager handles both output serialization (packing) and input deserialization (unpacking) — a prompt template produced by one function can be consumed as a typed input by another.

You will learn how to:

Implement pack_file() to serialize a ChatPromptTemplate to a human-readable JSON file
Implement unpack_file() to load it back from a DataItem
Test the full round-trip: one handler produces a template, another consumes it

Prerequisite

This tutorial requires langchain-core. Install it with:

pip install langchain

In this section

The problem
Setup
Create the custom packager
Test 1: pack a prompt template
Test 2: round-trip — consume as typed input
Recap

The problem#

Without a custom packager, returning a ChatPromptTemplate from an MLRun function causes it to be pickled. The resulting .pkl file is:

Opaque — you can't read or review the prompt structure
Fragile — pickle files break when LangChain versions change

Instead, the template should be saved as a readable JSON file that captures the message structure, and you can load it back as a ChatPromptTemplate in downstream functions.

Setup#

Create the project#

import mlrun

project = mlrun.get_or_create_project("round-trip-tutorial", "./")

Write a function#

This handler creates a ChatPromptTemplate with template variables ({domain}, {response}) and returns it. The packager will serialize it to JSON.

%%writefile build_prompt.py

from langchain_core.prompts import ChatPromptTemplate


def build_eval_prompt() -> ChatPromptTemplate:
    """
    Build a chat prompt template for evaluating AI agent responses.

    :returns: A ChatPromptTemplate with system + human messages.
    """
    return ChatPromptTemplate.from_messages([
        ("system", "You are an expert evaluator for the {domain} domain. "
                   "Rate the following response on a scale of 1 to 10."),
        ("human", "Agent response: {response}\n\nProvide your rating and reasoning."),
    ])

build_fn = project.set_function(
    "build_prompt.py",
    name="build-prompt",
    kind="job",
    image="mlrun/mlrun",
    handler="build_eval_prompt",
)

Create the custom packager#

Write the packager#

The packager serializes a ChatPromptTemplate as a JSON file containing each message's role and template string. On unpack, it reconstructs the template from the same JSON structure.

This approach is human-readable and independent of LangChain's internal serialization format.

%%writefile prompt_packager.py

import json
import os
import tempfile

from langchain_core.prompts import ChatPromptTemplate

from mlrun.artifacts import Artifact
from mlrun.package.packagers.default_packager import DefaultPackager
from mlrun.package.utils import ArtifactType


class ChatPromptTemplatePackager(DefaultPackager):
    """
    A custom packager for LangChain ChatPromptTemplate objects.

    Serializes the template as a JSON file containing the list of
    (role, template_string) message pairs.
    """

    PACKABLE_OBJECT_TYPE = ChatPromptTemplate
    DEFAULT_PACKING_ARTIFACT_TYPE = ArtifactType.FILE
    DEFAULT_UNPACKING_ARTIFACT_TYPE = ArtifactType.FILE

    def pack_file(
        self,
        obj: ChatPromptTemplate,
        key: str,
    ) -> tuple[Artifact, dict]:
        """
        Serialize a ChatPromptTemplate to a JSON file.

        The JSON contains a list of message dicts, each with 'role'
        and 'template' keys.

        :param obj: The ChatPromptTemplate to pack.
        :param key: The artifact key.

        :returns: The artifact and unpacking instructions.
        """
        # Extract messages as (role, template_string) pairs
        messages = []
        for message_prompt in obj.messages:
            role = message_prompt.__class__.__name__  # e.g. 'SystemMessagePromptTemplate'
            # Get the template string from the inner prompt
            template = message_prompt.prompt.template
            messages.append({"role": role, "template": template})

        # Write to a temporary JSON file
        temp_dir = tempfile.mkdtemp()
        file_path = os.path.join(temp_dir, f"{key}.json")
        with open(file_path, "w") as f:
            json.dump({"messages": messages}, f, indent=2)

        # Mark for cleanup
        self.add_future_clearing_path(temp_dir)

        artifact = Artifact(key=key, src_path=file_path)
        return artifact, {}

    def unpack_file(
        self,
        data_item,
    ) -> ChatPromptTemplate:
        """
        Deserialize a ChatPromptTemplate from a JSON file.

        :param data_item: The data item pointing to the JSON artifact.

        :returns: The reconstructed ChatPromptTemplate.
        """
        # Map class names back to role tuples that from_messages() expects
        role_map = {
            "SystemMessagePromptTemplate": "system",
            "HumanMessagePromptTemplate": "human",
            "AIMessagePromptTemplate": "ai",
        }

        local_path = self.get_data_item_local_path(data_item=data_item)
        with open(local_path) as f:
            data = json.load(f)

        # Reconstruct using from_messages()
        message_tuples = []
        for msg in data["messages"]:
            role = role_map.get(msg["role"], msg["role"])
            message_tuples.append((role, msg["template"]))

        return ChatPromptTemplate.from_messages(message_tuples)

Notice:

DEFAULT_PACKING_ARTIFACT_TYPE = "file" and DEFAULT_UNPACKING_ARTIFACT_TYPE = "file" — both packing and unpacking default to the file artifact type
pack_file() extracts each message's role and template string into a clean JSON structure; this is human-readable — you can inspect the artifact in any JSON viewer
unpack_file() reads the JSON and reconstructs the template using ChatPromptTemplate.from_messages() — the standard LangChain constructor
The role_map translates LangChain class names back to the short role strings ("system", "human", "ai") that from_messages() expects

Register the packager#

project.add_custom_packager(
    packager="prompt_packager.ChatPromptTemplatePackager", is_mandatory=True
)

Test 1: pack a prompt template#

Run the handler to produce a ChatPromptTemplate. The packager serializes it to a human-readable JSON artifact.

build_run = build_fn.run(
    local=True,
    returns=["eval_prompt"],  # default artifact type → file (JSON)
)

build_run.outputs

Test 2: round-trip — consume as typed input#

This handler receives the prompt template as a ChatPromptTemplate parameter. The packager automatically deserializes it from the JSON artifact — no manual file reading needed.

This demonstrates the round-trip: pack (Test 1) → artifact → unpack (Test 2).

%%writefile use_prompt.py

from langchain_core.prompts import ChatPromptTemplate


def format_prompt(
    template: ChatPromptTemplate,
    response: str = "Photosynthesis converts light energy into chemical energy.",
) -> str:
    """
    Format the prompt template with a sample response.

    :param template: The ChatPromptTemplate to use.
    :param response: A sample agent response to evaluate.

    :returns: The formatted prompt string.
    """
    messages = template.format_messages(
        domain="science",
        response=response,
    )
    # Join message contents for display
    return "\n---\n".join(msg.content for msg in messages)

use_fn = project.set_function(
    "use_prompt.py",
    name="use-prompt",
    kind="job",
    image="mlrun/mlrun",
    handler="format_prompt",
)

use_run = use_fn.run(
    local=True,
    inputs={"template": build_run.outputs["eval_prompt"]},
    params={"response": "Water boils at 100 degrees Celsius at sea level."},
    returns=["formatted_prompt"],
)

print(use_run.outputs["formatted_prompt"])

Recap#

In this tutorial you built a round-trip custom packager for ChatPromptTemplate:

Step	What you did
Subclass	`class ChatPromptTemplatePackager(DefaultPackager)`
Set type	`PACKABLE_OBJECT_TYPE = ChatPromptTemplate`
Pack	`pack_file()` — serializes messages to a readable JSON file
Unpack	`unpack_file()` — reconstructs from JSON via `from_messages()`
Register	`project.add_custom_packager(packager="module.Class", is_mandatory=True)`
Round-trip test	Handler 1 produces a template → Handler 2 consumes it as a typed input

The key difference from the PIL Image packager is that this packager implements both pack_file() and unpack_file(), enabling the template to flow between functions as a first-class typed object.

Next: See the EvalSuite packager tutorial for a packager that supports bundling and unbundling — decomposing a collection into separate artifacts.