diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index e90adc163e2..425ff7572d0 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -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 diff --git a/litellm/proxy/prompts/prompt_registry.py b/litellm/proxy/prompts/prompt_registry.py index 46caa30fdab..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,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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c22504c555c..bd319148d95 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) 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 67113ed9145..688b739fb5a 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints_crud.py @@ -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", 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..47f1ba13627 --- /dev/null +++ b/tests/test_litellm/proxy/prompts/test_prompt_registry.py @@ -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] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index f6aebc9546d..4266a13bf11 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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):