scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

chore(providers): guard URL-valued model destinations

Browse source
This commit is contained in:
user 2026-04-30 13:03:56 -07:00
parent 08541f49ee
commit 5ec925b02b
10 changed files with 232 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/llms/gemini/files/transformation.py
View file

@ -29,6 +29,8 @@ from litellm.types.utils import LlmProviders
from ..common_utils import GeminiModelInfo
_GEMINI_FILES_HOST = "generativelanguage.googleapis.com"
class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
def __init__(self):
@ -248,6 +250,13 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
"""
if file_id.startswith(("http://", "https://")):
parsed = urlparse(file_id)
if (
parsed.scheme != "https"
or parsed.hostname != _GEMINI_FILES_HOST
or parsed.username is not None
or parsed.password is not None
):
raise ValueError("Invalid Gemini file URL")
path = parsed.path.lstrip("/")
files_index = path.find("files/")
if files_index != -1:
@ -260,9 +269,22 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
normalized_file_id = normalized_file_id.strip("/")
if not normalized_file_id.startswith("files/"):
normalized_file_id = f"files/{normalized_file_id}"
self._validate_gemini_file_name(normalized_file_id)
return normalized_file_id
@staticmethod
def _validate_gemini_file_name(file_name: str) -> None:
parts = file_name.split("/")
if (
len(parts) != 2
or parts[0] != "files"
or not parts[1]
or parts[1] in {".", ".."}
or any(char in parts[1] for char in ("\\", "?", "#"))
):
raise ValueError("Invalid Gemini file name")
def transform_retrieve_file_response(
self,
raw_response: httpx.Response,
@ -337,13 +359,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig):
if not api_key:
raise ValueError("api_key is required")
# Extract file name from URI if full URI is provided
# file_id could be "files/abc123" or "https://generativelanguage.googleapis.com/v1beta/files/abc123"
if file_id.startswith("http"):
# Extract the file path from full URI
file_name = file_id.split("/v1beta/")[-1]
else:
file_name = file_id if file_id.startswith("files/") else f"files/{file_id}"
file_name = self._normalize_gemini_file_id(file_id)
# Construct the delete URL
url = f"{api_base}/v1beta/{file_name}"

15
litellm/llms/huggingface/chat/transformation.py
View file

@ -16,7 +16,11 @@ else:
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
from ..common_utils import HuggingFaceError, _fetch_inference_provider_mapping
from ..common_utils import (
HuggingFaceError,
_fetch_inference_provider_mapping,
validate_huggingface_model_identifier,
)
logger = logging.getLogger(__name__)
@ -76,9 +80,8 @@ class HuggingFaceChatConfig(OpenAIGPTConfig):
Do not add the chat/embedding/rerank extension here. Let the handler do this.
"""
if model.startswith(("http://", "https://")):
base_url = model
elif base_url is None:
validate_huggingface_model_identifier(model)
if base_url is None:
base_url = os.getenv("HF_API_BASE") or os.getenv("HUGGINGFACE_API_BASE", "")
return base_url
@ -95,6 +98,7 @@ class HuggingFaceChatConfig(OpenAIGPTConfig):
Get the complete URL for the API call.
For provider-specific routing through huggingface
"""
validate_huggingface_model_identifier(model)
# Check if api_base is provided
if api_base is not None:
complete_url = api_base
@ -103,9 +107,6 @@ class HuggingFaceChatConfig(OpenAIGPTConfig):
complete_url = str(os.getenv("HF_API_BASE")) or str(
os.getenv("HUGGINGFACE_API_BASE")
)
elif model.startswith(("http://", "https://")):
complete_url = model
complete_url = _build_chat_completion_url(complete_url)
# Default construction with provider
else:
# Parse provider and model

14
litellm/llms/huggingface/common_utils.py
View file

@ -9,6 +9,20 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
HF_HUB_URL = "https://huggingface.co"
def validate_huggingface_model_identifier(model: str) -> None:
"""Reject URL-valued model identifiers before provider credentials are added."""
if "://" not in model:
return
raise HuggingFaceError(
status_code=400,
message=(
"Invalid Hugging Face model identifier. Configure custom endpoints with "
"api_base or HF_API_BASE/HUGGINGFACE_API_BASE instead of passing a URL "
"as the model."
),
)
class HuggingFaceError(BaseLLMException):
def __init__(
self,

8
litellm/llms/huggingface/embedding/handler.py
View file

@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.types.utils import EmbeddingResponse
from ...base import BaseLLM
from ..common_utils import HuggingFaceError
from ..common_utils import HuggingFaceError, validate_huggingface_model_identifier
from .transformation import HuggingFaceEmbeddingConfig
config = HuggingFaceEmbeddingConfig()
@ -154,6 +154,7 @@ class HuggingFaceEmbedding(BaseLLM):
embed_url: str,
) -> dict:
data: Dict = {}
validate_huggingface_model_identifier(model)
## TRANSFORMATION ##
if "sentence-transformers" in model:
@ -334,6 +335,7 @@ class HuggingFaceEmbedding(BaseLLM):
headers={},
) -> EmbeddingResponse:
super().embedding()
validate_huggingface_model_identifier(model)
headers = config.validate_environment(
api_key=api_key,
headers=headers,
@ -348,9 +350,7 @@ class HuggingFaceEmbedding(BaseLLM):
)
# print_verbose(f"{model}, {task}")
embed_url = ""
if "https" in model:
embed_url = model
elif api_base:
if api_base:
embed_url = api_base
elif "HF_API_BASE" in os.environ:
embed_url = os.getenv("HF_API_BASE", "")

13
litellm/llms/huggingface/embedding/transformation.py
View file

@ -21,7 +21,13 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import Choices, Message, ModelResponse, Usage
from litellm.utils import token_counter
from ..common_utils import HuggingFaceError, hf_task_list, hf_tasks, output_parser
from ..common_utils import (
HuggingFaceError,
hf_task_list,
hf_tasks,
output_parser,
validate_huggingface_model_identifier,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -336,9 +342,8 @@ class HuggingFaceEmbeddingConfig(BaseConfig):
Do not add the chat/embedding/rerank extension here. Let the handler do this.
"""
if "https" in model:
completion_url = model
elif api_base is not None:
validate_huggingface_model_identifier(model)
if api_base is not None:
completion_url = api_base
elif "HF_API_BASE" in os.environ:
completion_url = os.getenv("HF_API_BASE", "")

12
litellm/llms/oobabooga/chat/oobabooga.py
View file

@ -5,7 +5,7 @@ import litellm
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.utils import EmbeddingResponse, ModelResponse, Usage
from ..common_utils import OobaboogaError
from ..common_utils import OobaboogaError, validate_oobabooga_model_identifier
from .transformation import OobaboogaConfig
oobabooga_config = OobaboogaConfig()
@ -26,6 +26,7 @@ def completion(
logger_fn=None,
default_max_tokens_to_sample=None,
):
validate_oobabooga_model_identifier(model)
headers = oobabooga_config.validate_environment(
api_key=api_key,
headers={},
@ -34,9 +35,7 @@ def completion(
optional_params=optional_params,
litellm_params=litellm_params,
)
if "https" in model:
completion_url = model
elif api_base:
if api_base:
completion_url = api_base
else:
raise OobaboogaError(
@ -96,9 +95,8 @@ def embedding(
encoding=None,
):
# Create completion URL
if "https" in model:
embeddings_url = model
elif api_base:
validate_oobabooga_model_identifier(model)
if api_base:
embeddings_url = f"{api_base}/v1/embeddings"
else:
raise OobaboogaError(

13
litellm/llms/oobabooga/common_utils.py
View file

@ -13,3 +13,16 @@ class OobaboogaError(BaseLLMException):
headers: Optional[Union[dict, httpx.Headers]] = None,
):
super().__init__(status_code=status_code, message=message, headers=headers)
def validate_oobabooga_model_identifier(model: str) -> None:
"""Oobabooga endpoints must be configured with api_base, not model URLs."""
if "://" not in model:
return
raise OobaboogaError(
status_code=400,
message=(
"Invalid Oobabooga model identifier. Configure the endpoint with "
"api_base instead of passing a URL as the model."
),
)

35
tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py
View file

@ -2,7 +2,6 @@
Test Google AI Studio (Gemini) files transformation functionality
"""
import os
from unittest.mock import Mock, patch
import httpx
@ -93,6 +92,27 @@ class TestGoogleAIStudioFilesTransformation:
assert "key=" not in url
assert params == {}
def test_transform_retrieve_file_request_rejects_untrusted_full_url(self):
file_id = "https://attacker.example/v1beta/files/test123"
litellm_params = {"api_key": "test-api-key"}
with pytest.raises(ValueError, match="Invalid Gemini file URL"):
self.handler.transform_retrieve_file_request(
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
)
def test_transform_retrieve_file_request_rejects_traversal_name(self):
litellm_params = {"api_key": "test-api-key"}
with pytest.raises(ValueError, match="Invalid Gemini file name"):
self.handler.transform_retrieve_file_request(
file_id="files/../secrets",
optional_params={},
litellm_params=litellm_params,
)
@patch.dict("os.environ", {}, clear=True)
@patch("litellm.llms.gemini.common_utils.get_secret_str", return_value=None)
def test_transform_retrieve_file_request_missing_api_key(self, mock_get_secret):
@ -322,3 +342,16 @@ class TestGoogleAIStudioFilesTransformation:
assert file_id in url
assert "generativelanguage.googleapis.com" in url
assert params == {}
def test_transform_delete_file_request_rejects_untrusted_full_url(self):
litellm_params = {
"api_key": "test-api-key",
"api_base": "https://generativelanguage.googleapis.com",
}
with pytest.raises(ValueError, match="Invalid Gemini file URL"):
self.handler.transform_delete_file_request(
file_id="https://attacker.example/v1beta/files/test123",
optional_params={},
litellm_params=litellm_params,
)

76
tests/test_litellm/llms/huggingface/test_huggingface_model_url_guard.py Normal file
View file

@ -0,0 +1,76 @@
from unittest.mock import MagicMock, patch
import pytest
from litellm.llms.huggingface.chat.transformation import HuggingFaceChatConfig
from litellm.llms.huggingface.common_utils import HuggingFaceError
from litellm.llms.huggingface.embedding.handler import HuggingFaceEmbedding
from litellm.llms.huggingface.embedding.transformation import (
HuggingFaceEmbeddingConfig,
)
def test_huggingface_chat_rejects_url_valued_model():
config = HuggingFaceChatConfig()
with pytest.raises(HuggingFaceError) as exc_info:
config.get_complete_url(
api_base=None,
api_key="hf-secret",
model="https://attacker.example/v1",
optional_params={},
litellm_params={},
)
assert exc_info.value.status_code == 400
def test_huggingface_chat_keeps_explicit_api_base_for_custom_endpoints():
config = HuggingFaceChatConfig()
complete_url = config.get_complete_url(
api_base="https://admin-configured.example",
api_key="hf-secret",
model="huggingface/mistral",
optional_params={},
litellm_params={},
)
assert complete_url == "https://admin-configured.example/v1/chat/completions"
def test_huggingface_embedding_config_rejects_url_valued_model():
config = HuggingFaceEmbeddingConfig()
with pytest.raises(HuggingFaceError) as exc_info:
config.get_api_base(
api_base=None,
model="prefixhttps://attacker.example/embeddings",
)
assert exc_info.value.status_code == 400
def test_huggingface_embedding_handler_rejects_before_task_lookup():
handler = HuggingFaceEmbedding()
logging_obj = MagicMock()
encoding = MagicMock()
encoding.encode.return_value = []
with patch(
"litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model"
) as mock_task_lookup:
with pytest.raises(HuggingFaceError) as exc_info:
handler.embedding(
model="https://attacker.example/embeddings",
input=["hello"],
model_response=MagicMock(),
optional_params={},
litellm_params={},
logging_obj=logging_obj,
encoding=encoding,
api_key="hf-secret",
)
assert exc_info.value.status_code == 400
mock_task_lookup.assert_not_called()

46
tests/test_litellm/llms/oobabooga/test_oobabooga_model_url_guard.py Normal file
View file

@ -0,0 +1,46 @@
from unittest.mock import MagicMock, patch
import pytest
from litellm.llms.oobabooga.chat.oobabooga import completion, embedding
from litellm.llms.oobabooga.common_utils import OobaboogaError
def test_oobabooga_completion_rejects_url_valued_model_before_request():
with patch("litellm.llms.oobabooga.chat.oobabooga._get_httpx_client") as mock_get:
with pytest.raises(OobaboogaError) as exc_info:
completion(
model="https://attacker.example/v1",
messages=[],
api_base="https://admin-configured.example",
model_response=MagicMock(),
print_verbose=MagicMock(),
encoding=MagicMock(),
api_key="ooba-secret",
logging_obj=MagicMock(),
optional_params={},
litellm_params={},
)
assert exc_info.value.status_code == 400
mock_get.assert_not_called()
def test_oobabooga_embedding_rejects_url_valued_model_before_request():
with patch(
"litellm.llms.oobabooga.chat.oobabooga.litellm.module_level_client.post"
) as mock_post:
with pytest.raises(OobaboogaError) as exc_info:
embedding(
model="prefixhttps://attacker.example/embeddings",
input=["hello"],
model_response=MagicMock(),
api_key="ooba-secret",
api_base="https://admin-configured.example",
logging_obj=MagicMock(),
optional_params={},
encoding=MagicMock(),
)
assert exc_info.value.status_code == 400
mock_post.assert_not_called()