Prompt Management Using Jinja

A guide on how Jinja2 templates can be used to manage prompts

Image generated using ChatGPT

This article is free to read. But, your contributions help me keep studying and creating content for you 😊.

Introduction

How do we usually manage prompts? We mostly rely on static prompts embedded within the code with a few custom variables here and there. The standard practice is to retype the whole prompt manually when making a change.

As the system scales up, it becomes a nightmare to manage multiple lengthy prompts within the code, and to ensure prompt quality.

I was recently going through a case study on how LinkedIn built a “Prompt Source of Truth” using the Jinja2 templating language. This allowed developers to put placeholders in place of the actual prompts and then fill them dynamically at runtime.

About Jinja

Jinja is a Python-based templating engine. In simpler terms, Jinja lets us create templates with placeholders where we can programmatically pass some data. This lets us alter the context dynamically for different use cases. It’s also a good way to organise our prompts.

This removes the need to modify our code every time. Jinja lets us manage prompts without redoing the whole prompt for every situation.

A templating engine processes templates with dynamic content, rendering them dynamically based on the context.

It also provides an advantage in a situation where multiple prompts need to be chained together. And in situations where the template needs to change based on certain conditions.

One more added benefit, is that the template can also act as a guardrail by ensuring that the complete context is being provided before the prompt is being passed to the LLM.

{%if tone == "formal" %}
Hello, {{name}}
{% else %}
Hey, {{name}}
{% endif %}

Let me show you a very basic example on how to use Jinja.

import jinja2

environment = jinja2.Environment()
template = environment.from_string("Hello, {{name}}")
print(template.render(name="Arunabh"))

Here, Hello{{name}}is the template and Arunabhis the value for name which we are passing in. We get the following output.

Hello, Arunabh

LinkedIn uses a “chain-based architecture”. More specifically, every business use case is mapped to a prompt “chain” which is a predefined sequence of prompts. Every chain accepts input as a dictionary of key-value string pairs and produces a structured response.

Why can’t this be handled by the LLM itself?

You might say, with the number of smart LLMs out there right now, why can’t this task be handled by the LLM itself. This is because AI is and always has been prone to hallucinations (this is because of their inherent non-deterministic nature). It generates varied content on every attempt which makes maintaining prompt quality extremely challenging.

Actual products that operate on an enterprise scale, require handling thousands of use cases and complex business logic, and they demand operational requirements such as reliability, performance at scale, observability, integrations, and compliance.

Sure, LLMs can be used directly to build a high-fidelity prototype. But a good product must be robust, reliable, scalable and fit for commercial use.

Our use case

Let’s consider a very basic use case for this article. We are working in an industry that deals with a huge amount of data stored in multiple databases. The databases are split based on the business function. Right now, the pain point is that users need to do an extensive manual search in all the databases to find the required information.

So we need to build something which gathers data from various sources and uses LLMs to generate key insights. It will help users understand key data points and make the decision-making process easier for them.

High-Level System Design

Let’s go over what we need to do one more time. We need to create templates which can adapt based on the user input or context. At the same time, we need to add a system for versioning the prompts so that any changes don’t break the existing system. As the prompts need to change dynamically, we might need to add some conditional logic within the templates (for example, some if else statements).

For example, if the user is from the finance department, the placeholder should change the department to finance. If the user is from Testing, the placeholder should change to testing.

Moreover, we want to reuse prompts across multiple use cases. Here Jinja’s template inheritance property plays a key role. We can create a base template for the main prompt body and then keep adding additional templates to make a “chain”.

In this article, it is mentioned that LinkedIn uses LangChain. So, I’ll also consider that for our case. I’ll use LCEL as a way to combine the prompt templates to the LLM and anything else that we need for the search. Since, LangChain cannot directly integrate with a Jinja2 template, we’ll have to convert the prompts into a LangChain PromptTemplate object for using in the code. I’ll write a custom function for that.

Since users belong to different domains, I’ll create a base prompt template with different “sub-templates” which change dynamically. Each use case will take input parameters as a dictionary of strings and generate a response.

Prompt chaining is useful to accomplish complex tasks which an LLM might struggle to address if prompted with a very detailed prompt.

High level design of a Chain. This response will be in the form of a PromptTemplate object which will be used by the LLM to generate a response.

Code

I’ll create a directory first to store my prompt templates.

.
└── jinja_project/
├── prompt_templates/
│ ├── base_template.txt
│ └── template_extended.txt
├── template_loader/
│ └── load_template.py
├── .env
└── main.py

base_template.txt

This file will serve as our base prompt.

<context>
- You need to choose the database based on input from the user.
- Please, answer using formal language.
- Use short and direct sentences.
- When interacting with the user call them by their name.
- Unless the user specifies in his question a specific number of examples they wish to obtain, always limit your query to at most {{top_k}} results.
- Only use the following tables: {{table_list}}
- The schema of the tables is as follows {{schema}}
- The context of the database in use is {{db_context}}
- The name of the user is {{ user_name }}.

{% block user_type %}
{% endblock %}

template_extended.txt

This template extends the base prompt and modifies it using conditional logic.

{% extends "prompt_templates/base_template.txt" %}

{% block user_type %}
{% if department == "finance" %}
- The user is from the Finance department, look into the Finance database.
{% elif department == "SCM" %}
- The user is from Supply Chain Management, look into the Supply Chain database.
{% elif department == "sales" %}
- The user is from Sales, look in the Sales database.
{% else %}
- The department is not specified. Tell the user to specify the department.
{% endif %}
</context>

<task>
Take the user's question and generate the most suitable SQL query based on the context provided.
</task>

<question>
{question}
</question>

{% endblock %}

load_template.py

This will be the file where we define the method to return the appropriate prompt template based on the context. In this case, the context will be the user’s name and department. So, every prompt chain is mapped to the user’s department in this example. Keep in mind that the chains can be mapped to multiple things at a time (department, location, designation, etc.)

from jinja2 import Environment, FileSystemLoader
from langchain_core.prompts import PromptTemplate

# Function to load .txt file as a Prompt Template
def load_jinja2_template(
prompt_path: str,
prompt_name: str,
input_variables: list[str]=[],
partial_variables: dict[str, str]={},
jinja2_placeholders: dict[str, str]={},
) -> PromptTemplate:

env = Environment(loader=FileSystemLoader(prompt_path))
template = env.get_template(prompt_name)
prompt_string = template.render(jinja2_placeholders)

prompt_template = PromptTemplate(
template=prompt_string,
input_variables=input_variables,
partial_variables=partial_variables,
)
return prompt_template

# Example usage
template = load_jinja2_template(
prompt_path="",
prompt_name="prompt_templates/template2.txt",
jinja2_placeholders={"user_name": "Alex", "department": "Finance"}
)

print(template)

which gives the following output:

>>input_variables=[] input_types={} partial_variables={} 
>>template='- When interacting with the user call him by his name.
>>n- The name of the user is Alex.nn n- The user is from the Finance
>>department, look into the Finance database.n'

So the final prompt template, based on this input context, is:

"When interacting with the user call him by his name.
The name of the user is Alex.
The user is from the Finance department, look into the Finance database."

An added advantage of using this approach is that we can get the template as a LangChain PromptTemplate object. This will let us take advantage of LangChain’s LCEL runnable interface to implement complex prompt chains.

main.py

This is where everything will come together. We will pass in some context as input. The template manager will select the appropriate prompt chain and the program will run the LCEL chain on the dataset to return the response.

Since I have my data in a Postgres database, I’ll use LangChain’s SQLDatabase class to connect to it.

from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from dotenv import load_dotenv
import os
load_dotenv()

# Connection string
# Load database parameters from environment variables
DB_NAME = os.getenv("DB_NAME")
DB_PASSWORD = os.getenv("DB_PASSWORD")
DB_USER = os.getenv("DB_USER")
DB_HOSTNAME = os.getenv("DB_HOSTNAME")
DB_PORT = os.getenv("DB_PORT")

# Load the database and return it as a SQLDatabase object
def load_database() -> SQLDatabase:
db = SQLDatabase.from_uri(
f"postgresql+psycopg2://{DB_NAME}:{DB_PASSWORD}@{DB_HOSTNAME}:{DB_PORT}/{DB_USER}"
)
return db

I’ll write a function which will load the Jinja2 template and put some parameters as arguments — like the name of the user, their department, and the schema of the database (providing the schema provides more context to the LLM).

# Load the Jinja prompt template as per the arguments and return a ChatPromptTemplate object
def load_template(username: str, dept: str, top_k: str, schema: str, table_list: List) -> PromptTemplate:
prompt_template = load_jinja2_template(
prompt_path="",
prompt_name="prompt_templates/template_extended.txt",
jinja2_placeholders={
"user_name": username,
"department": dept,
"top_k": top_k,
"schema": schema,
"table_list": table_list,
}
)
return prompt_template

LangChain’s SQLDatabase class has features that makes it easier for us to integrate LangChain functionality with those of SQL databases. I’ll use it to pass in more context along with the prompt template so that the model can zero in on the correct answer.

if __name__ == "__main__":
# Load the database
db = load_database()

# Get usable table names
table_schema = db.get_table_info()
table_list = db.get_usable_table_names()

# Since the username, department and the user question will be passed as arguments
username = "arunabh"
department = "sales"

# Load the prompt template using all arguments
template = load_template(
username= username,
dept=department,
top_k = 10,
schema=table_schema,
table_list=table_list,
)

# Define the runnable chain
chain = template | model | StrOutputParser()
result = chain.invoke({
"question": "What quantity of Enterprise Software has been sold to TechCorp Industries in 2024?"
})
print(result)

The SQL query which we get as a result is shown below:

SELECT quantity_sold 
FROM sales
WHERE customer_name = 'TechCorp Industries'
AND product_name = 'Enterprise Software Suite'
AND EXTRACT(YEAR FROM sale_date) = 2024;

This is a good response based on the context we provided. We can go ahead and run this query on the database using db.run(). I am sharing the link to the Github repository below.

GitHub – arunabh223/jinja-prompt-template

Conclusion

This is a brief overview of how we can use the Jinja templating language to generate dynamic prompt templates. In this case, it is finally being used in a Text-to-SQL system. But it can be extended to multiple other cases as well.

Once all this is in place, make sure you have a good test dataset to evaluate your prompts and then continuously iterate.

The “Prompt Source of Truth” provides multiple advantages:

  1. Developers can manage multiple prompt versions without breaking anything in production. Anyone who wants to modify the template can just alter the .txt files without altering the code.
  2. Prompt templates could be reused across multiple applications.
  3. Use of same templates across multiple applications ensures consistency.
  4. Jinja has functionality for if-statements, loops, and more.

References

  1. https://blog.bytebytego.com/p/the-evolution-of-linkedins-generative
  2. https://github.com/maylad31/jinja-prompt-manager
  3. https://www.linkedin.com/blog/engineering/generative-ai/behind-the-platform-the-journey-to-create-the-linkedin-genai-application-tech-stack
  4. https://medium.com/@alecgg27895/jinja2-prompting-a-guide-on-using-jinja2-templates-for-prompt-management-in-genai-applications-e36e5c1243cf


Prompt Management Using Jinja was originally published in Towards AI on Medium, where people are continuing the conversation by highlighting and responding to this story.

Liked Liked