scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(prompts): reject keyed prompt_data with prompt_id and populate prompt version

Browse source 
POST /prompts silently stored an empty template when litellm_params.prompt_id
was combined with prompt_data keyed by template name, because the loader
wrapped the already-keyed dict under prompt_id a second time. The loader now
wraps only a flat template (a dict carrying a content key), and create,
update, and patch reject the ambiguous keyed+prompt_id combination with a 400
that names both valid shapes. The API also returned version null on every
create and lost version, environment, and created_by on registry reload; both
now carry through. Versioned ids like my-prompt.v1, which the create API
itself returns, now resolve to their base template on the SDK prompt hooks,
and a flat DB prompt with no litellm_params.prompt_id registers under its base
API id instead of garbage.
This commit is contained in:
mateo-berri 2026-08-26 13:57:10 -07:00
parent f6571a653f
commit 6fa5164d86
7 changed files with 284 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/integrations/dotprompt/__init__.py
View file

@ -62,12 +62,16 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom
if dotprompt_content and not prompt_data and not prompt_file:
prompt_data = _get_prompt_data_from_dotprompt_content(dotprompt_content)
from .prompt_manager import strip_version_suffix
registration_prompt_id: Final = prompt_id or strip_version_suffix(prompt_spec.prompt_id) or prompt_spec.prompt_id
try:
dot_prompt_manager: Final = DotpromptManager(
prompt_directory=prompt_directory,
prompt_data=prompt_data,
prompt_file=prompt_file,
prompt_id=prompt_id,
prompt_id=registration_prompt_id,
)
return dot_prompt_manager

2
litellm/integrations/dotprompt/dotprompt_manager.py
View file

@ -96,7 +96,7 @@ class DotpromptManager(CustomPromptManagement):
if prompt_id is None:
return False
try:
return prompt_id in self.prompt_manager.list_prompts()
return self.prompt_manager.get_prompt(prompt_id) is not None
except Exception:
# If there's any error accessing prompts, don't run prompt management
return False

31
litellm/integrations/dotprompt/prompt_manager.py
View file

@ -11,6 +11,13 @@ from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
def strip_version_suffix(prompt_id: str) -> str | None:
base, separator, version = prompt_id.rpartition(".v")
if separator and base and version.isdigit():
return base
return None
class PromptTemplate:
"""Represents a single prompt template with metadata and content."""
@ -124,11 +131,13 @@ class PromptManager:
"content": "template content",
"metadata": {"model": "gpt-4", "temperature": 0.7, ...}
} + prompt_id
"""
if prompt_id:
prompt_data = {prompt_id: prompt_data}
for prompt_id, prompt_info in prompt_data.items():
A dict carrying a "content" key is a single flat template registered under
prompt_id; anything else is treated as already keyed by template ID.
"""
keyed_prompts: Final = {prompt_id: prompt_data} if prompt_id and "content" in prompt_data else prompt_data
for template_id, prompt_info in keyed_prompts.items():
try:
content = prompt_info.get("content", "")
metadata = prompt_info.get("metadata", {})
@ -136,11 +145,11 @@ class PromptManager:
template = PromptTemplate(
content=content,
metadata=metadata,
template_id=prompt_id,
template_id=template_id,
)
self.prompts[prompt_id] = template
self.prompts[template_id] = template
except Exception:
# Optional: print(f"Error loading prompt from JSON: {prompt_id}")
# Optional: print(f"Error loading prompt from JSON: {template_id}")
pass
def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate:
@ -272,8 +281,12 @@ class PromptManager:
if versioned_id in self.prompts:
return self.prompts[versioned_id]
# Fall back to base prompt_id
return self.prompts.get(prompt_id)
direct_match: Final = self.prompts.get(prompt_id)
if direct_match is not None:
return direct_match
base_prompt_id: Final = strip_version_suffix(prompt_id)
return self.prompts.get(base_prompt_id) if base_prompt_id else None
def list_prompts(self) -> list[str]:
"""Get a list of all available prompt IDs."""

31
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -323,6 +323,7 @@ def create_versioned_prompt_spec(db_prompt: _PromptRow) -> PromptSpec:
prompt_info=prompt_info,
created_at=row.created_at,
updated_at=row.updated_at,
version=row.version,
environment=row.environment,
created_by=row.created_by,
)
@ -334,6 +335,21 @@ class Prompt(BaseModel):
prompt_info: PromptInfo | None = None
AMBIGUOUS_PROMPT_DATA_ERROR: Final = (
"litellm_params.prompt_id cannot be combined with prompt_data keyed by template name. "
'Send a flat template, prompt_data={"content": "...", "metadata": {...}}, together with litellm_params.prompt_id, '
'or send prompt_data={"<template_id>": {"content": "...", "metadata": {...}}} without litellm_params.prompt_id.'
)
def is_ambiguous_keyed_prompt_data(litellm_params: PromptLiteLLMParams) -> bool:
extra_fields: Final = litellm_params.model_extra or {}
prompt_data: Final = extra_fields.get("prompt_data")
if not litellm_params.prompt_id or not isinstance(prompt_data, dict):
return False
return bool(prompt_data) and "content" not in prompt_data
class PatchPromptRequest(BaseModel):
litellm_params: PromptLiteLLMParams | None = None
prompt_info: PromptInfo | None = None
@ -737,11 +753,9 @@ async def create_prompt(
-d '{
"prompt_id": "my_prompt",
"litellm_params": {
"prompt_id": "json_prompt",
"prompt_id": "my_prompt",
"prompt_integration": "dotprompt",
### EITHER prompt_directory OR prompt_data MUST BE PROVIDED
"prompt_directory": "/path/to/dotprompt/folder",
"prompt_data": {"json_prompt": {"content": "This is a prompt", "metadata": {"model": "gpt-4"}}}
"prompt_data": {"content": "This is a prompt", "metadata": {"model": "gpt-4"}}
},
"prompt_info": {
"prompt_type": "config"
@ -763,6 +777,9 @@ async def create_prompt(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
if is_ambiguous_keyed_prompt_data(request.litellm_params):
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
try:
# Extract environment from request
environment: Final = (
@ -857,6 +874,9 @@ async def update_prompt(
if prisma_client is None:
raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
if is_ambiguous_keyed_prompt_data(request.litellm_params):
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
try:
# Strip version suffix from prompt_id if present (e.g., "jack_success.v1" -> "jack_success")
base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id)
@ -1143,6 +1163,9 @@ async def patch_prompt(
if updated_litellm_params is None:
raise HTTPException(status_code=400, detail="litellm_params cannot be None")
if is_ambiguous_keyed_prompt_data(updated_litellm_params):
raise HTTPException(status_code=400, detail=AMBIGUOUS_PROMPT_DATA_ERROR)
# Build update data dict
update_data: Final[dict[str, str]] = {
"litellm_params": updated_litellm_params.model_dump_json(),

3
litellm/proxy/prompts/prompt_registry.py
View file

@ -147,6 +147,9 @@ class InMemoryPromptRegistry:
prompt_info=prompt.prompt_info or PromptInfo(prompt_type="config"),
created_at=prompt.created_at,
updated_at=prompt.updated_at,
version=prompt.version,
environment=prompt.environment,
created_by=prompt.created_by,
)
# store references to the prompt in memory

80
tests/test_litellm/integrations/dotprompt/test_prompt_manager.py
View file

@ -577,3 +577,83 @@ async def test_dotprompt_with_prompt_version():
)
assert "Version 2:" in v2_rendered
assert "Test v2" in v2_rendered
def test_keyed_prompt_data_with_prompt_id_keeps_real_content():
prompt_data = {
"json_prompt": {
"content": "You are a pirate. Begin every reply with AHOY.",
"metadata": {"model": "gpt-4o-mini"},
}
}
manager = PromptManager(prompt_data=prompt_data, prompt_id="agent-prompt")
template = manager.get_prompt("json_prompt")
assert template is not None
assert template.content == "You are a pirate. Begin every reply with AHOY."
assert template.model == "gpt-4o-mini"
assert "agent-prompt" not in manager.prompts
def test_flat_prompt_data_with_prompt_id_registers_under_prompt_id():
manager = PromptManager(
prompt_data={"content": "Hello {{name}}", "metadata": {"model": "gpt-4o-mini"}},
prompt_id="flat-prompt",
)
template = manager.get_prompt("flat-prompt")
assert template is not None
assert template.content == "Hello {{name}}"
assert manager.render("flat-prompt", {"name": "world"}) == "Hello world"
def test_get_prompt_falls_back_to_base_id_for_versioned_id():
manager = PromptManager(
prompt_data={"content": "Hi", "metadata": {}},
prompt_id="my-prompt",
)
assert manager.get_prompt("my-prompt.v1") is not None
assert manager.get_prompt("my-prompt.v12") is not None
assert manager.get_prompt("my-prompt.vx") is None
assert manager.get_prompt("other-prompt.v1") is None
def test_should_run_prompt_management_accepts_versioned_id():
from litellm.integrations.dotprompt import DotpromptManager
dotprompt_manager = DotpromptManager(
prompt_data={"content": "Hi", "metadata": {}},
prompt_id="versioned-prompt",
)
assert dotprompt_manager.should_run_prompt_management("versioned-prompt", None, {}) is True
assert dotprompt_manager.should_run_prompt_management("versioned-prompt.v1", None, {}) is True
assert dotprompt_manager.should_run_prompt_management("missing-prompt", None, {}) is False
def test_prompt_initializer_registers_flat_db_prompt_under_base_id():
from litellm.integrations.dotprompt import DotpromptManager, prompt_initializer
from litellm.types.prompts.init_prompts import (
PromptInfo,
PromptLiteLLMParams,
PromptSpec,
)
litellm_params = PromptLiteLLMParams(
prompt_integration="dotprompt",
prompt_data={"content": "AHOY {{name}}", "metadata": {"model": "gpt-4o-mini"}},
)
prompt_spec = PromptSpec(
prompt_id="agent-prompt.v1",
litellm_params=litellm_params,
prompt_info=PromptInfo(prompt_type="db"),
)
dotprompt_manager = prompt_initializer(litellm_params, prompt_spec)
assert isinstance(dotprompt_manager, DotpromptManager)
template = dotprompt_manager.prompt_manager.get_prompt("agent-prompt")
assert template is not None
assert template.content == "AHOY {{name}}"

146
tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py
View file

@ -246,3 +246,149 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404():
exc_info.value.detail
== "Prompt with ID test_prompt not found in environment development"
)
def test_is_ambiguous_keyed_prompt_data_shapes():
from litellm.proxy.prompts.prompt_endpoints import is_ambiguous_keyed_prompt_data
keyed_with_id = PromptLiteLLMParams(
prompt_id="agent-prompt",
prompt_integration="dotprompt",
prompt_data={"json_prompt": {"content": "AHOY", "metadata": {}}},
)
flat_with_id = PromptLiteLLMParams(
prompt_id="agent-prompt",
prompt_integration="dotprompt",
prompt_data={"content": "AHOY", "metadata": {}},
)
keyed_without_id = PromptLiteLLMParams(
prompt_integration="dotprompt",
prompt_data={"json_prompt": {"content": "AHOY", "metadata": {}}},
)
no_prompt_data = PromptLiteLLMParams(
prompt_id="agent-prompt", prompt_integration="dotprompt"
)
empty_prompt_data = PromptLiteLLMParams(
prompt_id="agent-prompt", prompt_integration="dotprompt", prompt_data={}
)
assert is_ambiguous_keyed_prompt_data(keyed_with_id) is True
assert is_ambiguous_keyed_prompt_data(flat_with_id) is False
assert is_ambiguous_keyed_prompt_data(keyed_without_id) is False
assert is_ambiguous_keyed_prompt_data(no_prompt_data) is False
assert is_ambiguous_keyed_prompt_data(empty_prompt_data) is False
@pytest.mark.asyncio
async def test_create_prompt_rejects_keyed_prompt_data_with_prompt_id():
from fastapi import HTTPException
from litellm.proxy.prompts.prompt_endpoints import (
AMBIGUOUS_PROMPT_DATA_ERROR,
Prompt,
create_prompt,
)
mock_user_auth = UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN
)
request = Prompt(
prompt_id="agent-prompt",
litellm_params=PromptLiteLLMParams(
prompt_id="agent-prompt",
prompt_integration="dotprompt",
prompt_data={"json_prompt": {"content": "AHOY", "metadata": {}}},
),
)
with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()): # test-quality-ok: proxy_server module global is the endpoint's only injection point
with pytest.raises(HTTPException) as exc_info:
await create_prompt(request=request, user_api_key_dict=mock_user_auth)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == AMBIGUOUS_PROMPT_DATA_ERROR
@pytest.mark.asyncio
async def test_update_prompt_rejects_keyed_prompt_data_with_prompt_id():
from fastapi import HTTPException
from litellm.proxy.prompts.prompt_endpoints import (
AMBIGUOUS_PROMPT_DATA_ERROR,
Prompt,
update_prompt,
)
mock_user_auth = UserAPIKeyAuth(
api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN
)
request = Prompt(
prompt_id="agent-prompt",
litellm_params=PromptLiteLLMParams(
prompt_id="agent-prompt",
prompt_integration="dotprompt",
prompt_data={"json_prompt": {"content": "AHOY", "metadata": {}}},
),
)
with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()): # test-quality-ok: proxy_server module global is the endpoint's only injection point
with pytest.raises(HTTPException) as exc_info:
await update_prompt(
prompt_id="agent-prompt",
request=request,
user_api_key_dict=mock_user_auth,
)
assert exc_info.value.status_code == 400
assert exc_info.value.detail == AMBIGUOUS_PROMPT_DATA_ERROR
def test_create_versioned_prompt_spec_populates_version():
from litellm.proxy.prompts.prompt_endpoints import create_versioned_prompt_spec
db_prompt = MagicMock()
db_prompt.model_dump.return_value = {
"prompt_id": "agent-prompt",
"version": 3,
"environment": "development",
"created_by": "user-1",
"litellm_params": {
"prompt_id": "agent-prompt",
"prompt_integration": "dotprompt",
},
"prompt_info": {"prompt_type": "db"},
"created_at": None,
"updated_at": None,
}
prompt_spec = create_versioned_prompt_spec(db_prompt=db_prompt)
assert prompt_spec.prompt_id == "agent-prompt.v3"
assert prompt_spec.version == 3
def test_initialize_prompt_keeps_version_and_created_by():
import litellm
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry
registry = InMemoryPromptRegistry()
prompt_spec = PromptSpec(
prompt_id="agent-prompt.v3",
litellm_params=PromptLiteLLMParams(
prompt_id="agent-prompt",
prompt_integration="dotprompt",
prompt_data={"content": "AHOY", "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
version=3,
environment="development",
created_by="user-1",
)
with patch.object(litellm.logging_callback_manager, "add_litellm_callback"): # test-quality-ok: keeps initialize_prompt from registering a global callback that would leak across tests
initialized_prompt = registry.initialize_prompt(prompt=prompt_spec)
assert initialized_prompt is not None
assert initialized_prompt.version == 3
assert initialized_prompt.created_by == "user-1"
assert initialized_prompt.environment == "development"