scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(prompts): accept a string prompt_version and carry the viewed environment into code snippets

Browse source
This commit is contained in:
mateo-berri 2026-08-29 12:55:34 -07:00
parent 588f30950c
commit 3088db3f0e
10 changed files with 139 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/proxy/prompts/prompt_registry.py
View file

@ -74,6 +74,16 @@ def registry_key_for_prompt(prompt: PromptSpec) -> str:
return f"{prompt.prompt_id}::{prompt_environment_or_default(prompt.environment)}"
def parse_prompt_version(raw_version: object) -> int | None:
if isinstance(raw_version, bool):
return None
if isinstance(raw_version, int):
return raw_version
if isinstance(raw_version, str) and raw_version.isdigit():
return int(raw_version)
return None
def _spec_version(prompt: PromptSpec) -> int:
return prompt.version if prompt.version is not None else get_version_number(prompt_id=prompt.prompt_id)

5
litellm/proxy/utils.py
View file

@ -1749,7 +1749,6 @@ class ProxyLogging:
litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None))
prompt_id: Final[str | None] = data.get("prompt_id", None)
prompt_version: Final[int | None] = data.get("prompt_version", None)
## PROMPT TEMPLATE CHECK ##
@ -1759,11 +1758,13 @@ class ProxyLogging:
and prompt_id is not None
and (call_type == "completion" or call_type == "acompletion" or call_type == "aresponses")
):
from litellm.proxy.prompts.prompt_registry import parse_prompt_version
await self._process_prompt_template(
data=data,
litellm_logging_obj=litellm_logging_obj,
prompt_id=prompt_id,
prompt_version=prompt_version,
prompt_version=parse_prompt_version(data.get("prompt_version", None)),
call_type=call_type,
)

10
tests/test_litellm/proxy/prompts/test_prompt_registry.py
View file

@ -2,7 +2,7 @@ import pytest
import litellm
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry
from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry, parse_prompt_version
from litellm.types.prompts.init_prompts import PromptInfo, PromptLiteLLMParams, PromptSpec
@ -214,3 +214,11 @@ def test_remove_prompt_is_a_no_op_for_an_unknown_registry_key(isolated_callbacks
assert registry.resolve_prompt_spec("greeting") is not None
assert len(isolated_callbacks) == 1
@pytest.mark.parametrize(
("raw_version", "expected"),
[(2, 2), ("2", 2), (None, None), ("v2", None), (True, None), (2.0, None)],
)
def test_parse_prompt_version_accepts_integers_and_json_strings(raw_version: object, expected: int | None) -> None:
assert parse_prompt_version(raw_version) == expected

35
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -858,6 +858,41 @@ async def test_process_prompt_template_resolves_the_requested_environment(proxy_
assert data["messages"] == [{"role": "user", "content": "rendered"}]
@pytest.mark.asyncio
async def test_pre_call_hook_matches_a_prompt_version_sent_as_a_json_string(proxy_logging, monkeypatch):
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.prompts import prompt_registry
prompt_spec = MagicMock()
prompt_spec.litellm_params = MagicMock(prompt_id="greeting")
resolve_calls: list[dict] = []
def fake_resolve(prompt_id, version=None, environment=None):
resolve_calls.append({"prompt_id": prompt_id, "version": version, "environment": environment})
return prompt_spec
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", fake_resolve)
monkeypatch.setattr(
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "get_prompt_callback_for_prompt", lambda *a, **kw: MagicMock()
)
logging_obj = MagicMock()
logging_obj.async_get_chat_completion_prompt = AsyncMock(
return_value=("m", [{"role": "user", "content": "rendered"}], {})
)
data: Dict[str, Any] = {
"messages": [{"role": "user", "content": "orig"}],
"model": "m",
"prompt_id": "greeting",
"prompt_version": "2",
"litellm_logging_obj": logging_obj,
}
result = await proxy_logging.pre_call_hook(user_api_key_dict=UserAPIKeyAuth(), data=data, call_type="completion")
assert resolve_calls == [{"prompt_id": "greeting", "version": 2, "environment": None}]
assert result["messages"] == [{"role": "user", "content": "rendered"}]
@pytest.mark.asyncio
async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(proxy_logging, monkeypatch):
from litellm.proxy.prompts import prompt_registry

24
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.test.tsx
View file

@ -44,4 +44,28 @@ describe("PromptCodeSnippets", () => {
expect(screen.getByRole("combobox", { name: "Language" })).toHaveTextContent("Python (OpenAI SDK)");
});
it("includes the viewed environment in every generated request", async () => {
const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
render(
<PromptCodeSnippets
promptId="welcome"
model="gpt-4o"
accessToken="token"
version="2"
environment="development"
/>,
);
await user.click(screen.getByRole("button", { name: /get code/i }));
await screen.findByText("Generated Code");
await user.click(screen.getByRole("button", { name: /copy to clipboard/i }));
expect(await navigator.clipboard.readText()).toContain('"prompt_environment": "development"');
await user.click(screen.getByRole("tab", { name: "With Version" }));
await user.click(screen.getByRole("button", { name: /copy to clipboard/i }));
const versionSnippet = await navigator.clipboard.readText();
expect(versionSnippet).toContain('"prompt_environment": "development"');
expect(versionSnippet).toContain('"prompt_version": 2');
});
});

29
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptCodeSnippets.tsx
View file

@ -22,6 +22,7 @@ interface PromptCodeSnippetsProps {
promptVariables?: Record<string, string>;
accessToken: string | null;
version?: string;
environment?: string;
proxySettings?: {
PROXY_BASE_URL?: string;
LITELLM_UI_API_DOC_BASE_URL?: string | null;
@ -34,6 +35,7 @@ const PromptCodeSnippets: React.FC<PromptCodeSnippetsProps> = ({
promptVariables = {},
accessToken,
version = "1",
environment,
proxySettings,
}) => {
const syntaxTheme = useSyntaxTheme(coy);
@ -64,6 +66,9 @@ const PromptCodeSnippets: React.FC<PromptCodeSnippetsProps> = ({
// Generate code based on selected language and tab
const generateCode = () => {
const hasVariables = Object.keys(promptVariables).length > 0;
const curlEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : "";
const pythonEnvironment = environment ? `,\n "prompt_environment": "${environment}"` : "";
const jsEnvironment = environment ? `,\n prompt_environment: "${environment}"` : "";
if (selectedLanguage === "curl") {
if (selectedTab === "basic") {
@ -72,7 +77,7 @@ const PromptCodeSnippets: React.FC<PromptCodeSnippetsProps> = ({
-H 'Authorization: Bearer ${effectiveApiKey}' \\
-d '{
"model": "${model}",
"prompt_id": "${promptId}"${
"prompt_id": "${promptId}"${curlEnvironment}${
hasVariables
? `,
"prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}`
@ -85,7 +90,7 @@ const PromptCodeSnippets: React.FC<PromptCodeSnippetsProps> = ({
-H 'Authorization: Bearer ${effectiveApiKey}' \\
-d '{
"model": "${model}",
"prompt_id": "${promptId}"${
"prompt_id": "${promptId}"${curlEnvironment}${
hasVariables
? `,
"prompt_variables": ${JSON.stringify(promptVariables, null, 6).replace(/\n/g, "\n ")}`
@ -104,7 +109,7 @@ const PromptCodeSnippets: React.FC<PromptCodeSnippetsProps> = ({
-H 'Authorization: Bearer ${effectiveApiKey}' \\
-d '{
"model": "${model}",
"prompt_id": "${promptId}",
"prompt_id": "${promptId}"${curlEnvironment},
"prompt_version": ${version},
"messages": [
{
@ -127,7 +132,7 @@ client = openai.OpenAI(
response = client.chat.completions.create(
model="${model}",
extra_body={
"prompt_id": "${promptId}"${
"prompt_id": "${promptId}"${pythonEnvironment}${
hasVariables
? `,
"prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}`
@ -145,7 +150,7 @@ response = client.chat.completions.create(
{"role": "user", "content": "hi"}
],
extra_body={
"prompt_id": "${promptId}"${
"prompt_id": "${promptId}"${pythonEnvironment}${
hasVariables
? `,
"prompt_variables": ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}`
@ -163,7 +168,7 @@ response = client.chat.completions.create(
{"role": "user", "content": "Who are u"}
],
extra_body={
"prompt_id": "${promptId}",
"prompt_id": "${promptId}"${pythonEnvironment},
"prompt_version": ${version}
}
)
@ -186,9 +191,9 @@ async function main() {
model: "${model}",
${
hasVariables
? `prompt_id: "${promptId}",
? `prompt_id: "${promptId}"${jsEnvironment},
prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}`
: `prompt_id: "${promptId}"`
: `prompt_id: "${promptId}"${jsEnvironment}`
}
});
@ -206,9 +211,9 @@ async function main() {
],
${
hasVariables
? `prompt_id: "${promptId}",
? `prompt_id: "${promptId}"${jsEnvironment},
prompt_variables: ${JSON.stringify(promptVariables, null, 8).replace(/\n/g, "\n ")}`
: `prompt_id: "${promptId}"`
: `prompt_id: "${promptId}"${jsEnvironment}`
}
});
@ -224,7 +229,7 @@ async function main() {
messages: [
{ role: "user", content: "Who are u" }
],
prompt_id: "${promptId}",
prompt_id: "${promptId}"${jsEnvironment},
prompt_version: ${version}
});
@ -241,7 +246,7 @@ main();`;
if (isModalVisible) {
setGeneratedCode(generateCode());
}
}, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables]);
}, [isModalVisible, selectedLanguage, selectedTab, promptId, model, promptVariables, version, environment]);
return (
<>

5
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.test.tsx
View file

@ -2,7 +2,9 @@ import { fireEvent, render, screen } from "@testing-library/react";
import { describe, expect, it, vi } from "vitest";
import PromptEditorHeader from "./PromptEditorHeader";
vi.mock("./PromptCodeSnippets", () => ({ default: () => <button>Get Code</button> }));
vi.mock("./PromptCodeSnippets", () => ({
default: ({ environment }: { environment?: string }) => <button data-environment={environment}>Get Code</button>,
}));
describe("PromptEditorHeader", () => {
it("preserves navigation, naming, and save actions", () => {
@ -48,5 +50,6 @@ describe("PromptEditorHeader", () => {
);
expect(screen.getByRole("combobox", { name: "Environment" })).toHaveTextContent(label);
expect(screen.getByRole("button", { name: "Get Code" })).toHaveAttribute("data-environment", environment);
});
});

1
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_editor_view/PromptEditorHeader.tsx
View file

@ -89,6 +89,7 @@ const PromptEditorHeader: React.FC<PromptEditorHeaderProps> = ({
promptVariables={promptVariables}
accessToken={accessToken}
version={version?.replace("v", "") || "1"}
environment={environment}
proxySettings={proxySettings}
/>
{editMode && onShowHistory && (

36
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.test.tsx
View file

@ -12,7 +12,9 @@ vi.mock("@/components/networking", () => ({
}));
vi.mock("./prompt_editor_view/PromptCodeSnippets", () => ({
default: () => <div data-testid="prompt-code-snippets" />,
default: ({ environment }: { environment?: string }) => (
<div data-testid="prompt-code-snippets" data-environment={environment} />
),
}));
const promptWithoutTemplate = {
@ -58,6 +60,38 @@ describe("PromptInfoView environment scoping", () => {
});
});
describe("PromptInfoView code snippets", () => {
beforeEach(() => {
vi.mocked(networking.getPromptVersions).mockReset().mockResolvedValue({ prompts: [] });
});
it.each([
["a prompt with several environments", "staging", ["development", "staging"]],
["a config prompt with no environment list", "development", []],
])("hands the viewed environment of %s to the code snippets", async (_label, environment, environments) => {
vi.mocked(networking.getPromptInfo)
.mockReset()
.mockResolvedValue({
...promptWithoutTemplate,
prompt_spec: { ...promptWithoutTemplate.prompt_spec, environment },
environments,
});
render(
<PromptInfoView
promptId="support-reply"
initialEnvironment={environment}
onClose={vi.fn()}
accessToken="sk-test"
isAdmin={true}
/>,
);
await screen.findByRole("tab", { name: "Raw JSON" });
expect(screen.getByTestId("prompt-code-snippets")).toHaveAttribute("data-environment", environment);
});
});
describe("PromptInfoView tabs", () => {
beforeEach(() => {
vi.mocked(networking.getPromptInfo).mockReset().mockResolvedValue(promptWithoutTemplate);

1
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/prompt_info.tsx
View file

@ -221,6 +221,7 @@ const PromptInfoView: React.FC<PromptInfoProps> = ({
promptVariables={extractTemplateVariables(promptTemplate?.content)}
accessToken={accessToken}
version={currentVersion}
environment={selectedEnv ?? promptData.environment}
/>
<Button onClick={() => onEdit?.(rawApiResponse)} className="flex items-center">
<Pencil />