diff --git a/litellm/litellm_core_utils/deprecated_package.py b/litellm/litellm_core_utils/deprecated_package.py index 542686ecf19..f9915af4880 100644 --- a/litellm/litellm_core_utils/deprecated_package.py +++ b/litellm/litellm_core_utils/deprecated_package.py @@ -1,53 +1,73 @@ import importlib +import importlib.abc +import importlib.machinery +import importlib.util import sys import warnings -from itertools import chain -from pathlib import Path +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial from types import ModuleType from typing import Final -def _dotted_name(root: Path, path: Path, package_name: str) -> str: - parts: Final = path.relative_to(root).with_suffix("").parts - return ".".join((package_name, *parts)).removesuffix(".__init__") +@dataclass(frozen=True, slots=True) +class _AliasLoader(importlib.abc.Loader): + new_dotted: str + canonical_spec: importlib.machinery.ModuleSpec + old_dotted: str + + def create_module(self, spec: importlib.machinery.ModuleSpec) -> ModuleType: + return importlib.import_module(self.new_dotted) + + def exec_module(self, module: ModuleType) -> None: + module.__spec__ = self.canonical_spec # rebind-ok: [LIT011] restore import metadata overwritten by importlib + if self.canonical_spec.submodule_search_locations is not None: + _alias_attributes(self.old_dotted, self.new_dotted) -def _parents_and_self(dotted: str, package_name: str) -> tuple[str, ...]: - parts: Final = dotted.removeprefix(f"{package_name}.").split(".") - return tuple(f"{package_name}.{'.'.join(parts[:depth])}" for depth in range(1, len(parts) + 1)) +@dataclass(frozen=True, slots=True) +class _AliasFinder(importlib.abc.MetaPathFinder): + old_name: str + new_name: str - -def _module_names(package: ModuleType) -> tuple[str, ...]: - roots: Final = tuple(Path(root) for root in package.__path__) - files: Final = tuple( - (root, path) for root in roots for path in root.rglob("*.py") - ) # comprehension-ok: nested walk over each package root - leaves: Final = tuple( - name - for name in (_dotted_name(root, path, package.__name__) for root, path in files) - if name != package.__name__ - ) - return tuple( - sorted( - set(chain.from_iterable(_parents_and_self(leaf, package.__name__) for leaf in leaves)) - {package.__name__} + def find_spec( + self, + fullname: str, + path: Sequence[str] | None = None, + target: ModuleType | None = None, + ) -> importlib.machinery.ModuleSpec | None: + if not fullname.startswith(f"{self.old_name}."): + return None + new_dotted: Final = f"{self.new_name}{fullname.removeprefix(self.old_name)}" + new_spec: Final = importlib.util.find_spec(new_dotted) + if new_spec is None: + return None + return importlib.util.spec_from_loader( + fullname, + _AliasLoader(new_dotted, new_spec, fullname), + is_package=new_spec.submodule_search_locations is not None, ) - ) -def _alias(old_name: str, new_package_name: str, new_dotted: str) -> None: - module: Final = importlib.import_module(new_dotted) - old_dotted: Final = f"{old_name}{new_dotted.removeprefix(new_package_name)}" - old_parent, _, child = old_dotted.rpartition(".") - sys.modules[old_dotted] = module - setattr(sys.modules[old_parent], child, module) +def _deprecated_attribute(old_name: str, new_name: str, name: str) -> ModuleType: + if importlib.util.find_spec(f"{new_name}.{name}") is None: + raise AttributeError(f"module {old_name!r} has no attribute {name!r}") + return importlib.import_module(f"{old_name}.{name}") + + +def _alias_attributes(old_name: str, new_name: str) -> None: + if "__getattr__" not in vars(sys.modules[old_name]): + sys.modules[old_name].__getattr__ = partial(_deprecated_attribute, old_name, new_name) def alias_deprecated_package(old_name: str, new_package: ModuleType) -> None: - """Make every module under ``new_package`` importable and attribute-reachable under ``old_name`` and warn once.""" warnings.warn( f"{old_name} is deprecated, import {new_package.__name__} instead", DeprecationWarning, stacklevel=3, ) - for new_dotted in _module_names(new_package): - _alias(old_name, new_package.__name__, new_dotted) + finder: Final = _AliasFinder(old_name, new_package.__name__) + if finder not in sys.meta_path: + sys.meta_path.insert(0, finder) + _alias_attributes(old_name, new_package.__name__) diff --git a/tests/unit/experimental_mcp_client/test_init.py b/tests/unit/experimental_mcp_client/test_init.py index 9d2f1137470..1ce2a236d9d 100644 --- a/tests/unit/experimental_mcp_client/test_init.py +++ b/tests/unit/experimental_mcp_client/test_init.py @@ -19,6 +19,8 @@ def test_old_package_warns_and_deep_imports_resolve_to_mcp_client(monkeypatch: p with pytest.warns(DeprecationWarning, match="litellm.mcp_client"): import litellm.experimental_mcp_client as old + assert old.client is litellm.mcp_client.client + assert old.tools is litellm.mcp_client.tools assert old.__all__ == litellm.mcp_client.__all__ assert old.load_mcp_tools is litellm.mcp_client.load_mcp_tools assert old.call_openai_tool is litellm.mcp_client.call_openai_tool diff --git a/tests/unit/proxy/_experimental/mcp_server/test_init.py b/tests/unit/proxy/_experimental/mcp_server/test_init.py index 2f933e11351..5ecb79b9813 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_init.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_init.py @@ -1,6 +1,7 @@ import importlib import sys from collections.abc import Mapping +from typing import Final import pytest @@ -16,9 +17,17 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa if name.startswith(OLD_PACKAGE): monkeypatch.delitem(sys.modules, name) + monkeypatch.delattr(litellm.proxy.mcp_server.auth, "user_api_key_auth_mcp") + monkeypatch.delitem(sys.modules, "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp") + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): old = importlib.import_module(OLD_PACKAGE) + assert old.server is litellm.proxy.mcp_server.server + assert old.auth is litellm.proxy.mcp_server.auth + assert old.auth.user_api_key_auth_mcp is importlib.import_module( + "litellm.proxy.mcp_server.auth.user_api_key_auth_mcp" + ) assert importlib.import_module(f"{OLD_PACKAGE}.server") is litellm.proxy.mcp_server.server assert importlib.import_module(f"{OLD_PACKAGE}.utils") is litellm.proxy.mcp_server.utils assert ( @@ -28,3 +37,43 @@ def test_old_package_warns_and_deep_imports_resolve_to_proxy_mcp_server(monkeypa attrs: Mapping[str, object] = vars(old) assert attrs["server"] is litellm.proxy.mcp_server.server assert attrs["auth"] is litellm.proxy.mcp_server.auth + + +def test_old_stdio_helper_import_works_without_optional_mcp_dependency(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delattr(litellm.proxy, "mcp_server") + old_parent: Final = importlib.import_module("litellm.proxy._experimental") + monkeypatch.delattr(old_parent, "mcp_server", raising=False) + monkeypatch.delattr(litellm, "mcp_client", raising=False) + prefixes: Final = (OLD_PACKAGE, "litellm.proxy.mcp_server", "litellm.mcp_client", "mcp") + for name in tuple(sys.modules): + if any(name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes): + monkeypatch.delitem(sys.modules, name) + monkeypatch.setitem(sys.modules, "mcp", None) + + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): + old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate") + new: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate") + + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") + assert old.is_mcp_stdio_enabled() is True + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "false") + assert old.is_mcp_stdio_enabled() is False + assert old is new + + +def test_legacy_child_import_keeps_canonical_module_reload_working(monkeypatch: pytest.MonkeyPatch) -> None: + for name in tuple(sys.modules): + if name.startswith(OLD_PACKAGE): + monkeypatch.delitem(sys.modules, name) + + helper: Final = importlib.import_module("litellm.proxy.mcp_server.stdio_gate") + with pytest.warns(DeprecationWarning, match="litellm.proxy.mcp_server"): + old: Final = importlib.import_module(f"{OLD_PACKAGE}.stdio_gate") + + monkeypatch.setattr(helper, "MCP_STDIO_ENABLED_ENV_VAR", "RELOAD_PROBE") + monkeypatch.setenv("LITELLM_ENABLE_MCP_STDIO", "true") + monkeypatch.setenv("RELOAD_PROBE", "false") + importlib.reload(helper) + + assert old is helper + assert old.is_mcp_stdio_enabled() is True