mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
7bac4a41af
LiteLLM_PromptTable is unique on (prompt_id, version, environment) and
version numbering restarts at 1 per environment, but the in-memory
registry keyed prompts as {prompt_id}.v{version} with no environment, so
environments sharing a prompt id shadowed each other and only one
environment's template ever served.
Registry entries are now keyed {versioned_id}::{environment}, and serve
time resolution goes through resolve_prompt_spec(base_id, version,
environment): production > staging > development when no environment is
requested, latest version within the chosen environment when no version
is requested. Chat requests can pin an environment with a new optional
prompt_environment body param, filtered from provider-bound params like
prompt_id and prompt_version. The newest-updated_at dedupe in
_init_prompts_in_db is dropped since registry keys can no longer
collide, and the key-parsing serve helpers plus dead registry getters
are removed
253 lines
9 KiB
Python
253 lines
9 KiB
Python
import json
|
|
import pytest
|
|
from unittest.mock import MagicMock
|
|
from litellm.types.prompts.init_prompts import (
|
|
PromptInfo,
|
|
PromptSpec,
|
|
PromptLiteLLMParams,
|
|
)
|
|
|
|
|
|
def test_prompt_info_default_environment():
|
|
"""PromptInfo should default environment to 'development'."""
|
|
info = PromptInfo(prompt_type="db")
|
|
assert info.environment == "development"
|
|
|
|
|
|
def test_prompt_info_custom_environment():
|
|
"""PromptInfo should accept a custom environment."""
|
|
info = PromptInfo(prompt_type="db", environment="production")
|
|
assert info.environment == "production"
|
|
|
|
|
|
def test_prompt_spec_includes_environment_and_created_by():
|
|
"""PromptSpec should carry environment and created_by fields."""
|
|
spec = PromptSpec(
|
|
prompt_id="test",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="test", prompt_integration="dotprompt"
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db", environment="staging"),
|
|
environment="staging",
|
|
created_by="user-123",
|
|
)
|
|
assert spec.environment == "staging"
|
|
assert spec.created_by == "user-123"
|
|
|
|
|
|
def test_prompt_spec_default_environment():
|
|
"""PromptSpec environment should default to 'development'."""
|
|
spec = PromptSpec(
|
|
prompt_id="test",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="test", prompt_integration="dotprompt"
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
)
|
|
assert spec.environment == "development"
|
|
assert spec.created_by is None
|
|
|
|
|
|
def test_create_versioned_prompt_spec_includes_environment():
|
|
"""create_versioned_prompt_spec should populate environment and created_by from DB row."""
|
|
from litellm.proxy.prompts.prompt_endpoints import create_versioned_prompt_spec
|
|
|
|
mock_db_prompt = MagicMock()
|
|
mock_db_prompt.model_dump.return_value = {
|
|
"id": "uuid-123",
|
|
"prompt_id": "test_prompt",
|
|
"version": 2,
|
|
"environment": "staging",
|
|
"created_by": "user-456",
|
|
"litellm_params": json.dumps(
|
|
{
|
|
"prompt_id": "test_prompt",
|
|
"prompt_integration": "dotprompt",
|
|
}
|
|
),
|
|
"prompt_info": json.dumps({"prompt_type": "db", "environment": "staging"}),
|
|
"created_at": None,
|
|
"updated_at": None,
|
|
}
|
|
spec = create_versioned_prompt_spec(mock_db_prompt)
|
|
assert spec.environment == "staging"
|
|
assert spec.created_by == "user-456"
|
|
assert spec.prompt_id == "test_prompt.v2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_prompt_stores_environment_and_created_by():
|
|
"""create_prompt should pass environment and created_by to the DB."""
|
|
from unittest.mock import AsyncMock, patch
|
|
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
|
|
from litellm.proxy.prompts.prompt_endpoints import create_prompt, Prompt
|
|
|
|
mock_user_auth = UserAPIKeyAuth(
|
|
api_key="sk-1234",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
user_id="user-789",
|
|
)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_db_entry = MagicMock()
|
|
mock_db_entry.model_dump.return_value = {
|
|
"id": "uuid-1",
|
|
"prompt_id": "my_prompt",
|
|
"version": 1,
|
|
"environment": "staging",
|
|
"created_by": "user-789",
|
|
"litellm_params": json.dumps(
|
|
{
|
|
"prompt_id": "my_prompt",
|
|
"prompt_integration": "dotprompt",
|
|
}
|
|
),
|
|
"prompt_info": json.dumps({"prompt_type": "db", "environment": "staging"}),
|
|
"created_at": None,
|
|
"updated_at": None,
|
|
}
|
|
mock_prisma_client.db.litellm_prompttable.create = AsyncMock(
|
|
return_value=mock_db_entry
|
|
)
|
|
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(return_value=[])
|
|
|
|
request = Prompt(
|
|
prompt_id="my_prompt",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="my_prompt", prompt_integration="dotprompt"
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db", environment="staging"),
|
|
)
|
|
|
|
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
|
with patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry:
|
|
mock_registry.initialize_prompt.return_value = PromptSpec(
|
|
prompt_id="my_prompt.v1",
|
|
litellm_params=request.litellm_params,
|
|
prompt_info=request.prompt_info,
|
|
environment="staging",
|
|
created_by="user-789",
|
|
)
|
|
await create_prompt(request=request, user_api_key_dict=mock_user_auth)
|
|
|
|
create_call = mock_prisma_client.db.litellm_prompttable.create.call_args
|
|
data = create_call.kwargs["data"]
|
|
assert data["environment"] == "staging"
|
|
assert data["created_by"] == "user-789"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_prompt_stores_environment_and_created_by():
|
|
"""update_prompt should pass environment and created_by to new version."""
|
|
from unittest.mock import AsyncMock, patch
|
|
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
|
|
from litellm.proxy.prompts.prompt_endpoints import update_prompt, Prompt
|
|
|
|
mock_user_auth = UserAPIKeyAuth(
|
|
api_key="sk-1234",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
user_id="user-update",
|
|
)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_existing = MagicMock()
|
|
mock_existing.version = 1
|
|
mock_prisma_client.db.litellm_prompttable.find_many = AsyncMock(
|
|
return_value=[mock_existing]
|
|
)
|
|
|
|
mock_db_entry = MagicMock()
|
|
mock_db_entry.model_dump.return_value = {
|
|
"id": "uuid-2",
|
|
"prompt_id": "my_prompt",
|
|
"version": 2,
|
|
"environment": "production",
|
|
"created_by": "user-update",
|
|
"litellm_params": json.dumps(
|
|
{
|
|
"prompt_id": "my_prompt",
|
|
"prompt_integration": "dotprompt",
|
|
}
|
|
),
|
|
"prompt_info": json.dumps({"prompt_type": "db", "environment": "production"}),
|
|
"created_at": None,
|
|
"updated_at": None,
|
|
}
|
|
mock_prisma_client.db.litellm_prompttable.create = AsyncMock(
|
|
return_value=mock_db_entry
|
|
)
|
|
|
|
request = Prompt(
|
|
prompt_id="my_prompt",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="my_prompt", prompt_integration="dotprompt"
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db", environment="production"),
|
|
)
|
|
|
|
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
|
with patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry:
|
|
mock_registry.has_config_prompt.return_value = False
|
|
mock_registry.initialize_prompt.return_value = PromptSpec(
|
|
prompt_id="my_prompt.v2",
|
|
litellm_params=request.litellm_params,
|
|
prompt_info=request.prompt_info,
|
|
environment="production",
|
|
created_by="user-update",
|
|
)
|
|
await update_prompt(
|
|
prompt_id="my_prompt", request=request, user_api_key_dict=mock_user_auth
|
|
)
|
|
|
|
create_call = mock_prisma_client.db.litellm_prompttable.create.call_args
|
|
data = create_call.kwargs["data"]
|
|
assert data["environment"] == "production"
|
|
assert data["created_by"] == "user-update"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_prompt_scoped_to_environment():
|
|
"""delete_prompt with environment param should scope deletion."""
|
|
from unittest.mock import AsyncMock, patch
|
|
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
|
|
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(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry:
|
|
prompt_spec = PromptSpec(
|
|
prompt_id="test_prompt.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="test_prompt", prompt_integration="dotprompt"
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
environment="staging",
|
|
)
|
|
mock_registry.resolve_prompt_spec.return_value = prompt_spec
|
|
mock_registry.has_config_prompt.return_value = False
|
|
|
|
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
|
|
await delete_prompt(
|
|
prompt_id="test_prompt",
|
|
user_api_key_dict=mock_user_auth,
|
|
environment="staging",
|
|
)
|
|
|
|
mock_prisma_client.db.litellm_prompttable.delete_many.assert_called_once_with(
|
|
where={"prompt_id": "test_prompt", "environment": "staging"}
|
|
)
|
|
mock_registry.delete_prompts_by_base_id.assert_called_once_with(
|
|
base_prompt_id="test_prompt", environment="staging"
|
|
)
|