From d5fc63f63f2b3437093481dee1b897fa2f04ba6c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 12 Mar 2026 03:19:12 +0000 Subject: [PATCH] fix(ci): fix deprecated model refs and schema validation in unit tests - Replace gemini-pro with gemini-3-pro-preview in test_cost_discount_vertex_ai (gemini-pro removed from cost map) - Replace github/claude-3-5-sonnet-latest with github/claude-3-7-sonnet-20250219 in test_supports_function_calling_github_anthropic_alias (model removed) - Add supports_multimodal, uses_embed_content, input/output_cost_per_token_above_256k_tokens to JSON schema in test_utils.py (new properties added to model cost map) Co-authored-by: yuneng-jiang --- tests/test_litellm/test_cost_calculator.py | 8 +- tests/test_litellm/test_utils.py | 6 +- .../guardrail_info_helpers.test.tsx | 202 ++++++++++++++++++ 3 files changed, 211 insertions(+), 5 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.test.tsx diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index b991dcaf4e..1204b11934 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -970,12 +970,12 @@ def test_cost_discount_vertex_ai(): # Save original config original_discount_config = litellm.cost_discount_config.copy() - # Create mock response + # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( id="test-id", choices=[], created=1234567890, - model="gemini-pro", + model="gemini-3-pro-preview", object="chat.completion", usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), ) @@ -984,7 +984,7 @@ def test_cost_discount_vertex_ai(): litellm.cost_discount_config = {} cost_without_discount = completion_cost( completion_response=response, - model="vertex_ai/gemini-pro", + model="vertex_ai/gemini-3-pro-preview", custom_llm_provider="vertex_ai", ) @@ -994,7 +994,7 @@ def test_cost_discount_vertex_ai(): # Calculate cost with discount cost_with_discount = completion_cost( completion_response=response, - model="vertex_ai/gemini-pro", + model="vertex_ai/gemini-3-pro-preview", custom_llm_provider="vertex_ai", ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index fec516336f..64488e2fb6 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -90,7 +90,7 @@ def test_supports_function_calling_github_openai_alias(): def test_supports_function_calling_github_anthropic_alias(): assert ( litellm.utils.supports_function_calling( - model="github/claude-3-5-sonnet-latest" + model="github/claude-3-7-sonnet-20250219" ) is True ) @@ -619,6 +619,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_image_above_128k_tokens": {"type": "number"}, "input_cost_per_image_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, "cache_read_input_token_cost_flex": {"type": "number"}, "cache_read_input_token_cost_priority": {"type": "number"}, @@ -700,6 +701,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_token": {"type": "number"}, "output_cost_per_token_above_128k_tokens": {"type": "number"}, "output_cost_per_token_above_200k_tokens": {"type": "number"}, + "output_cost_per_token_above_256k_tokens": {"type": "number"}, "output_cost_per_token_above_272k_tokens": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels": {"type": "number"}, "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": { @@ -738,6 +740,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, "supports_url_context": {"type": "boolean"}, + "supports_multimodal": {"type": "boolean"}, + "uses_embed_content": {"type": "boolean"}, "supports_reasoning": {"type": "boolean"}, "supports_none_reasoning_effort": {"type": "boolean"}, "supports_xhigh_reasoning_effort": {"type": "boolean"}, diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.test.tsx new file mode 100644 index 0000000000..d9e01acaad --- /dev/null +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.test.tsx @@ -0,0 +1,202 @@ +import { describe, expect, it, beforeEach } from "vitest"; +import { + populateGuardrailProviders, + populateGuardrailProviderMap, + getGuardrailProviders, + shouldRenderPIIConfigSettings, + shouldRenderContentFilterConfigSettings, + shouldRenderAzureTextModerationConfigSettings, + getGuardrailLogoAndName, + DynamicGuardrailProviders, + guardrail_provider_map, + GuardrailProviders, +} from "./guardrail_info_helpers"; + +describe("guardrail_info_helpers", () => { + // Reset mutable module state between tests + beforeEach(() => { + // Clear DynamicGuardrailProviders by repopulating with empty + Object.keys(DynamicGuardrailProviders).forEach( + (key) => delete DynamicGuardrailProviders[key] + ); + // Remove any dynamically added keys from guardrail_provider_map + const staticKeys = new Set([ + "PresidioPII", + "Bedrock", + "Lakera", + "LitellmContentFilter", + "ToolPermission", + "BlockCodeExecution", + ]); + Object.keys(guardrail_provider_map).forEach((key) => { + if (!staticKeys.has(key)) delete guardrail_provider_map[key]; + }); + }); + + describe("populateGuardrailProviders", () => { + it("should populate dynamic providers from API response while preserving legacy providers", () => { + const apiResponse = { + zscaler_ai_guard: { + ui_friendly_name: "Zscaler AI Guard", + some_param: { required: true }, + }, + aporia_ai: { + ui_friendly_name: "Aporia AI", + }, + }; + + const result = populateGuardrailProviders(apiResponse); + + // Legacy providers preserved + expect(result.PresidioPII).toBe("Presidio PII"); + expect(result.Bedrock).toBe("Bedrock Guardrail"); + expect(result.Lakera).toBe("Lakera"); + + // Dynamic providers added with PascalCase keys + expect(result.ZscalerAiGuard).toBe("Zscaler AI Guard"); + expect(result.AporiaAi).toBe("Aporia AI"); + + // Should also update the module-level DynamicGuardrailProviders + expect(DynamicGuardrailProviders).toEqual(result); + }); + + it("should skip entries without ui_friendly_name", () => { + const apiResponse = { + valid_provider: { ui_friendly_name: "Valid Provider" }, + invalid_provider: { some_field: "no ui_friendly_name" }, + string_value: "not an object", + }; + + const result = populateGuardrailProviders(apiResponse); + + expect(result.ValidProvider).toBe("Valid Provider"); + expect(result.InvalidProvider).toBeUndefined(); + expect(result.StringValue).toBeUndefined(); + }); + }); + + describe("getGuardrailProviders", () => { + it("should return legacy GuardrailProviders enum when no dynamic providers are populated", () => { + const result = getGuardrailProviders(); + + expect(result).toEqual(GuardrailProviders); + expect(result).toHaveProperty("PresidioPII", "Presidio PII"); + }); + + it("should return dynamic providers when populated", () => { + populateGuardrailProviders({ + custom_guardrail: { ui_friendly_name: "Custom Guardrail" }, + }); + + const result = getGuardrailProviders(); + + // Returns dynamic (which includes legacy + custom) + expect(result.CustomGuardrail).toBe("Custom Guardrail"); + expect(result.PresidioPII).toBe("Presidio PII"); + }); + }); + + describe("shouldRenderPIIConfigSettings", () => { + it("should return true for PresidioPII provider key", () => { + expect(shouldRenderPIIConfigSettings("PresidioPII")).toBe(true); + }); + + it("should return false for non-Presidio providers", () => { + expect(shouldRenderPIIConfigSettings("Bedrock")).toBe(false); + expect(shouldRenderPIIConfigSettings("Lakera")).toBe(false); + }); + + it("should return false for null provider", () => { + expect(shouldRenderPIIConfigSettings(null)).toBe(false); + }); + }); + + describe("shouldRenderContentFilterConfigSettings", () => { + it("should return true when dynamic providers include LiteLLM Content Filter", () => { + populateGuardrailProviders({ + litellm_content_filter: { + ui_friendly_name: "LiteLLM Content Filter", + }, + }); + + expect( + shouldRenderContentFilterConfigSettings("LitellmContentFilter") + ).toBe(true); + }); + + it("should return false for unrelated providers", () => { + expect(shouldRenderContentFilterConfigSettings("PresidioPII")).toBe( + false + ); + }); + + it("should return false for null", () => { + expect(shouldRenderContentFilterConfigSettings(null)).toBe(false); + }); + }); + + describe("shouldRenderAzureTextModerationConfigSettings", () => { + it("should return true when dynamic providers include Azure Content Safety Text Moderation", () => { + populateGuardrailProviders({ + azure_content_safety: { + ui_friendly_name: "Azure Content Safety Text Moderation", + }, + }); + + expect( + shouldRenderAzureTextModerationConfigSettings("AzureContentSafety") + ).toBe(true); + }); + + it("should return false for null", () => { + expect(shouldRenderAzureTextModerationConfigSettings(null)).toBe(false); + }); + }); + + describe("getGuardrailLogoAndName", () => { + it("should return correct logo and display name for a known provider value", () => { + const result = getGuardrailLogoAndName("presidio"); + + expect(result.displayName).toBe("Presidio PII"); + expect(result.logo).toContain("microsoft_azure.svg"); + }); + + it("should return the raw value as displayName when provider is unknown", () => { + const result = getGuardrailLogoAndName("unknown_provider"); + + expect(result.displayName).toBe("unknown_provider"); + expect(result.logo).toBe(""); + }); + + it("should return fallback for empty string", () => { + const result = getGuardrailLogoAndName(""); + + expect(result.displayName).toBe("-"); + expect(result.logo).toBe(""); + }); + + it("should handle case-insensitive matching of provider values", () => { + const lower = getGuardrailLogoAndName("presidio"); + const upper = getGuardrailLogoAndName("PRESIDIO"); + const mixed = getGuardrailLogoAndName("Presidio"); + + expect(lower.displayName).toBe("Presidio PII"); + expect(upper.displayName).toBe("Presidio PII"); + expect(mixed.displayName).toBe("Presidio PII"); + }); + + it("should work with dynamically populated providers", () => { + populateGuardrailProviders({ + noma: { ui_friendly_name: "Noma Security" }, + }); + populateGuardrailProviderMap({ + noma: { ui_friendly_name: "Noma Security" }, + }); + + const result = getGuardrailLogoAndName("noma"); + + expect(result.displayName).toBe("Noma Security"); + expect(result.logo).toContain("noma_security.png"); + }); + }); +});