mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge pull request #38434 from BerriAI/litellm_propagate_prompt_deletes
fix(prompts): propagate prompt deletes to every worker and pod
This commit is contained in:
Mateo Wang
GitHub
commit
aedaf4d0b0
6 changed files with 260 additions and 21 deletions
|
|
@ -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}"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue