mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 19:04:59 +00:00
e8fcb01215
* Cato Networks guardrail, based on Aim (#26597) * Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim * Add more tests * Move test so they are reached by codecov coverage * base URL trailing slashes * Support Lemonade runtime context metadata (#28135) * Support Lemonade runtime context metadata * Add provider hook for runtime model metadata * Address provider model info review feedback Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise. Co-authored-by: openhands <openhands@all-hands.dev> * Fix CI after staging rebase Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec. Co-authored-by: openhands <openhands@all-hands.dev> * Normalize Lemonade runtime model metadata * Avoid leaking Ollama metadata auth * Avoid leaking Lemonade metadata auth --------- Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com> Co-authored-by: openhands <openhands@all-hands.dev> * fix(cato): address guardrail review feedback Use proxy-authenticated user identity, forward moderation hook return values, and ensure streaming sender tasks are cancelled and awaited on exit. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of #28010 (#28846) * fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path Fixes #26083 vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma- prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up and should_use_openai_handler routes it to the OpenAI-compatible /endpoints/openapi/chat/completions URL. No gemma-detection exclusion needed (the "gemma/" check uses a slash, which google/gemma-... doesn't match). No OpenAIGPTConfig subclass needed — works with the base handler. * fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified) * fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup * test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS Addresses oss-pr-review-agent-shin feedback on PR #28010: supports_function_calling, supports_tool_choice, and supports_vision were marked true but had no tests proving the payloads actually reached the OpenAI-compatible endpoint. Added: - test_gemma_maas_supports_function_calling — verifies the utility returns True when the model_cost entry carries supports_function_calling=true - test_gemma_maas_supports_vision — same for supports_vision - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice appear in the JSON body POSTed to /endpoints/openapi/chat/completions - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts survive transformation and reach the global endpoint URL * fix: Delete uv.lock * test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS Addresses oss-pr-review-agent-shin feedback on PR #28010: P1 (patch target): Added a comment explaining why patching litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct — get_async_httpx_client() (defined in http_handler.py) instantiates AsyncHTTPHandler within that module's scope, so the definition-site patch intercepts it. Without the mock the test raises AuthenticationError, confirming it never silently passes. P2 (partner-provider regression guard): Added test_gemma_routes_through_openai_handler() which calls VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's routing to VertexPartnerProvider.llama ever changes the URL-shape tests below it become a real regression guard rather than an unanchored unit test. Also added: - test_gemma_maas_supports_function_calling / supports_vision — capability flag checks via patch.dict(litellm.model_cost) - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice forwarded in the request body - test_vertex_ai_gemma_vision_passthrough — image_url part survives transformation to the global endpoint Added: - test_gemma_maas_supports_function_calling — verifies the utility returns True when the model_cost entry carries supports_function_calling=true - test_gemma_maas_supports_vision — same for supports_vision - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice appear in the JSON body POSTed to /endpoints/openapi/chat/completions - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts survive transformation and reach the global endpoint URL * fix: proper patch for unit tests --------- Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local> * fix(cato): guardrail all completion choices on output When n > 1, only choices[0] was analyzed and redacted. Iterate every Choices entry so block and anonymize actions apply to all completions. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix review * fix(cato_networks): harden output anonymize handling and restructure nested UI routes Guard against empty redacted_output and empty all_redacted_messages from Cato. Restructure nested admin UI HTML exports to index.html so extensionless routes work. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix mypy * fix(cato): guard missing policy_drill_down and all_redacted_messages keys * fix(cato): avoid KeyError bypassing block action on missing analysis_result * fix(cato): preserve non-text message fields during anonymize Rebuild redacted messages from the original messages, overwriting only content, so tool_calls, tool_call_id, name and multimodal fields survive the anonymize action. * fix(cato): preserve trailing messages when fewer redacted messages returned Avoid silently truncating the conversation in _anonymize_request when Cato returns fewer redacted messages than were sent, and isolate the no-api-key config test from a pre-existing CATO_API_KEY environment variable. * fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup Suppress ConnectionClosed (alongside CancelledError) when tearing down the Cato streaming sender task so a backend ConnectionClosed cannot mask the original StreamingCallbackError (e.g. a guardrail block) raised by the receive loop. Thread api_key through get_model_info -> _get_model_info_helper so an explicit key reaches a provider's dynamic get_model_info for a caller-supplied api_base. Previously only api_base was forwarded, so authenticated Ollama and Lemonade servers at a custom base could only be queried unauthenticated. * fix(cato): surface mid-stream forwarding errors instead of blocking on recv If the upstream LLM stream errors mid-flight, the sender task dies before sending the terminal done frame, so the consumer would block on websocket.recv() until Cato closes the connection. Race recv against the sender task and raise the stored sender exception promptly as a StreamingCallbackError. * fix(cato): drop spoofable end_user_id from guardrail user identity Only the key/JWT-bound user_email is a trusted identity. end_user_id is resolved from caller-supplied request fields (OpenAI user param, headers, metadata), so an authenticated caller with no bound user_email could set it to another user's email and have LiteLLM forward x-cato-user-email for that victim, poisoning Cato audit and policy attribution. Forward only user_email and omit the header otherwise. * fix(cato): harden output anonymize path against missing content key * fix(cato): fall back to original message when redacted content key is missing * refactor(model-info): drop unused api_key from cached model-info helper _cached_get_model_info_helper is only called by the cost-tracking hot path, which never authenticates, so the api_key parameter was never populated. Keeping it in the lru_cache key offered no benefit and risked fragmenting the high-RPS cache and retaining credential strings per entry. * fix(cato): preserve None content on tool-call-only choices in output hook * fix(ollama): respect static-model guard in OllamaConfig.get_model_info Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama models short-circuit before the /api/show network call instead of hitting the server unconditionally. * fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds An empty-string api_key was treated as an explicit key, so it passed the guard meant to keep server-side credentials off caller-supplied bases and then fell back through the env/global key chain. A caller could point api_base at a server they control and send api_key="" to receive the configured provider key in the Authorization header. Gate the credential fallback on the api_key being truthy instead of merely not-None. * fix(cato): inspect and redact Responses-API input, not just messages The guardrail only read data["messages"], so /v1/responses requests, which carry their text in data["input"], reached Cato as an empty message list and bypassed inspection entirely. Send build_inspection_messages(data) so both shapes are analyzed, and write anonymized results back with apply_redacted_messages_back when the request used input. * perf(utils): keep api_key out of get_model_info lru_cache key * fix(cato): propagate ssl_verify to streaming WebSocket connection The streaming hook applied ssl_verify only to the HTTP handler; the websockets.connect() call used default verification, so a custom Cato instance behind TLS with a self-signed cert worked for non-streaming calls but failed every streaming request. Resolve the ssl_verify setting into the connect() ssl argument, mirroring the HTTP handler. * refactor(utils): rename shadowing local in _get_model_info_helper * fix(cato): flatten multimodal chat content before inspection Chat Completions requests whose message content is a multimodal parts array were posted to Cato as the raw OpenAI parts, so text inside content: [{"type":"text", ...}] reached the model without Cato ever inspecting the string. Flatten each message's list content to plain text while keeping the list 1:1 with the request so the index-based redaction write-back stays valid; Responses-API input requests still go through build_inspection_messages. * test(lemonade): clear get_model_info cache around api_base test * fix(cato): inspect and redact Responses-API input even when messages present _inspection_messages returned early once messages was non-empty, so a /v1/responses caller could place benign text in messages and disallowed text in input and have only messages reach Cato while the model used input. Inspect both fields and write anonymize redactions back to input as well as the index-aligned messages. * test(log_db_metrics): assert table_name event_metadata contract log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata (table_name only, function_name/function_kwargs/function_args dropped as redundant with call_type and unsafe to stamp on a span). The success-path test still asserted function_name membership and crashed with TypeError on the None metadata returned when no table_name is passed. Pass a table_name and assert the surfaced contract instead. * fix(cato): inspect and redact completion prompt and Responses-API instructions The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from. * fix(cato): serialize non-str/bytes websocket chunks before forwarding * fix(cato): inspect tool descriptions and tool-call arguments * fix(cato): map redacted output by assistant index; restore get_model_info.cache_info * fix(cato): block output even when detection_message is null/empty A block_action returned by Cato on the output hook whose detection_message was null or empty was let through to the caller: the truthiness guard on detection_message skipped the HTTPException and the unblocked response was returned. Raise the HTTPException directly in _handle_block_action_on_output so the output path blocks unconditionally, mirroring the input path. * fix(cato): inspect and redact nested tool param and legacy function descriptions Tool/function parameter descriptions and the legacy functions[] array are forwarded to the model but were not seen by Cato, so blocked text hidden there bypassed inspection and anonymization. Recursively walk every description string in tools[].function and functions[] schemas for both the analyze payload and the anonymize write-back. * fix(cato): traverse schema descriptions iteratively to satisfy recursive detector The nested walk() generator recursed over tool/function JSON schemas with no depth bound, which the recursive_detector code-quality gate rejects. Replace it with an explicit-stack DFS that yields the same (container, key) refs in the same pre-order, so schema description redaction is unchanged. * fix(cato): inspect and redact response_format JSON schema descriptions response_format json_schema descriptions are forwarded to the model, so blocked text hidden in nested schema descriptions could bypass Cato inspection and redaction. Extend the schema-description walk to cover response_format alongside tools and legacy functions. * fix(cato): skip output rewrite when Cato returns no redaction Return None from call_cato_guardrail_on_output on monitor/no-action so the post-call hook only mutates the message when there is an actual redaction, instead of redundantly re-writing the original content. * refactor(utils): resolve explicit api_key model info without the cache Move the model-info build into a non-cached _build_model_info helper and drop api_key from the lru-cached _cached_get_model_info signature. Both cached helpers now take the same (model, provider, api_base) key and never forward api_key, while explicit per-caller keys are resolved through the builder directly instead of reaching into the cache wrapper's __wrapped__. * fix(cato): inspect and redact non-description schema string values Tool, function and response_format JSON schemas forward more than just description text to the model. enum, const, default, examples and title values are sent verbatim, so blocked content hidden in any of them bypassed Cato inspection and redaction. Walk those schema string values alongside descriptions on both the inspection and anonymize paths. * fix(model-info): surface swallowed dynamic model-info errors The provider-specific get_model_info dispatch falls back to the static cost map when a provider's dynamic lookup raises, which is intentional graceful degradation. Previously the exception was discarded with a bare debug line, so a real failure (e.g. a provider whose get_model_info signature does not accept api_key) was invisible. Log the exception at warning level with the model and provider context so the fallback is diagnosable. * fix(cato): inspect and redact Responses API output in post-call hook The post-call success hook only handled ModelResponse, so /v1/responses (which returns a ResponsesAPIResponse) bypassed the Cato output guardrail. Extract and inspect/redact every output_text content block and function-call arguments string, blocking on a block action, so generated text cannot escape inspection by using the Responses API. * chore: reset _experimental/out folder * chore(ui): remove orphaned prebuilt dashboard chunk files The _experimental/out manifests are byte-identical to the base branch, so the served dashboard already matches base. 436 unreferenced Next.js chunk files had accumulated in the directory and are not loaded by any manifest; removing them restores the committed UI artifacts to the base build and drops the artifact churn from this PR's diff. * fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show --------- Co-authored-by: Alex Yaroslavsky <trexinc@gmail.com> Co-authored-by: Graham Neubig <neubig@gmail.com> Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com> Co-authored-by: openhands <openhands@all-hands.dev> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Piotr Placzko <piotr@icep-design.com> Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
236 lines
8.9 KiB
Python
236 lines
8.9 KiB
Python
"""
|
|
Unit tests for per-service SSL support in LiteLLM.
|
|
|
|
These tests verify that ssl_verify parameters are correctly propagated
|
|
through the call stack without requiring live API credentials.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
|
|
# Add litellm to path
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
import litellm.proxy.guardrails.guardrail_hooks.aim.aim as _aim_module
|
|
import litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks as _cato_networks_module
|
|
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
|
from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM
|
|
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
|
from litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks import CatoNetworksGuardrail
|
|
|
|
|
|
class TestBaseAWSLLMSSLVerify:
|
|
"""Test SSL verification parameter handling in BaseAWSLLM."""
|
|
|
|
def test_get_ssl_verify_with_parameter(self):
|
|
"""Test that _get_ssl_verify accepts and uses the ssl_verify parameter."""
|
|
base_llm = BaseAWSLLM()
|
|
|
|
# Test with True
|
|
result = base_llm._get_ssl_verify(ssl_verify=True)
|
|
assert result is True
|
|
|
|
# Test with False
|
|
result = base_llm._get_ssl_verify(ssl_verify=False)
|
|
assert result is False
|
|
|
|
# Test with cert path
|
|
cert_path = "/path/to/cert.pem"
|
|
result = base_llm._get_ssl_verify(ssl_verify=cert_path)
|
|
assert result == cert_path
|
|
|
|
def test_get_ssl_verify_without_parameter(self):
|
|
"""Test that _get_ssl_verify falls back to environment/global when no parameter."""
|
|
base_llm = BaseAWSLLM()
|
|
|
|
# Should fall back to environment or global litellm.ssl_verify
|
|
result = base_llm._get_ssl_verify()
|
|
# Result depends on environment, just verify it doesn't crash
|
|
assert result is not None or result is None # Can be None, True, False, or path
|
|
|
|
@patch("boto3.client")
|
|
def test_get_credentials_propagates_ssl_verify(self, mock_boto_client):
|
|
"""Test that get_credentials propagates ssl_verify to boto3 clients."""
|
|
base_llm = BaseAWSLLM()
|
|
|
|
# Mock the boto3 client
|
|
mock_sts_client = Mock()
|
|
mock_sts_client.assume_role.return_value = {
|
|
"Credentials": {
|
|
"AccessKeyId": "test_key",
|
|
"SecretAccessKey": "test_secret",
|
|
"SessionToken": "test_token",
|
|
"Expiration": "2026-01-20T00:00:00Z",
|
|
}
|
|
}
|
|
mock_boto_client.return_value = mock_sts_client
|
|
|
|
# Call get_credentials with ssl_verify parameter
|
|
cert_path = "/path/to/cert.pem"
|
|
try:
|
|
base_llm.get_credentials(
|
|
aws_access_key_id="test_key",
|
|
aws_secret_access_key="test_secret",
|
|
aws_region_name="us-east-1",
|
|
ssl_verify=cert_path,
|
|
)
|
|
except Exception:
|
|
# May fail due to missing credentials, but we're checking the call
|
|
pass
|
|
|
|
# Verify boto3.client was called with verify parameter
|
|
# Note: This test verifies the parameter is accepted, actual propagation
|
|
# is tested in integration tests
|
|
assert True # If we got here without error, parameter was accepted
|
|
|
|
|
|
class TestBedrockLLMSSLVerify:
|
|
"""Test SSL verification parameter handling in BedrockLLM."""
|
|
|
|
def test_bedrock_llm_accepts_ssl_verify_in_optional_params(self):
|
|
"""Test that BedrockLLM can receive ssl_verify in optional_params."""
|
|
# This is a simple test to verify the parameter is accepted
|
|
# The actual propagation is tested in integration tests
|
|
bedrock_llm = BedrockLLM()
|
|
|
|
# Verify the class exists and can be instantiated
|
|
assert bedrock_llm is not None
|
|
|
|
# Verify _get_ssl_verify method exists and works
|
|
result = bedrock_llm._get_ssl_verify(ssl_verify="/path/to/cert.pem")
|
|
assert result == "/path/to/cert.pem"
|
|
|
|
|
|
class TestAimGuardrailSSLVerify:
|
|
"""Test SSL verification parameter handling in AimGuardrail."""
|
|
|
|
def test_init_accepts_ssl_verify(self):
|
|
"""Test that AimGuardrail.__init__ accepts and uses ssl_verify parameter."""
|
|
mock_handler = Mock()
|
|
|
|
# Use patch.object on the actual module reference for reliable patching
|
|
# across different import orders / CI environments
|
|
with patch.object(
|
|
_aim_module, "get_async_httpx_client", return_value=mock_handler
|
|
) as mock_get_client:
|
|
# Initialize with ssl_verify
|
|
cert_path = "/path/to/aim_cert.pem"
|
|
AimGuardrail(
|
|
api_key="test_key",
|
|
api_base="https://test.aim.api",
|
|
ssl_verify=cert_path,
|
|
)
|
|
|
|
# Verify get_async_httpx_client was called with ssl_verify in params
|
|
assert mock_get_client.called
|
|
call_kwargs = mock_get_client.call_args[1]
|
|
assert "params" in call_kwargs
|
|
assert call_kwargs["params"] is not None
|
|
assert call_kwargs["params"]["ssl_verify"] == cert_path
|
|
|
|
def test_init_without_ssl_verify(self):
|
|
"""Test that AimGuardrail works without ssl_verify parameter."""
|
|
mock_handler = Mock()
|
|
|
|
# Use patch.object on the actual module reference for reliable patching
|
|
with patch.object(
|
|
_aim_module, "get_async_httpx_client", return_value=mock_handler
|
|
) as mock_get_client:
|
|
# Initialize without ssl_verify
|
|
AimGuardrail(api_key="test_key", api_base="https://test.aim.api")
|
|
|
|
# Should still work, just without custom SSL
|
|
assert mock_get_client.called
|
|
|
|
|
|
class TestCatoNetworksGuardrailSSLVerify:
|
|
"""Test SSL verification parameter handling in CatoNetworksGuardrail."""
|
|
|
|
def test_init_accepts_ssl_verify(self):
|
|
"""Test that CatoNetworksGuardrail.__init__ accepts and uses ssl_verify parameter."""
|
|
mock_handler = Mock()
|
|
|
|
# Use patch.object on the actual module reference for reliable patching
|
|
# across different import orders / CI environments
|
|
with patch.object(
|
|
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
|
|
) as mock_get_client:
|
|
# Initialize with ssl_verify
|
|
cert_path = "/path/to/cato_cert.pem"
|
|
CatoNetworksGuardrail(
|
|
api_key="test_key",
|
|
api_base="https://test.catonetworks.api",
|
|
ssl_verify=cert_path,
|
|
)
|
|
|
|
# Verify get_async_httpx_client was called with ssl_verify in params
|
|
assert mock_get_client.called
|
|
call_kwargs = mock_get_client.call_args[1]
|
|
assert "params" in call_kwargs
|
|
assert call_kwargs["params"] is not None
|
|
assert call_kwargs["params"]["ssl_verify"] == cert_path
|
|
|
|
def test_init_without_ssl_verify(self):
|
|
"""Test that CatoNetworksGuardrail works without ssl_verify parameter."""
|
|
mock_handler = Mock()
|
|
|
|
# Use patch.object on the actual module reference for reliable patching
|
|
with patch.object(
|
|
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
|
|
) as mock_get_client:
|
|
# Initialize without ssl_verify
|
|
CatoNetworksGuardrail(api_key="test_key", api_base="https://test.catonetworks.api")
|
|
|
|
# Should still work, just without custom SSL
|
|
assert mock_get_client.called
|
|
|
|
|
|
class TestHTTPHandlerSSLVerify:
|
|
"""Test SSL verification parameter handling in HTTP handlers."""
|
|
|
|
def test_get_async_httpx_client_accepts_ssl_verify_in_params(self):
|
|
"""Test that get_async_httpx_client accepts ssl_verify in params dict."""
|
|
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
|
from litellm.types.llms.custom_http import httpxSpecialProvider
|
|
|
|
# Call with ssl_verify in params
|
|
cert_path = "/path/to/cert.pem"
|
|
client = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
|
params={"ssl_verify": cert_path},
|
|
)
|
|
|
|
# Verify client was created (actual SSL config is tested in integration tests)
|
|
assert client is not None
|
|
|
|
|
|
def test_ssl_verify_parameter_types():
|
|
"""Test that various ssl_verify parameter types are handled correctly."""
|
|
base_llm = BaseAWSLLM()
|
|
|
|
# Test boolean True
|
|
result = base_llm._get_ssl_verify(ssl_verify=True)
|
|
assert result is True
|
|
|
|
# Test boolean False
|
|
result = base_llm._get_ssl_verify(ssl_verify=False)
|
|
assert result is False
|
|
|
|
# Test string path
|
|
cert_path = "/path/to/cert.pem"
|
|
result = base_llm._get_ssl_verify(ssl_verify=cert_path)
|
|
assert result == cert_path
|
|
|
|
# Test None (should fall back to environment/global)
|
|
result = base_llm._get_ssl_verify(ssl_verify=None)
|
|
# Result depends on environment
|
|
assert result is not None or result is None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Run tests
|
|
pytest.main([__file__, "-v", "--tb=short"])
|