diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 40e36a900d3..a6f5ee9c6f4 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -10,5 +10,9 @@ "disable_aiohttp_trust_env", "disable_aiohttp_transport", "user_agent" + ], + "url_policy": [ + "user_url_validation", + "user_url_allowed_hosts" ] } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 2ab6517b61d..77542855fec 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -24,6 +24,7 @@ pub(crate) fn call_config( asynchronous: bool, ) -> PyResult { decline_live_clients(kwargs)?; + decline_custom_url_policy(&PythonSettings::UrlPolicy.read(py)?)?; let configured = settings(&PythonSettings::Http.read(py)?)? .with_environment(&|name| std::env::var(name).ok()); let settings = for_call(configured, call_ssl_verify(kwargs)?, asynchronous) @@ -63,6 +64,23 @@ pub(crate) fn decline_live_clients(kwargs: &Bound<'_, PyDict>) -> PyResult<()> { Ok(()) } +#[derive(FromPyObject)] +struct PythonUrlPolicy { + user_url_validation: bool, + user_url_allowed_hosts: Vec, +} + +fn decline_custom_url_policy(value: &Bound<'_, PyAny>) -> PyResult<()> { + match value.extract::() { + Ok(policy) if policy.user_url_validation && policy.user_url_allowed_hosts.is_empty() => { + Ok(()) + } + Ok(_) | Err(_) => Err(RustBridgeDeclined::new_err( + "litellm.user_url_validation / user_url_allowed_hosts are applied by the Python route", + )), + } +} + #[derive(FromPyObject)] struct PythonHttpSettings<'py> { ssl_verify: Bound<'py, PyAny>, @@ -245,6 +263,37 @@ user_agent='litellm/9.9.9', }); } + fn url_policy<'py>(py: Python<'py>, fields: &str) -> Bound<'py, PyAny> { + let source = std::ffi::CString::new(format!( + "import types\npolicy = types.SimpleNamespace({fields})" + )) + .unwrap(); + let locals = PyDict::new(py); + py.run(&source, Some(&locals), Some(&locals)).unwrap(); + locals.get_item("policy").unwrap().unwrap() + } + + #[test] + fn default_url_policy_stays_on_the_rust_route() { + Python::initialize(); + Python::attach(|py| { + let policy = url_policy(py, "user_url_validation=True, user_url_allowed_hosts=[]"); + decline_custom_url_policy(&policy).unwrap(); + }); + } + + #[rstest] + #[case::validation_off("user_url_validation=False, user_url_allowed_hosts=[]")] + #[case::allowlist("user_url_validation=True, user_url_allowed_hosts=['docs.internal']")] + #[case::mistyped("user_url_validation=True, user_url_allowed_hosts=None")] + fn custom_url_policy_declines_so_python_applies_it(#[case] fields: &str) { + Python::initialize(); + Python::attach(|py| { + let error = decline_custom_url_policy(&url_policy(py, fields)).unwrap_err(); + assert!(error.is_instance_of::(py)); + }); + } + #[test] fn mistyped_python_settings_decline_instead_of_raising() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index dcb46e7d2b5..b7855566850 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -5,15 +5,17 @@ const MODULE: &str = "litellm.rust_bridge.settings"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) enum PythonSettings { Http, + UrlPolicy, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 1] = [Self::Http]; + pub(crate) const ALL: [Self; 2] = [Self::Http, Self::UrlPolicy]; pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", + Self::UrlPolicy => "url_policy", } } diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 491312c97b6..bccfd01ec73 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass @@ -17,6 +18,21 @@ class HttpSettings: user_agent: str +@dataclass(frozen=True, slots=True) +class UrlPolicy: + user_url_validation: bool + user_url_allowed_hosts: Sequence[str] + + +def url_policy() -> UrlPolicy: + import litellm + + return UrlPolicy( + user_url_validation=litellm.user_url_validation, + user_url_allowed_hosts=litellm.user_url_allowed_hosts, + ) + + def http_settings() -> HttpSettings: import litellm from litellm.llms.custom_httpx.http_handler import default_user_agent diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index f4f9cbc8eec..7e7b1c6743b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -15,7 +15,20 @@ CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-b def test_the_rust_contract_matches_the_returned_fields() -> None: contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) - assert contract == {"http_settings": [field.name for field in dataclasses.fields(settings.http_settings())]} + assert contract == { + "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], + "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], + } + + +def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) + + assert settings.url_policy() == settings.UrlPolicy( + user_url_validation=False, + user_url_allowed_hosts=["docs.internal:8443"], + ) def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: