fix(prompts): propagate PATCHed prompt templates to every worker and pod

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

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

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)
 

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"

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

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