scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #38411 from BerriAI/litellm_fix_prompt_patch_sync

Browse source 
fix(prompts): propagate PATCHed prompt templates to every worker and pod
This commit is contained in:
Mateo Wang 2026-08-26 17:20:27 -07:00 • committed by GitHub
commit 53a607e088
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 381 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1045,15 +1045,8 @@ async def delete_prompt(
raise HTTPException(status_code=500, detail=str(e))
def _reload_prompt_in_registry(
registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec
) -> PromptSpec:
"""Remove stale entry and re-initialize the prompt in the in-memory registry."""
if versioned_id in registry.IN_MEMORY_PROMPTS:
del registry.IN_MEMORY_PROMPTS[versioned_id]
if versioned_id in registry.prompt_id_to_custom_prompt:
del registry.prompt_id_to_custom_prompt[versioned_id]
initialized: Final = registry.initialize_prompt(prompt=updated_prompt_spec, config_file_path=None)
def _reload_prompt_in_registry(registry: "InMemoryPromptRegistry", updated_prompt_spec: PromptSpec) -> PromptSpec:
initialized: Final = registry.reload_prompt(prompt=updated_prompt_spec)
if initialized is None:
raise HTTPException(status_code=500, detail="Failed to patch prompt")
return initialized
@ -1146,25 +1139,15 @@ async def patch_prompt(
detail="Cannot update config prompts.",
)
# Use existing prompt from memory or build from DB row for field merging
if existing_prompt:
current_litellm_params = existing_prompt.litellm_params
current_prompt_info = existing_prompt.prompt_info
else:
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
current_litellm_params = current_spec.litellm_params
current_prompt_info = current_spec.prompt_info
current_spec: Final = create_versioned_prompt_spec(db_prompt=target_row)
# Update fields if provided
updated_litellm_params: Final = (
request.litellm_params if request.litellm_params is not None else current_litellm_params
request.litellm_params if request.litellm_params is not None else current_spec.litellm_params
)
updated_prompt_info: Final = request.prompt_info if request.prompt_info is not None else current_prompt_info
# Ensure we have valid litellm_params
if updated_litellm_params is None:
raise HTTPException(status_code=400, detail="litellm_params cannot be None")
updated_prompt_info: Final = (
request.prompt_info if request.prompt_info is not None else current_spec.prompt_info
)
# Build update data dict
update_data: Final[dict[str, str]] = {
@ -1188,7 +1171,7 @@ async def patch_prompt(
updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry)
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec)
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, updated_prompt_spec)
except HTTPException as e:
raise e

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

@ -118,7 +118,16 @@ class InMemoryPromptRegistry:
verbose_proxy_logger.debug("prompt_id already exists in IN_MEMORY_PROMPTS")
return self.IN_MEMORY_PROMPTS[prompt_id]
custom_prompt_callback: CustomPromptManagement | None = None
parsed_prompt, custom_prompt_callback = self._build_prompt_callback(prompt=prompt)
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
# store references to the prompt in memory
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
return parsed_prompt
def _build_prompt_callback(self, prompt: PromptSpec) -> tuple[PromptSpec, CustomPromptManagement]:
litellm_params_data: Final = prompt.litellm_params
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
@ -132,17 +141,17 @@ class InMemoryPromptRegistry:
raise ValueError("prompt_integration is required")
initializer: Final = prompt_initializer_registry.get(prompt_integration)
if initializer:
custom_prompt_callback = initializer(litellm_params, prompt)
if not isinstance(custom_prompt_callback, CustomPromptManagement):
raise ValueError(f"CustomPromptManagement is required, got {type(custom_prompt_callback)}")
litellm.logging_callback_manager.add_litellm_callback(custom_prompt_callback)
else:
if initializer is None:
raise ValueError(f"Unsupported prompt: {prompt_integration}")
custom_prompt_callback: Final = initializer(litellm_params, prompt)
if not isinstance(custom_prompt_callback, CustomPromptManagement):
raise ValueError( # noqa: TRY004 # prompt endpoints map ValueError to HTTP 400; keep the existing contract
f"CustomPromptManagement is required, got {type(custom_prompt_callback)}"
)
parsed_prompt: Final = PromptSpec(
prompt_id=prompt_id,
prompt_id=prompt.prompt_id,
litellm_params=litellm_params,
prompt_info=prompt.prompt_info or PromptInfo(prompt_type="config"),
created_at=prompt.created_at,
@ -151,13 +160,29 @@ class InMemoryPromptRegistry:
environment=prompt.environment,
created_by=prompt.created_by,
)
return parsed_prompt, custom_prompt_callback
# store references to the prompt in memory
self.IN_MEMORY_PROMPTS[prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt_id] = custom_prompt_callback
def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None:
import litellm
parsed_prompt, new_callback = self._build_prompt_callback(prompt=prompt)
stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt.prompt_id, None)
self.IN_MEMORY_PROMPTS.pop(prompt.prompt_id, None)
if stale_callback is not None:
litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback)
litellm.logging_callback_manager.add_litellm_callback(new_callback)
self.IN_MEMORY_PROMPTS[prompt.prompt_id] = parsed_prompt
self.prompt_id_to_custom_prompt[prompt.prompt_id] = new_callback
return parsed_prompt
def sync_prompt_from_db(self, prompt: PromptSpec) -> PromptSpec | None:
existing: Final = self.IN_MEMORY_PROMPTS.get(prompt.prompt_id)
if existing is None:
return self.initialize_prompt(prompt=prompt)
if existing.litellm_params == prompt.litellm_params and existing.prompt_info == prompt.prompt_info:
return existing
return self.reload_prompt(prompt=prompt)
def get_prompt_by_id(self, prompt_id: str) -> PromptSpec | None:
"""
Get a prompt by its ID from memory

36
litellm/proxy/proxy_server.py
View file

@ -7257,12 +7257,40 @@ class ProxyConfig:
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.types.prompts.init_prompts import PromptSpec
def parse_row(db_prompt: object) -> PromptSpec | None:
try:
return self._get_prompt_spec_for_db_prompt(db_prompt=db_prompt)
except Exception as row_error: # noqa: BLE001 # a malformed row must not block syncing the remaining prompts
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to parse prompt row %s: %s",
getattr(db_prompt, "prompt_id", None),
row_error,
)
return None
try:
prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many()
for prompt in prompts_in_db:
# Convert DB object to dict and create versioned prompt_id
prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt)
IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=prompt_spec)
parsed_specs: Final[tuple[PromptSpec, ...]] = tuple(
spec for row in prompts_in_db if (spec := parse_row(row)) is not None
)
newest_spec_per_id: Final[Mapping[str, PromptSpec]] = MappingProxyType(
{
spec.prompt_id: spec
for spec in sorted(
parsed_specs,
key=lambda s: s.updated_at.timestamp() if s.updated_at else float("-inf"),
)
}
)
for prompt_spec in newest_spec_per_id.values():
try:
IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec)
except Exception as prompt_sync_error: # noqa: BLE001 # one poisoned row must not block syncing the remaining prompts
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - failed to sync prompt %s: %s",
prompt_spec.prompt_id,
prompt_sync_error,
)
except Exception as e:
verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e)

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

@ -1,4 +1,5 @@
import json
import pytest
from unittest.mock import MagicMock, AsyncMock, patch
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
@ -9,6 +10,27 @@ from litellm.types.prompts.init_prompts import (
)
def _db_row(content: str) -> MagicMock:
row = MagicMock()
row.id = "row-1"
row.version = 1
row.model_dump.return_value = {
"prompt_id": "test_prompt",
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": {
"prompt_id": "test_prompt",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
},
"prompt_info": {"prompt_type": "db"},
"created_at": None,
"updated_at": None,
}
return row
@pytest.mark.asyncio
async def test_delete_prompt_success():
"""
@ -209,9 +231,7 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404():
api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN
)
target_row = MagicMock()
target_row.id = "row-1"
target_row.version = 1
target_row = _db_row("Begin every reply with AHOY")
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(
@ -249,6 +269,48 @@ async def test_patch_prompt_row_deleted_mid_update_returns_404():
)
@pytest.mark.asyncio
async def test_patch_prompt_merges_unsent_fields_from_db_row_not_stale_memory():
from litellm.proxy.prompts.prompt_endpoints import PatchPromptRequest, patch_prompt
mock_user_auth = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
db_row = _db_row("Begin every reply with HOWDY")
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row])
mock_prisma_client.db.litellm_prompttable.update = AsyncMock(return_value=db_row)
stale_in_memory = PromptSpec(
prompt_id="test_prompt.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="test_prompt",
prompt_integration="dotprompt",
prompt_data={"content": "Begin every reply with AHOY", "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point
patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint writes and reloads
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
) as mock_registry,
):
mock_registry.get_prompt_by_id.return_value = stale_in_memory
mock_registry.reload_prompt.side_effect = lambda prompt: prompt
response = await patch_prompt(
prompt_id="test_prompt",
request=PatchPromptRequest(prompt_info=PromptInfo(prompt_type="db")),
user_api_key_dict=mock_user_auth,
)
written_params = json.loads(mock_prisma_client.db.litellm_prompttable.update.call_args.kwargs["data"]["litellm_params"])
assert written_params["prompt_data"]["content"] == "Begin every reply with HOWDY"
reloaded_spec = mock_registry.reload_prompt.call_args.kwargs["prompt"]
assert reloaded_spec.prompt_id == "test_prompt.v1"
assert reloaded_spec.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY"
assert response.litellm_params.prompt_data["content"] == "Begin every reply with HOWDY"
def test_is_ambiguous_keyed_prompt_data_shapes():
from litellm.proxy.prompts.prompt_endpoints import is_ambiguous_keyed_prompt_data
@ -358,6 +420,14 @@ async def test_patch_prompt_info_only_keeps_legacy_keyed_row_patchable():
target_row = MagicMock()
target_row.id = "row-1"
target_row.version = 1
target_row.model_dump.return_value = {
"prompt_id": "agent-prompt",
"version": 1,
"environment": "production",
"created_by": None,
"litellm_params": legacy_params.model_dump_json(),
"prompt_info": PromptInfo(prompt_type="db", environment="production").model_dump_json(),
}
updated_row = MagicMock()
updated_row.model_dump.return_value = {
"prompt_id": "agent-prompt",

90
tests/test_litellm/proxy/prompts/test_prompt_registry.py Normal file
View file

@ -0,0 +1,90 @@
import pytest
import litellm
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry
from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec
def _db_prompt_spec(content: str) -> PromptSpec:
return PromptSpec(
prompt_id="greeting.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="greeting",
prompt_integration="dotprompt",
prompt_data={"content": content, "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
def _served_content(registry: InMemoryPromptRegistry) -> str:
callback = registry.get_prompt_callback_by_id("greeting.v1")
assert callback is not None
return callback.prompt_manager.get_prompt("greeting").content
@pytest.fixture
def isolated_callbacks(monkeypatch: pytest.MonkeyPatch) -> list:
monkeypatch.setattr(litellm, "callbacks", [])
return litellm.callbacks
def test_sync_prompt_from_db_reloads_row_edited_elsewhere(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
stale_callback = registry.get_prompt_callback_by_id("greeting.v1")
assert _served_content(registry) == "begin every reply with AHOY"
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with HOWDY"))
assert _served_content(registry) == "begin every reply with HOWDY"
assert registry.get_prompt_by_id("greeting.v1").litellm_params.prompt_data["content"] == "begin every reply with HOWDY"
assert stale_callback not in isolated_callbacks
assert isolated_callbacks == [registry.get_prompt_callback_by_id("greeting.v1")]
def test_sync_prompt_from_db_keeps_unchanged_row_in_place(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
first_callback = registry.get_prompt_callback_by_id("greeting.v1")
registry.sync_prompt_from_db(prompt=_db_prompt_spec("begin every reply with AHOY"))
assert registry.get_prompt_callback_by_id("greeting.v1") is first_callback
assert isolated_callbacks == [first_callback]
def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY"))
stale_callback = registry.get_prompt_callback_by_id("greeting.v1")
reloaded = registry.reload_prompt(prompt=_db_prompt_spec("begin every reply with HOWDY"))
assert reloaded is not None
assert _served_content(registry) == "begin every reply with HOWDY"
assert stale_callback not in isolated_callbacks
assert len(isolated_callbacks) == 1
def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolated_callbacks: list) -> None:
registry = InMemoryPromptRegistry()
registry.initialize_prompt(prompt=_db_prompt_spec("begin every reply with AHOY"))
old_callback = registry.get_prompt_callback_by_id("greeting.v1")
broken = PromptSpec(
prompt_id="greeting.v1",
litellm_params=PromptLiteLLMParams(
prompt_id="greeting",
prompt_integration="does_not_exist",
prompt_data={"content": "begin every reply with HOWDY", "metadata": {}},
),
prompt_info=PromptInfo(prompt_type="db"),
)
with pytest.raises(ValueError, match="Unsupported prompt"):
registry.reload_prompt(prompt=broken)
assert registry.get_prompt_callback_by_id("greeting.v1") is old_callback
assert _served_content(registry) == "begin every reply with AHOY"
assert isolated_callbacks == [old_callback]

141
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -11351,6 +11351,147 @@ async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_re
assert not GUARDRAIL_RECONCILE_LOCK.locked()
@pytest.mark.asyncio
async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(content: str) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": "greeting_sync",
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": "greeting_sync",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": None,
}
return row
def served_content() -> str:
callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")
assert callback is not None
return callback.prompt_manager.get_prompt("greeting_sync").content
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with AHOY")])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert served_content() == "Begin every reply with AHOY"
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[db_row("Begin every reply with HOWDY")])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert served_content() == "Begin every reply with HOWDY"
assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_sync.v1")]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_sync")
@pytest.mark.asyncio
async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(prompt_id: str, integration: str) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": prompt_id,
"version": 1,
"environment": "development",
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": prompt_id,
"prompt_integration": integration,
"prompt_data": {"content": "Begin every reply with AHOY", "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": None,
}
return row
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(
return_value=[db_row("broken_sync", "does_not_exist"), db_row("healthy_sync", "dotprompt")]
)
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("broken_sync.v1") is None
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1") is not None
assert litellm.callbacks == [IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("healthy_sync.v1")]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("healthy_sync")
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("broken_sync")
@pytest.mark.asyncio
async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collide_on_a_versioned_id(monkeypatch):
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
from litellm.proxy.proxy_server import ProxyConfig
monkeypatch.setattr(litellm, "callbacks", [])
def db_row(environment: str, content: str, updated_at: datetime) -> MagicMock:
row = MagicMock()
row.model_dump.return_value = {
"prompt_id": "greeting_env",
"version": 1,
"environment": environment,
"created_by": None,
"litellm_params": json.dumps(
{
"prompt_id": "greeting_env",
"prompt_integration": "dotprompt",
"prompt_data": {"content": content, "metadata": {}},
}
),
"prompt_info": json.dumps({"prompt_type": "db"}),
"created_at": None,
"updated_at": updated_at,
}
return row
freshly_patched = db_row(
"production", "Begin every reply with HOWDY", datetime(2026, 8, 26, 12, 0, tzinfo=timezone.utc)
)
stale_sibling = db_row(
"development", "Begin every reply with AHOY", datetime(2026, 8, 26, 11, 0, tzinfo=timezone.utc)
)
prisma_client = MagicMock()
try:
prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[freshly_patched, stale_sibling])
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
first_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1")
assert first_callback is not None
assert first_callback.prompt_manager.get_prompt("greeting_env").content == "Begin every reply with HOWDY"
await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client)
assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_env.v1") is first_callback
assert litellm.callbacks == [first_callback]
finally:
IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env")
class TestEmbeddingsFailureHookRequestData:
@pytest.mark.asyncio
async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):