Merge pull request #26336 from BerriAI/litellm_yj_apr22

[IInfra] Merge dev branch

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -15,6 +15,7 @@ import litellm.types
 import litellm.types.llms
 from litellm import verbose_logger
 from litellm._uuid import uuid
+from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get
 from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
 from litellm.types.files import get_file_extension_from_mime_type
 from litellm.types.llms.anthropic import *
@@ -3324,7 +3325,7 @@ def _load_image_from_url(image_url):
     try:
         # Send a GET request to the image URL
         client = HTTPHandler(concurrent_limit=1)
-        response = client.get(image_url)
+        response = safe_get(client, image_url)
         response.raise_for_status()  # Raise an exception for HTTP errors
 
         # Check the response's content type to ensure it is an image
@@ -3562,7 +3563,7 @@ class BedrockImageProcessor:
                 params={"concurrent_limit": 1},
             )
             # Send a GET request to the image URL
-            response = await client.get(image_url, follow_redirects=True)
+            response = await async_safe_get(client, image_url)
             response.raise_for_status()  # Raise an exception for HTTP errors
 
             return BedrockImageProcessor._post_call_image_processing(
@@ -3577,7 +3578,7 @@ class BedrockImageProcessor:
         try:
             client = HTTPHandler(concurrent_limit=1)
             # Send a GET request to the image URL
-            response = client.get(image_url, follow_redirects=True)
+            response = safe_get(client, image_url)
             response.raise_for_status()  # Raise an exception for HTTP errors
 
             return BedrockImageProcessor._post_call_image_processing(

litellm/litellm_core_utils/token_counter.py

@@ -23,6 +23,7 @@ from litellm.constants import (
     DEFAULT_IMAGE_HEIGHT,
     DEFAULT_IMAGE_TOKEN_COUNT,
     DEFAULT_IMAGE_WIDTH,
+    MAX_IMAGE_URL_DOWNLOAD_SIZE_MB,
     MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES,
     MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES,
     MAX_TILE_HEIGHT,
@@ -215,7 +216,14 @@ def get_image_dimensions(
         try:
             client = _get_httpx_client()
             response = safe_get(client, data)
-            img_data = response.read()
+            max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024)
+            content_length = response.headers.get("Content-Length")
+            if content_length is not None and int(content_length) > max_bytes:
+                pass  # skip download; img_data stays None
+            else:
+                body = response.read()
+                if len(body) <= max_bytes:
+                    img_data = body
         except Exception:
             pass
     if img_data is None:

litellm/llms/bedrock/batches/transformation.py

@@ -1,4 +1,5 @@
 import os
+import re
 import time
 from typing import Any, Dict, List, Literal, Optional, Union, cast
 
@@ -294,7 +295,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
             raise ValueError(f"Invalid ARN format: {batch_id}")
 
         region = arn_parts[3]
-        # arn_parts[5] contains "model-invocation-job/{jobId}"
+        if not re.match(r"^[a-z][a-z0-9-]*$", region):
+            raise ValueError(f"Invalid region in ARN: {batch_id}")
 
         # Build the endpoint URL for GetModelInvocationJob
         # AWS API format: GET /model-invocation-job/{jobIdentifier}

litellm/llms/black_forest_labs/image_edit/transformation.py

@@ -14,7 +14,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 import httpx
 from httpx._types import RequestFiles
 
+import litellm
 from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+from litellm.litellm_core_utils.url_utils import safe_get
 from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.images.main import ImageEditOptionalRequestParams
@@ -208,14 +210,14 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig):
             )
         elif isinstance(image, str):
             if image.startswith(("http://", "https://")):
-                # Download image from URL
-                response = httpx.get(image, timeout=60.0)
+                response = safe_get(litellm.module_level_client, image, timeout=60.0)
                 response.raise_for_status()
                 return response.content
             else:
-                # Assume it's a file path
-                with open(image, "rb") as f:
-                    return f.read()
+                raise ValueError(
+                    "Unsupported image input: plain string values that are not URLs are not accepted. "
+                    "Provide image bytes or a file-like object."
+                )
         elif hasattr(image, "read"):
             # File-like object
             pos = getattr(image, "tell", lambda: 0)()

litellm/llms/s3_vectors/vector_stores/transformation.py

@@ -1,3 +1,4 @@
+import re
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import httpx
@@ -66,6 +67,8 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
         aws_region_name = litellm_params.get("aws_region_name")
         if not aws_region_name:
             raise ValueError("aws_region_name is required for S3 Vectors")
+        if not re.match(r"^[a-z][a-z0-9-]*$", aws_region_name):
+            raise ValueError("Invalid aws_region_name format")
         return f"https://s3vectors.{aws_region_name}.api.aws"
 
     def transform_search_vector_store_request(

litellm/llms/snowflake/utils.py

@@ -1,3 +1,4 @@
+import re
 from typing import TYPE_CHECKING, Any, List, Optional, Tuple
 
 from litellm.secret_managers.main import get_secret_str
@@ -61,6 +62,8 @@ class SnowflakeBaseConfig:
                 account_id = get_secret_str("SNOWFLAKE_ACCOUNT_ID")
             if account_id is None:
                 raise ValueError("Missing snowflake account_id")
+            if not re.match(r"^[a-zA-Z0-9_-]+$", account_id):
+                raise ValueError("Invalid account_id format")
             api_base = f"https://{account_id}.snowflakecomputing.com/api/v2"
 
         api_base = api_base.rstrip("/")

litellm/llms/vertex_ai/common_utils.py

@@ -232,8 +232,11 @@ def get_vertex_base_url(
     """
     if vertex_location == "global":
         return "https://aiplatform.googleapis.com"
-    else:
-        return f"https://{vertex_location}-aiplatform.googleapis.com"
+    if vertex_location is None:
+        raise ValueError("vertex_location is required")
+    if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
+        raise ValueError("Invalid vertex_location format")
+    return f"https://{vertex_location}-aiplatform.googleapis.com"
 
 
 def _get_embedding_url(

litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py

@@ -368,13 +368,16 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM):
             if stream_pos is not None:
                 image.seek(stream_pos)
             return data
-        if isinstance(image, (str, Path)):
-            path_obj = Path(image)
-            if not path_obj.exists():
-                raise ValueError(
-                    f"Mask/image path does not exist for Vertex AI Imagen image edit: {path_obj}"
-                )
-            return path_obj.read_bytes()
+        if isinstance(image, str):
+            raise ValueError(
+                "Unsupported image input: plain string values are not accepted for "
+                "Vertex AI Imagen image edit. Provide image bytes or a file-like object."
+            )
+        if isinstance(image, Path):
+            raise ValueError(
+                "Unsupported image input: filesystem paths are not accepted for "
+                "Vertex AI Imagen image edit. Provide image bytes or a file-like object."
+            )
         if hasattr(image, "read"):
             data = image.read()
             if isinstance(data, str):

litellm/proxy/auth/auth_utils.py

@@ -151,7 +151,15 @@ def is_request_body_safe(
     A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key.
     Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997
     """
-    banned_params = ["api_base", "base_url", "user_config"]
+    banned_params = [
+        "api_base",
+        "base_url",
+        "user_config",
+        "aws_sts_endpoint",
+        "aws_web_identity_token",
+        "aws_role_name",
+        "vertex_credentials",
+    ]
 
     for param in banned_params:
         if (

litellm/proxy/image_endpoints/endpoints.py

@@ -285,6 +285,13 @@ async def image_edit_api(
     if mask_files:
         data["mask"] = mask_files
 
+    for _field in ("image", "mask"):
+        if _field in data and isinstance(data[_field], str):
+            raise HTTPException(
+                status_code=422,
+                detail=f"'{_field}' must be provided as a multipart file upload, not a string.",
+            )
+
     # Ensure prompt exists in data (default to None for models that don't require it)
     if "prompt" not in data:
         data["prompt"] = None

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -8,6 +8,7 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
 
 import json
 import os
+import re
 from typing import Any, Optional, Tuple, Union, cast
 
 import httpx
@@ -1500,6 +1501,10 @@ def get_vertex_base_url(vertex_location: Optional[str]) -> str:
     """
     if vertex_location == "global":
         return "https://aiplatform.googleapis.com/"
+    if vertex_location is None:
+        raise ValueError("vertex_location is required")
+    if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
+        raise ValueError("Invalid vertex_location format")
     return f"https://{vertex_location}-aiplatform.googleapis.com/"