[Feat] Prompt Management - Add support for versioning prompts (#16836)

* test_dotprompt_auto_detection_with_model_only

* fix _auto_detect_prompt_management_logger

* test_dotprompt_with_prompt_version

* add v1, v2 tests

* add _compile_prompt_helper

* fix _compile_prompt_helper

* test_dotprompt_with_prompt_version

* test_dotprompt_with_prompt_version, test_get_prompt_with_version
This commit is contained in:
Ishaan Jaff
2025-11-19 13:16:03 -08:00
committed by GitHub
parent 1f8fe007a1
commit 3ebe489082
5 changed files with 209 additions and 12 deletions
@@ -108,7 +108,7 @@ class DotpromptManager(CustomPromptManagement):
Compile a .prompt file into a PromptManagementClient structure.
This method:
1. Loads the prompt template from the .prompt file
1. Loads the prompt template from the .prompt file (with optional version)
2. Renders it with the provided variables
3. Converts the rendered text into chat messages
4. Extracts model and optional parameters from metadata
@@ -116,13 +116,22 @@ class DotpromptManager(CustomPromptManagement):
try:
# Get the prompt template
template = self.prompt_manager.get_prompt(prompt_id)
# Get the prompt template (versioned or base)
template = self.prompt_manager.get_prompt(
prompt_id=prompt_id, version=prompt_version
)
if template is None:
raise ValueError(f"Prompt '{prompt_id}' not found in prompt directory")
version_str = f" (version {prompt_version})" if prompt_version else ""
raise ValueError(
f"Prompt '{prompt_id}'{version_str} not found in prompt directory"
)
# Render the template with variables
rendered_content = self.prompt_manager.render(prompt_id, prompt_variables)
# Render the template with variables (pass version for proper lookup)
rendered_content = self.prompt_manager.render(
prompt_id=prompt_id,
prompt_variables=prompt_variables,
version=prompt_version,
)
# Convert rendered content to chat messages
messages = self._convert_to_messages(rendered_content)
@@ -183,7 +183,10 @@ class PromptManager:
return frontmatter, template_content
def render(
self, prompt_id: str, prompt_variables: Optional[Dict[str, Any]] = None
self,
prompt_id: str,
prompt_variables: Optional[Dict[str, Any]] = None,
version: Optional[int] = None,
) -> str:
"""
Render a prompt template with the given variables.
@@ -191,6 +194,7 @@ class PromptManager:
Args:
prompt_id: The ID of the prompt template to render
prompt_variables: Variables to substitute in the template
version: Optional version number. If provided, looks for {prompt_id}.v{version}
Returns:
The rendered prompt string
@@ -199,13 +203,16 @@ class PromptManager:
KeyError: If prompt_id is not found
ValueError: If template rendering fails
"""
if prompt_id not in self.prompts:
# Get the template (versioned or base)
template = self.get_prompt(prompt_id=prompt_id, version=version)
if template is None:
available_prompts = list(self.prompts.keys())
version_str = f" (version {version})" if version else ""
raise KeyError(
f"Prompt '{prompt_id}' not found. Available prompts: {available_prompts}"
f"Prompt '{prompt_id}'{version_str} not found. Available prompts: {available_prompts}"
)
template = self.prompts[prompt_id]
variables = prompt_variables or {}
# Validate input variables against schema if defined
@@ -254,8 +261,26 @@ class PromptManager:
return type_mapping.get(schema_type.lower(), str) # type: ignore
def get_prompt(self, prompt_id: str) -> Optional[PromptTemplate]:
"""Get a prompt template by ID."""
def get_prompt(
self, prompt_id: str, version: Optional[int] = None
) -> Optional[PromptTemplate]:
"""
Get a prompt template by ID and optional version.
Args:
prompt_id: The base prompt ID
version: Optional version number. If provided, looks for {prompt_id}.v{version}
Returns:
The prompt template if found, None otherwise
"""
if version is not None:
# Try versioned prompt first: prompt_id.v{version}
versioned_id = f"{prompt_id}.v{version}"
if versioned_id in self.prompts:
return self.prompts[versioned_id]
# Fall back to base prompt_id
return self.prompts.get(prompt_id)
def list_prompts(self) -> List[str]:
@@ -0,0 +1,10 @@
---
model: gpt-3.5-turbo
temperature: 0.5
max_tokens: 100
input:
schema:
user_message: string
---
Version 1: {{user_message}}
@@ -0,0 +1,10 @@
---
model: gpt-4
temperature: 0.9
max_tokens: 200
input:
schema:
user_message: string
---
Version 2: {{user_message}}
@@ -195,6 +195,29 @@ def test_get_prompt_metadata():
assert "output" in metadata
def test_get_prompt_with_version():
"""Test that get_prompt correctly retrieves versioned prompts."""
prompt_dir = Path(__file__).parent
manager = PromptManager(prompt_directory=str(prompt_dir))
# Get base prompt (no version)
base_prompt = manager.get_prompt(prompt_id="chat_prompt")
assert base_prompt is not None
assert "User: {{user_message}}" in base_prompt.content
# Get version 1
v1_prompt = manager.get_prompt(prompt_id="chat_prompt", version=1)
assert v1_prompt is not None
assert "Version 1:" in v1_prompt.content
assert v1_prompt.model == "gpt-3.5-turbo"
# Get version 2
v2_prompt = manager.get_prompt(prompt_id="chat_prompt", version=2)
assert v2_prompt is not None
assert "Version 2:" in v2_prompt.content
assert v2_prompt.model == "gpt-4"
def test_add_prompt_programmatically():
"""Test adding prompts programmatically."""
prompt_dir = Path(
@@ -597,3 +620,123 @@ async def test_dotprompt_auto_detection_with_model_only():
finally:
# Restore original callbacks
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_dotprompt_with_prompt_version():
"""
Test that dotprompt can load and use specific prompt versions.
Versions are stored as separate files with .v{version}.prompt naming convention.
"""
from litellm.integrations.dotprompt import DotpromptManager
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
prompt_dir = Path(__file__).parent
dotprompt_manager = DotpromptManager(prompt_directory=str(prompt_dir))
# Register the dotprompt manager in callbacks
original_callbacks = litellm.callbacks.copy()
litellm.callbacks = [dotprompt_manager]
try:
# Mock the HTTP handler to avoid actual API calls
with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post:
mock_response_data = litellm.ModelResponse(
choices=[
litellm.Choices(
message=litellm.Message(content="Hello!"),
index=0,
finish_reason="stop",
)
]
).model_dump()
# Create a proper mock response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.text = json.dumps(mock_response_data)
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = mock_response_data
mock_post.return_value = mock_response
# Test version 1
await litellm.acompletion(
model="gpt-3.5-turbo",
prompt_id="chat_prompt",
prompt_version=1,
prompt_variables={"user_message": "Test v1"},
messages=[],
)
assert mock_post.call_count >= 1
data_str = mock_post.call_args.kwargs.get("data", "{}")
request_body = json.loads(data_str)
print(f"Version 1 request body: {json.dumps(request_body, indent=2)}")
# Verify version 1 prompt was used
# chat_prompt.v1.prompt has: model: gpt-3.5-turbo, temperature: 0.5, max_tokens: 100
assert request_body["model"] == "gpt-3.5-turbo"
# Verify the message contains "Version 1:" prefix from v1 template
messages = request_body["messages"]
assert len(messages) >= 1
first_message_content = messages[0]["content"]
print(f"Version 1 message: {first_message_content}")
assert "Version 1:" in first_message_content
assert "Test v1" in first_message_content
# Reset mock for version 2 test
mock_post.reset_mock()
# Test version 2
with patch("litellm.llms.custom_httpx.llm_http_handler.AsyncHTTPHandler.post") as mock_post:
mock_response_data = litellm.ModelResponse(
choices=[
litellm.Choices(
message=litellm.Message(content="Hello!"),
index=0,
finish_reason="stop",
)
]
).model_dump()
# Create a proper mock response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.text = json.dumps(mock_response_data)
mock_response.headers = {"Content-Type": "application/json"}
mock_response.json.return_value = mock_response_data
mock_post.return_value = mock_response
await litellm.acompletion(
model="gpt-4",
prompt_id="chat_prompt",
prompt_version=2,
prompt_variables={"user_message": "Test v2"},
messages=[],
)
mock_post.assert_called_once()
data_str = mock_post.call_args.kwargs.get("data", "{}")
request_body = json.loads(data_str)
print(f"Version 2 request body: {json.dumps(request_body, indent=2)}")
# Verify version 2 prompt was used
# chat_prompt.v2.prompt has: model: gpt-4, temperature: 0.9, max_tokens: 200
assert request_body["model"] == "gpt-4"
# Verify the message contains "Version 2:" prefix from v2 template
messages = request_body["messages"]
assert len(messages) >= 1
first_message_content = messages[0]["content"]
print(f"Version 2 message: {first_message_content}")
assert "Version 2:" in first_message_content
assert "Test v2" in first_message_content
finally:
# Restore original callbacks
litellm.callbacks = original_callbacks