From b687fe2b50ff6662d45d6b930ab0f8bfcca3bbcb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:03:32 -0700 Subject: [PATCH 1/3] fix(prompts): propagate PATCHed prompt templates to every worker and pod --- litellm/proxy/prompts/prompt_endpoints.py | 33 +++------ litellm/proxy/prompts/prompt_registry.py | 17 +++++ litellm/proxy/proxy_server.py | 2 +- .../prompts/test_prompt_endpoints_crud.py | 69 ++++++++++++++++++- .../proxy/prompts/test_prompt_registry.py | 67 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 48 +++++++++++++ 6 files changed, 207 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/proxy/prompts/test_prompt_registry.py diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index a289ed7cbfb..cebfd5022a8 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -1025,15 +1025,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 @@ -1123,25 +1116,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]] = { @@ -1165,7 +1148,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 diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 695bdabfe83..ec7c98a068a 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -155,6 +155,23 @@ class InMemoryPromptRegistry: return parsed_prompt + def reload_prompt(self, prompt: PromptSpec) -> PromptSpec | None: + import litellm + + 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) + return self.initialize_prompt(prompt=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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e55b254ab8e..8c08342f491 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7237,7 +7237,7 @@ class ProxyConfig: 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) + IN_MEMORY_PROMPT_REGISTRY.sync_prompt_from_db(prompt=prompt_spec) except Exception as e: verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py index 3e8e1e9dff8..c0be93b3dcc 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -1,3 +1,5 @@ +import json + import pytest from unittest.mock import MagicMock, AsyncMock, patch from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles @@ -8,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(): """ @@ -208,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( @@ -246,3 +267,45 @@ 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" ) + + +@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" diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py new file mode 100644 index 00000000000..0533a6c11a8 --- /dev/null +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -0,0 +1,67 @@ +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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index afc42e8db45..ec2b79908ec 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11195,6 +11195,54 @@ 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") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): From 6df307fef86fd07a73c4ec85f97cea5b62afd068 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:49:31 -0700 Subject: [PATCH 2/3] fix(prompts): validate a prompt replacement before swapping and isolate per-row sync failures --- litellm/proxy/prompts/prompt_registry.py | 40 +++++++++++------- litellm/proxy/proxy_server.py | 12 ++++-- .../proxy/prompts/test_prompt_registry.py | 23 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 42 +++++++++++++++++++ 4 files changed, 98 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 49d61ef70a2..d4342773a85 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -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,21 +160,20 @@ class InMemoryPromptRegistry: environment=prompt.environment, created_by=prompt.created_by, ) - - # 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 + return parsed_prompt, 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) - return self.initialize_prompt(prompt=prompt) + 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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 88ef62ecb2d..afc29255d58 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7242,9 +7242,15 @@ class ProxyConfig: 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.sync_prompt_from_db(prompt=prompt_spec) + try: + prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) + 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", + getattr(prompt, "prompt_id", None), + prompt_sync_error, + ) except Exception as e: verbose_proxy_logger.debug("litellm.proxy.proxy_server.py::ProxyConfig:_init_prompts_in_db - %s", e) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 0533a6c11a8..47f1ba13627 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -65,3 +65,26 @@ def test_reload_prompt_replaces_callback_without_leaking_the_old_one(isolated_ca 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] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7eea7e0652b..42cff844513 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11327,6 +11327,48 @@ async def test_init_prompts_in_db_reloads_rows_patched_on_another_worker(monkeyp 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") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self): From 5461bb3b48925a6a64e585c0be3ddb177b0ba707 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:00:24 -0700 Subject: [PATCH 3/3] fix(prompts): sync only the newest row when environments share a versioned prompt id --- litellm/proxy/proxy_server.py | 28 ++++++++-- tests/test_litellm/proxy/test_proxy_server.py | 51 +++++++++++++++++++ 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index afc29255d58..b1ffce9c15a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7239,16 +7239,38 @@ 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: + 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: - prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) 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", - getattr(prompt, "prompt_id", None), + prompt_spec.prompt_id, prompt_sync_error, ) except Exception as e: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 42cff844513..fbf71829abc 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11369,6 +11369,57 @@ async def test_init_prompts_in_db_syncs_remaining_rows_when_one_row_fails(monkey 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):