mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 08:23:17 +00:00
[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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user