Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/types_utils/utils.py: 23%
103 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1import asyncio
2import importlib
3import importlib.util
4import os
5from collections.abc import Callable
6from typing import Any, Final, Literal, get_type_hints
9def get_instance_fn(value: str, config_file_path: str | None = None) -> Any:
10 module_name = value
11 instance_name = None
12 try:
13 # Check if value starts with s3:// or gcs://
14 if value.startswith("s3://") or value.startswith("gcs://"): 14 ↛ 22line 14 didn't jump to line 22 because the condition on line 14 was never true
15 # Remote module loading is a documented operator feature when
16 # invoked from config-file load (``config_file_path`` carries
17 # the YAML path). Without that signal the URL is request-body
18 # data on an admin endpoint — a one-step admin-to-RCE primitive
19 # via ``_load_instance_from_remote_storage``'s ``exec_module``.
20 # Register the module under ``litellm_settings`` in the
21 # config.yaml instead.
22 if config_file_path is None:
23 raise ValueError(
24 "Remote module loading (s3://, gcs://) is only "
25 "permitted from the config-file load path. Register "
26 "the module under ``litellm_settings`` in your "
27 "config.yaml instead."
28 )
29 return _load_instance_from_remote_storage(value, config_file_path)
31 # Split the path by dots to separate module from instance
32 parts: Final = value.split(".")
34 # The module path is all but the last part, and the instance_name is the last part
35 module_name = ".".join(parts[:-1])
36 instance_name = parts[-1]
38 module_file_path = None
39 if config_file_path is not None: 39 ↛ 40line 39 didn't jump to line 40 because the condition on line 39 was never true
40 directory: Final = os.path.dirname(config_file_path)
41 module_file_path = os.path.join(directory, *module_name.split(".")) + ".py"
43 if module_file_path is not None and os.path.exists(module_file_path): 43 ↛ 44line 43 didn't jump to line 44 because the condition on line 43 was never true
44 spec: Final = importlib.util.spec_from_file_location(module_name, module_file_path)
45 if spec is None:
46 raise ImportError(f"Could not find a module specification for {module_file_path}")
47 module = importlib.util.module_from_spec(spec)
48 if spec.loader is None:
49 raise ImportError(f"Could not find a module loader for {module_file_path}")
50 spec.loader.exec_module(module)
51 else:
52 module = importlib.import_module(module_name)
54 # Get the instance from the module
55 instance: Final[object] = getattr(module, instance_name)
57 return instance
58 except ImportError as e:
59 # Re-raise the exception with a user-friendly message
60 if instance_name and module_name:
61 raise ImportError(f"Could not import {instance_name} from {module_name}") from e
62 else:
63 raise e
64 except Exception as e:
65 raise e
68def _load_instance_from_remote_storage(remote_url: str, config_file_path: str | None = None) -> Any:
69 """
70 Load custom logger instance from S3 or GCS URL.
72 Expected format:
73 - s3://bucket-name/path/to/module.instance_name
74 - gcs://bucket-name/path/to/module.instance_name
76 Args:
77 remote_url (str): The s3:// or gcs:// URL
78 config_file_path (str): Optional config file path for temp directory context
80 Returns:
81 Any: The loaded instance
82 """
83 try:
84 from litellm._logging import verbose_proxy_logger
86 # Parse the URL
87 if remote_url.startswith("s3://"):
88 storage_type = "s3"
89 url_without_prefix = remote_url[5:] # Remove 's3://'
90 elif remote_url.startswith("gcs://"):
91 storage_type = "gcs"
92 url_without_prefix = remote_url[6:] # Remove 'gcs://'
93 else:
94 raise ValueError(f"Unsupported URL scheme in {remote_url}")
96 # Split bucket and path
97 parts: Final = url_without_prefix.split("/", 1)
98 if len(parts) < 2:
99 raise ValueError(
100 f"Invalid URL format: {remote_url}. Expected: {storage_type}://bucket-name/path/to/module.instance"
101 )
103 bucket_name: Final = parts[0]
104 path_and_module: Final = parts[1]
106 # Extract module path and instance name
107 # Example: "loggers/custom_callbacks.proxy_handler_instance"
108 # Handle case where user accidentally includes .py extension
109 if path_and_module.endswith(".py"):
110 module_name_without_py: Final = path_and_module[:-3] # Remove .py
111 raise ValueError(
112 f"Invalid URL format in {remote_url}. "
113 f"Don't include '.py' extension and you must specify the instance name. "
114 f"Expected format: {storage_type}://{bucket_name}/{module_name_without_py}.instance_name "
115 f"(e.g., {storage_type}://{bucket_name}/{module_name_without_py}.proxy_handler_instance)"
116 )
118 # Split by last dot to separate module from instance
119 module_parts: Final = path_and_module.split(".")
120 if len(module_parts) < 2:
121 raise ValueError(f"Invalid module specification in {remote_url}. Expected: path/to/module.instance_name")
123 instance_name: Final = module_parts[-1]
124 module_path: Final = ".".join(module_parts[:-1])
126 # Create object key (file path in bucket)
127 object_key: Final = f"{module_path}.py"
129 verbose_proxy_logger.debug(
130 "Loading custom logger from %s: bucket=%s, object_key=%s, instance=%s",
131 storage_type,
132 bucket_name,
133 object_key,
134 instance_name,
135 )
137 import tempfile
139 # Create temporary file for the downloaded module using the actual module name
140 temp_file: Final = tempfile.NamedTemporaryFile(suffix=".py", delete=False)
141 local_file_path: Final = temp_file.name
142 temp_file.close() # Close the file so we can write to it
144 # Download the file
145 if storage_type == "s3":
146 from litellm.proxy.common_utils.load_config_utils import (
147 download_python_file_from_s3,
148 )
150 success = download_python_file_from_s3(
151 bucket_name=bucket_name,
152 object_key=object_key,
153 local_file_path=local_file_path,
154 )
155 else: # gcs
156 success = asyncio.run(_download_gcs_file_wrapper(bucket_name, object_key, local_file_path))
158 if not success:
159 raise ImportError(f"Failed to download {object_key} from {storage_type} bucket {bucket_name}")
161 # Load the module from the downloaded file using the actual module name
162 spec: Final = importlib.util.spec_from_file_location(module_path, local_file_path)
163 if spec is None or spec.loader is None:
164 raise ImportError(f"Could not create module spec for {local_file_path}")
166 module: Final = importlib.util.module_from_spec(spec)
167 spec.loader.exec_module(module)
169 # Get the instance
170 instance: Final[object] = getattr(module, instance_name)
172 # Clean up the temporary file
173 try:
174 os.remove(local_file_path)
175 except Exception as cleanup_error:
176 verbose_proxy_logger.warning("Could not clean up temporary file %s: %s", local_file_path, cleanup_error)
178 verbose_proxy_logger.info("Successfully loaded custom logger from %s", remote_url)
179 return instance
181 except Exception as e:
182 raise ImportError(f"Failed to load custom logger from {remote_url}: {e}") from e
185async def _download_gcs_file_wrapper(bucket_name: str, object_key: str, local_file_path: str) -> bool:
186 """Wrapper for GCS download to handle async properly"""
187 try:
188 from litellm.proxy.common_utils.load_config_utils import (
189 download_python_file_from_gcs,
190 )
192 return await download_python_file_from_gcs(bucket_name, object_key, local_file_path)
193 except Exception as e:
194 from litellm._logging import verbose_proxy_logger
196 verbose_proxy_logger.error("Error downloading from GCS: %s", e)
197 return False
200def validate_custom_validate_return_type(
201 fn: Callable[..., Any] | None,
202) -> Callable[..., Literal[True]] | None:
203 if fn is None:
204 return None
206 hints: Final = get_type_hints(fn)
207 return_type: Final = hints.get("return")
209 if return_type != Literal[True]:
210 raise TypeError(f"Custom validator must be annotated to return Literal[True], got {return_type}")
212 return fn