diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index 425ff7572d0..9cfd6959a66 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -1021,19 +1021,7 @@ async def delete_prompt( # Delete versions from the database (scoped to environment if provided) await _prompt_table(prisma_client).delete_many(where=delete_where) - # Remove matching prompts from memory — scope to environment if provided - if environment: - prompts_to_delete: Final = [ - pid - for pid, prompt in IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.items() - if get_base_prompt_id(prompt_id=pid) == base_prompt_id and prompt.environment == environment - ] - for pid in prompts_to_delete: - del IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS[pid] - if pid in IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt: - del IN_MEMORY_PROMPT_REGISTRY.prompt_id_to_custom_prompt[pid] - else: - IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id) + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id(base_prompt_id, environment=environment or None) env_msg: Final = f" from {environment}" if environment else "" return {"message": f"Prompt {base_prompt_id} deleted successfully{env_msg}"} diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index d4342773a85..addfb3f80d5 100644 --- a/litellm/proxy/prompts/prompt_registry.py +++ b/litellm/proxy/prompts/prompt_registry.py @@ -195,12 +195,22 @@ class InMemoryPromptRegistry: """ return self.prompt_id_to_custom_prompt.get(prompt_id) - def delete_prompts_by_base_id(self, base_prompt_id: str) -> list[str]: + def remove_prompt(self, prompt_id: str) -> None: + import litellm + + self.IN_MEMORY_PROMPTS.pop(prompt_id, None) + stale_callback: Final = self.prompt_id_to_custom_prompt.pop(prompt_id, None) + if stale_callback is not None: + litellm.logging_callback_manager.remove_callback_from_all_lists(stale_callback) + + def delete_prompts_by_base_id(self, base_prompt_id: str, environment: str | None = None) -> list[str]: """ - Delete all prompts matching the given base prompt ID from memory. + Delete all prompts matching the given base prompt ID from memory, along with their + registered callbacks; scoped to one environment when given. Args: base_prompt_id: The base prompt ID (without version suffix) + environment: When set, only delete prompts deployed to this environment Returns: List of prompt IDs that were deleted @@ -208,13 +218,14 @@ class InMemoryPromptRegistry: from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id prompts_to_delete: Final = [ - pid for pid in self.IN_MEMORY_PROMPTS if get_base_prompt_id(prompt_id=pid) == base_prompt_id + pid + for pid, prompt in self.IN_MEMORY_PROMPTS.items() + if get_base_prompt_id(prompt_id=pid) == base_prompt_id + and (environment is None or prompt.environment == environment) ] for pid in prompts_to_delete: - del self.IN_MEMORY_PROMPTS[pid] - if pid in self.prompt_id_to_custom_prompt: - del self.prompt_id_to_custom_prompt[pid] + self.remove_prompt(prompt_id=pid) return prompts_to_delete diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1d0f0477767..990682f10a5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7268,6 +7268,7 @@ class ProxyConfig: return None try: + prompt_ids_loaded_before_db_read: Final = frozenset(IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS) prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many() parsed_specs: Final[tuple[PromptSpec, ...]] = tuple( spec for row in prompts_in_db if (spec := parse_row(row)) is not None @@ -7290,6 +7291,18 @@ class ProxyConfig: prompt_spec.prompt_id, prompt_sync_error, ) + # An unparsable row still exists in the DB, so skip the sweep rather than unload its in-memory copy + every_row_parsed: Final = len(parsed_specs) == len(prompts_in_db) + if every_row_parsed: + deleted_db_prompt_ids: Final = tuple( + prompt_id + for prompt_id in prompt_ids_loaded_before_db_read + if (loaded_spec := IN_MEMORY_PROMPT_REGISTRY.IN_MEMORY_PROMPTS.get(prompt_id)) is not None + and loaded_spec.prompt_info.prompt_type == "db" + and prompt_id not in newest_spec_per_id + ) + for deleted_prompt_id in deleted_db_prompt_ids: + IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id=deleted_prompt_id) 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 688b739fb5a..b5792ac7572 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -79,7 +79,7 @@ async def test_delete_prompt_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id + expected_base_id, environment=None ) assert response == { @@ -150,7 +150,7 @@ async def test_delete_prompt_by_base_id_success(): # 2. Memory deletion should use base ID mock_registry.delete_prompts_by_base_id.assert_called_once_with( - expected_base_id + expected_base_id, environment=None ) assert response == { @@ -158,6 +158,37 @@ async def test_delete_prompt_by_base_id_success(): } +@pytest.mark.asyncio +async def test_delete_prompt_environment_scope_reaches_db_and_registry(): + from litellm.proxy.prompts.prompt_endpoints import delete_prompt + + mock_user_auth = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN) + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_prompttable.delete_many = AsyncMock(return_value=None) + + with patch( # test-quality-ok: stubs the collaborator so the test pins what the endpoint deletes + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: + mock_registry.get_prompt_by_id.return_value = PromptSpec( + prompt_id="test_prompt.v2", + litellm_params=PromptLiteLLMParams(prompt_id="test_prompt", prompt_integration="dotprompt"), + 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 + response = await delete_prompt( + prompt_id="test_prompt.v2", + environment="production", + user_api_key_dict=mock_user_auth, + ) + + mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with( + where={"prompt_id": "test_prompt", "environment": "production"} + ) + mock_registry.delete_prompts_by_base_id.assert_called_once_with("test_prompt", environment="production") + assert response == {"message": "Prompt test_prompt deleted successfully from production"} + + @pytest.mark.asyncio async def test_get_prompt_info_by_base_id(): """ diff --git a/tests/test_litellm/proxy/prompts/test_prompt_registry.py b/tests/test_litellm/proxy/prompts/test_prompt_registry.py index 47f1ba13627..3008821974e 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_registry.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -88,3 +88,55 @@ def test_reload_prompt_keeps_the_old_template_when_the_replacement_fails(isolate 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] + + +def _versioned_prompt_spec(version: int, environment: str) -> PromptSpec: + return PromptSpec( + prompt_id=f"greeting.v{version}", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting", + prompt_integration="dotprompt", + prompt_data={"content": f"begin every reply with AHOY v{version}", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="db", environment=environment), + version=version, + environment=environment, + ) + + +def test_delete_prompts_by_base_id_removes_the_callbacks_from_litellm_callbacks(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) + registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "development")) + assert len(isolated_callbacks) == 1 + + deleted = registry.delete_prompts_by_base_id("greeting") + + assert sorted(deleted) == ["greeting.v1", "greeting.v2"] + assert registry.get_prompt_by_id("greeting.v1") is None + assert registry.get_prompt_callback_by_id("greeting.v2") is None + assert isolated_callbacks == [] + + +def test_delete_prompts_by_base_id_environment_scope_keeps_other_environments(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) + registry.initialize_prompt(prompt=_versioned_prompt_spec(2, "production")) + production_callback = registry.get_prompt_callback_by_id("greeting.v2") + + deleted = registry.delete_prompts_by_base_id("greeting", environment="development") + + assert deleted == ["greeting.v1"] + assert registry.get_prompt_by_id("greeting.v1") is None + assert registry.get_prompt_by_id("greeting.v2") is not None + assert registry.get_prompt_callback_by_id("greeting.v2") is production_callback + + +def test_remove_prompt_is_a_no_op_for_an_unknown_id(isolated_callbacks: list) -> None: + registry = InMemoryPromptRegistry() + registry.initialize_prompt(prompt=_versioned_prompt_spec(1, "development")) + + registry.remove_prompt(prompt_id="not_there.v1") + + assert registry.get_prompt_by_id("greeting.v1") is not None + 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 3ebaf951e26..f51648faf80 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11492,6 +11492,150 @@ async def test_init_prompts_in_db_serves_the_newest_row_when_environments_collid IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_env") +def _prompt_db_row(prompt_id: str, litellm_params: str) -> MagicMock: + row = MagicMock() + row.model_dump.return_value = { + "prompt_id": prompt_id, + "version": 1, + "environment": "development", + "created_by": None, + "litellm_params": litellm_params, + "prompt_info": json.dumps({"prompt_type": "db"}), + "created_at": None, + "updated_at": None, + } + return row + + +def _dotprompt_params(prompt_id: str) -> str: + return json.dumps( + { + "prompt_id": prompt_id, + "prompt_integration": "dotprompt", + "prompt_data": {"content": "Begin every reply with AHOY", "metadata": {}}, + } + ) + + +@pytest.mark.asyncio +async def test_init_prompts_in_db_unloads_rows_deleted_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", []) + + prisma_client = MagicMock() + try: + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[_prompt_db_row("greeting_del", _dotprompt_params("greeting_del"))] + ) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is not None + + prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[]) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_del.v1") is None + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_del.v1") is None + assert litellm.callbacks == [] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_del") + + +@pytest.mark.asyncio +async def test_init_prompts_in_db_keeps_config_prompts_when_their_id_has_no_db_row(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec + + monkeypatch.setattr(litellm, "callbacks", []) + + config_prompt = PromptSpec( + prompt_id="greeting_cfg", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting_cfg", + prompt_integration="dotprompt", + prompt_data={"content": "Begin every reply with AHOY", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="config"), + ) + + prisma_client = MagicMock() + try: + IN_MEMORY_PROMPT_REGISTRY.initialize_prompt(prompt=config_prompt) + prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[]) + + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_cfg") is not None + assert len(litellm.callbacks) == 1 + finally: + IN_MEMORY_PROMPT_REGISTRY.remove_prompt(prompt_id="greeting_cfg") + + +@pytest.mark.asyncio +async def test_init_prompts_in_db_keeps_the_in_memory_copy_when_a_row_fails_to_parse(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr(litellm, "callbacks", []) + + prisma_client = MagicMock() + try: + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[_prompt_db_row("greeting_broken", _dotprompt_params("greeting_broken"))] + ) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + loaded_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1") + assert loaded_callback is not None + + prisma_client.db.litellm_prompttable.find_many = AsyncMock( + return_value=[_prompt_db_row("greeting_broken", "this is not json")] + ) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_broken.v1") is loaded_callback + assert litellm.callbacks == [loaded_callback] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_broken") + + +@pytest.mark.asyncio +async def test_init_prompts_in_db_keeps_a_prompt_created_while_the_sync_was_reading(monkeypatch): + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.proxy_server import ProxyConfig + from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec + + monkeypatch.setattr(litellm, "callbacks", []) + + prisma_client = MagicMock() + try: + + async def create_prompt_behind_the_select() -> list: + IN_MEMORY_PROMPT_REGISTRY.initialize_prompt( + prompt=PromptSpec( + prompt_id="greeting_race.v1", + litellm_params=PromptLiteLLMParams( + prompt_id="greeting_race", + prompt_integration="dotprompt", + prompt_data={"content": "Begin every reply with AHOY", "metadata": {}}, + ), + prompt_info=PromptInfo(prompt_type="db"), + ) + ) + return [] + + prisma_client.db.litellm_prompttable.find_many = AsyncMock(side_effect=create_prompt_behind_the_select) + await ProxyConfig()._init_prompts_in_db(prisma_client=prisma_client) + + surviving_callback = IN_MEMORY_PROMPT_REGISTRY.get_prompt_callback_by_id("greeting_race.v1") + assert IN_MEMORY_PROMPT_REGISTRY.get_prompt_by_id("greeting_race.v1") is not None + assert surviving_callback is not None + assert litellm.callbacks == [surviving_callback] + finally: + IN_MEMORY_PROMPT_REGISTRY.delete_prompts_by_base_id("greeting_race") + + class TestEmbeddingsFailureHookRequestData: @pytest.mark.asyncio async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):