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] 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):