mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
e9d40a8f73
A name bound twice keeps only the second binding. In `tests/` that is nearly always a repeated import, harmless but misleading, and the same rule is what catches the cases that are not harmless: a local that shadows an import the module still calls, and a second `def test_x` that quietly replaces the first. 311 of the 344 sites were repeated imports and came out with ruff's own fix. The remaining 33 needed a decision. Four modules imported a name they never used because a local definition below already shadowed it. Two comprehensions bound `call` over `unittest.mock.call`, which those modules import and use. One test rebound the two module handles its nested reload closure had captured. One class attribute shadowed an unused `status` import. The load-test fixtures move to a conftest, which is how pytest is meant to share them, so the test module no longer imports three fixture names it never calls. The nine `prisma_client` parameters keep a narrow `noqa`: pytest resolves that fixture by name before the body runs, so the parameter never shadows anything.
462 lines
16 KiB
Python
462 lines
16 KiB
Python
"""
|
|
Test prompt endpoints for version filtering and history
|
|
"""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm.types.prompts.init_prompts import (
|
|
PromptInfo,
|
|
PromptLiteLLMParams,
|
|
PromptSpec,
|
|
)
|
|
|
|
|
|
class TestPromptVersioning:
|
|
"""
|
|
Test prompt versioning functionality
|
|
"""
|
|
|
|
def test_get_latest_prompt_versions(self):
|
|
"""
|
|
Test that get_latest_prompt_versions returns only the latest version of each prompt
|
|
"""
|
|
from litellm.proxy.prompts.prompt_endpoints import get_latest_prompt_versions
|
|
|
|
# Create mock prompts with different versions
|
|
prompts = [
|
|
PromptSpec(
|
|
prompt_id="jack.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v1 content",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
PromptSpec(
|
|
prompt_id="jack.v2",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v2 content",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
PromptSpec(
|
|
prompt_id="jane.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jane",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="jane v1",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
PromptSpec(
|
|
prompt_id="jack.v3",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v3 content",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
]
|
|
|
|
# Get latest versions
|
|
latest = get_latest_prompt_versions(prompts=prompts)
|
|
|
|
# Should return 2 prompts (jack.v3 and jane.v1)
|
|
assert len(latest) == 2
|
|
|
|
# Find jack and jane in results
|
|
jack_prompt = next((p for p in latest if "jack" in p.prompt_id), None)
|
|
jane_prompt = next((p for p in latest if "jane" in p.prompt_id), None)
|
|
|
|
assert jack_prompt is not None
|
|
assert jack_prompt.prompt_id == "jack.v3"
|
|
assert jack_prompt.litellm_params.dotprompt_content == "v3 content"
|
|
|
|
assert jane_prompt is not None
|
|
assert jane_prompt.prompt_id == "jane.v1"
|
|
|
|
def test_get_version_number(self):
|
|
"""
|
|
Test that get_version_number correctly extracts version numbers
|
|
"""
|
|
from litellm.proxy.prompts.prompt_endpoints import get_version_number
|
|
|
|
assert get_version_number(prompt_id="jack.v1") == 1
|
|
assert get_version_number(prompt_id="jack.v2") == 2
|
|
assert get_version_number(prompt_id="jack.v10") == 10
|
|
assert get_version_number(prompt_id="jack") == 1
|
|
assert get_version_number(prompt_id="jack.vinvalid") == 1
|
|
|
|
def test_get_base_prompt_id(self):
|
|
"""
|
|
Test that get_base_prompt_id correctly strips version suffixes
|
|
"""
|
|
from litellm.proxy.prompts.prompt_endpoints import get_base_prompt_id
|
|
|
|
assert get_base_prompt_id(prompt_id="jack.v1") == "jack"
|
|
assert get_base_prompt_id(prompt_id="jack.v2") == "jack"
|
|
assert get_base_prompt_id(prompt_id="jack") == "jack"
|
|
assert get_base_prompt_id(prompt_id="my_prompt.v10") == "my_prompt"
|
|
|
|
def test_get_latest_version_prompt_id(self):
|
|
"""
|
|
Test that get_latest_version_prompt_id returns the highest version
|
|
"""
|
|
from litellm.proxy.prompts.prompt_endpoints import get_latest_version_prompt_id
|
|
|
|
# Mock prompt IDs dictionary
|
|
all_prompt_ids = {
|
|
"jack.v1": {},
|
|
"jack.v2": {},
|
|
"jack.v3": {},
|
|
"jane.v1": {},
|
|
"simple_prompt": {},
|
|
}
|
|
|
|
# Test with base prompt ID - should return latest version
|
|
assert (
|
|
get_latest_version_prompt_id(
|
|
prompt_id="jack", all_prompt_ids=all_prompt_ids
|
|
)
|
|
== "jack.v3"
|
|
)
|
|
|
|
# Test with versioned prompt ID - should still return latest version
|
|
assert (
|
|
get_latest_version_prompt_id(
|
|
prompt_id="jack.v1", all_prompt_ids=all_prompt_ids
|
|
)
|
|
== "jack.v3"
|
|
)
|
|
|
|
# Test with single version
|
|
assert (
|
|
get_latest_version_prompt_id(
|
|
prompt_id="jane", all_prompt_ids=all_prompt_ids
|
|
)
|
|
== "jane.v1"
|
|
)
|
|
|
|
# Test with non-versioned prompt
|
|
assert (
|
|
get_latest_version_prompt_id(
|
|
prompt_id="simple_prompt", all_prompt_ids=all_prompt_ids
|
|
)
|
|
== "simple_prompt"
|
|
)
|
|
|
|
# Test with non-existent prompt
|
|
assert (
|
|
get_latest_version_prompt_id(
|
|
prompt_id="nonexistent", all_prompt_ids=all_prompt_ids
|
|
)
|
|
== "nonexistent"
|
|
)
|
|
|
|
def test_construct_versioned_prompt_id(self):
|
|
"""
|
|
Test that construct_versioned_prompt_id correctly builds versioned IDs
|
|
"""
|
|
from litellm.proxy.prompts.prompt_endpoints import construct_versioned_prompt_id
|
|
|
|
# Test with base prompt ID and version
|
|
assert (
|
|
construct_versioned_prompt_id(prompt_id="jack_success", version=4)
|
|
== "jack_success.v4"
|
|
)
|
|
|
|
# Test with None version - should return base ID unchanged
|
|
assert (
|
|
construct_versioned_prompt_id(prompt_id="jack_success", version=None)
|
|
== "jack_success"
|
|
)
|
|
|
|
# Test with existing versioned ID - should replace version
|
|
assert (
|
|
construct_versioned_prompt_id(prompt_id="jack_success.v2", version=4)
|
|
== "jack_success.v4"
|
|
)
|
|
|
|
# Test with hyphenated prompt ID
|
|
assert (
|
|
construct_versioned_prompt_id(prompt_id="my-prompt", version=1)
|
|
== "my-prompt.v1"
|
|
)
|
|
|
|
# Test with double-digit version
|
|
assert (
|
|
construct_versioned_prompt_id(prompt_id="test_prompt", version=10)
|
|
== "test_prompt.v10"
|
|
)
|
|
|
|
|
|
class TestPromptVersionsEndpoint:
|
|
"""
|
|
Test the /prompts/{prompt_id}/versions endpoint
|
|
"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_prompt_versions_returns_all_versions(self):
|
|
"""
|
|
Test that get_prompt_versions returns all versions of a prompt sorted by version number
|
|
"""
|
|
from unittest.mock import patch
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.prompts.prompt_endpoints import get_prompt_versions
|
|
|
|
# Mock user with admin role
|
|
mock_user = UserAPIKeyAuth(
|
|
api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN
|
|
)
|
|
|
|
# Create mock prompt registry with multiple versions
|
|
mock_prompts = {
|
|
"jack.v1": PromptSpec(
|
|
prompt_id="jack.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v1",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jack.v2": PromptSpec(
|
|
prompt_id="jack.v2",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v2",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jack.v3": PromptSpec(
|
|
prompt_id="jack.v3",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v3",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jane.v1": PromptSpec(
|
|
prompt_id="jane.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jane",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="jane",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
}
|
|
|
|
# Force the in-memory path so this test is isolated from any leaked prisma mocks.
|
|
with (
|
|
patch("litellm.proxy.proxy_server.prisma_client", None),
|
|
patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry,
|
|
):
|
|
mock_registry.IN_MEMORY_PROMPTS = mock_prompts
|
|
|
|
# Test with base prompt ID
|
|
response = await get_prompt_versions(
|
|
prompt_id="jack", user_api_key_dict=mock_user
|
|
)
|
|
|
|
# Should return 3 versions of jack, sorted newest first
|
|
assert len(response.prompts) == 3
|
|
assert response.prompts[0].prompt_id == "jack"
|
|
assert response.prompts[0].version == 3
|
|
assert response.prompts[1].prompt_id == "jack"
|
|
assert response.prompts[1].version == 2
|
|
assert response.prompts[2].prompt_id == "jack"
|
|
assert response.prompts[2].version == 1
|
|
|
|
# Test with versioned prompt ID (should strip version)
|
|
response = await get_prompt_versions(
|
|
prompt_id="jack.v1", user_api_key_dict=mock_user
|
|
)
|
|
|
|
assert len(response.prompts) == 3
|
|
assert response.prompts[0].prompt_id == "jack"
|
|
assert response.prompts[0].version == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_prompt_versions_not_found(self):
|
|
"""
|
|
Test that get_prompt_versions raises 404 when prompt doesn't exist
|
|
"""
|
|
from unittest.mock import patch
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.prompts.prompt_endpoints import get_prompt_versions
|
|
|
|
mock_user = UserAPIKeyAuth(
|
|
api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN
|
|
)
|
|
|
|
with (
|
|
patch("litellm.proxy.proxy_server.prisma_client", None),
|
|
patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry,
|
|
):
|
|
mock_registry.IN_MEMORY_PROMPTS = {}
|
|
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
await get_prompt_versions(
|
|
prompt_id="nonexistent", user_api_key_dict=mock_user
|
|
)
|
|
|
|
assert exc_info.value.status_code == 404
|
|
assert "No versions found" in exc_info.value.detail
|
|
|
|
|
|
class TestAdminViewerReadAccess:
|
|
"""
|
|
proxy_admin_viewer has READ parity with proxy_admin on the prompt read endpoints
|
|
"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list_prompts_returns_all_prompts_for_admin_viewer(self):
|
|
"""A role without admin view falls through to the empty-list branch here."""
|
|
from unittest.mock import patch
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.prompts.prompt_endpoints import list_prompts
|
|
|
|
viewer = UserAPIKeyAuth(
|
|
api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
|
)
|
|
|
|
mock_prompts = {
|
|
"jack.v1": PromptSpec(
|
|
prompt_id="jack.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v1",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jack.v2": PromptSpec(
|
|
prompt_id="jack.v2",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v2",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jane.v1": PromptSpec(
|
|
prompt_id="jane.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jane",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="jane",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
}
|
|
|
|
with patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry:
|
|
mock_registry.IN_MEMORY_PROMPTS = mock_prompts
|
|
|
|
response = await list_prompts(user_api_key_dict=viewer)
|
|
|
|
assert sorted(p.prompt_id for p in response.prompts) == ["jack", "jane"]
|
|
jack = next(p for p in response.prompts if p.prompt_id == "jack")
|
|
assert jack.litellm_params.dotprompt_content == "v2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_prompt_versions_allows_admin_viewer(self):
|
|
"""Version history used to 403 anyone who was not exactly proxy_admin."""
|
|
from unittest.mock import patch
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.prompts.prompt_endpoints import get_prompt_versions
|
|
|
|
viewer = UserAPIKeyAuth(
|
|
api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
|
)
|
|
|
|
mock_prompts = {
|
|
"jack.v1": PromptSpec(
|
|
prompt_id="jack.v1",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v1",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
"jack.v2": PromptSpec(
|
|
prompt_id="jack.v2",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v2",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
),
|
|
}
|
|
|
|
with (
|
|
patch("litellm.proxy.proxy_server.prisma_client", None),
|
|
patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry,
|
|
):
|
|
mock_registry.IN_MEMORY_PROMPTS = mock_prompts
|
|
|
|
response = await get_prompt_versions(
|
|
prompt_id="jack", user_api_key_dict=viewer
|
|
)
|
|
|
|
assert [p.version for p in response.prompts] == [2, 1]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_prompt_info_allows_admin_viewer(self):
|
|
"""Prompt info used to 403 anyone who was not exactly proxy_admin."""
|
|
from unittest.mock import patch
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.prompts.prompt_endpoints import get_prompt_info
|
|
|
|
viewer = UserAPIKeyAuth(
|
|
api_key="test_key", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
|
|
)
|
|
|
|
with (
|
|
patch("litellm.proxy.proxy_server.prisma_client", None),
|
|
patch(
|
|
"litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY"
|
|
) as mock_registry,
|
|
):
|
|
mock_registry.get_prompt_by_id.return_value = PromptSpec(
|
|
prompt_id="jack.v2",
|
|
litellm_params=PromptLiteLLMParams(
|
|
prompt_id="jack",
|
|
prompt_integration="dotprompt",
|
|
dotprompt_content="v2",
|
|
),
|
|
prompt_info=PromptInfo(prompt_type="db"),
|
|
)
|
|
mock_registry.IN_MEMORY_PROMPTS = {"jack.v1": {}, "jack.v2": {}}
|
|
mock_registry.get_prompt_callback_by_id.return_value = None
|
|
|
|
response = await get_prompt_info(prompt_id="jack", user_api_key_dict=viewer)
|
|
|
|
assert response.prompt_spec.prompt_id == "jack"
|
|
assert response.prompt_spec.version == 2
|