"""
Utility functions for managed resources.
This module provides common utility functions that can be used across
different managed resource types (files, vector stores, etc.).
"""
import base64
import re
from typing import List, Optional, Union, Literal
def is_base64_encoded_unified_id(
resource_id: str,
prefix: str = "litellm_proxy:",
) -> Union[str, Literal[False]]:
"""
Check if a resource ID is a base64 encoded unified ID.
Args:
resource_id: The resource ID to check
prefix: The expected prefix for unified IDs
Returns:
Decoded string if valid unified ID, False otherwise
"""
# Ensure resource_id is a string
if not isinstance(resource_id, str):
return False
# Add padding back if needed
padded = resource_id + "=" * (-len(resource_id) % 4)
# Decode from base64
try:
decoded = base64.urlsafe_b64decode(padded).decode()
if decoded.startswith(prefix):
return decoded
else:
return False
except Exception:
return False
def extract_target_model_names_from_unified_id(
unified_id: str,
) -> List[str]:
"""
Extract target model names from a unified resource ID.
Args:
unified_id: The unified resource ID (decoded or encoded)
Returns:
List of target model names
Example:
unified_id = "litellm_proxy:vector_store;unified_id,uuid;target_model_names,gpt-4,gemini-2.0"
returns: ["gpt-4", "gemini-2.0"]
"""
try:
# Ensure unified_id is a string
if not isinstance(unified_id, str):
return []
# Decode if it's base64 encoded
decoded_id = is_base64_encoded_unified_id(unified_id)
if decoded_id:
unified_id = decoded_id
# Extract model names using regex
match = re.search(r"target_model_names,([^;]+)", unified_id)
if match:
# Split on comma and strip whitespace from each model name
return [model.strip() for model in match.group(1).split(",")]
return []
except Exception:
return []
def extract_resource_type_from_unified_id(
unified_id: str,
) -> Optional[str]:
"""
Extract resource type from a unified resource ID.
Args:
unified_id: The unified resource ID (decoded or encoded)
Returns:
Resource type string or None
Example:
unified_id = "litellm_proxy:vector_store;unified_id,uuid;..."
returns: "vector_store"
"""
try:
# Ensure unified_id is a string
if not isinstance(unified_id, str):
return None
# Decode if it's base64 encoded
decoded_id = is_base64_encoded_unified_id(unified_id)
if decoded_id:
unified_id = decoded_id
# Extract resource type (comes after prefix and before first semicolon)
match = re.search(r"litellm_proxy:([^;]+)", unified_id)
if match:
return match.group(1).strip()
return None
except Exception:
return None
def extract_unified_uuid_from_unified_id(
unified_id: str,
) -> Optional[str]:
"""
Extract the UUID from a unified resource ID.
Args:
unified_id: The unified resource ID (decoded or encoded)
Returns:
UUID string or None
Example:
unified_id = "litellm_proxy:vector_store;unified_id,abc-123;..."
returns: "abc-123"
"""
try:
# Ensure unified_id is a string
if not isinstance(unified_id, str):
return None
# Decode if it's base64 encoded
decoded_id = is_base64_encoded_unified_id(unified_id)
if decoded_id:
unified_id = decoded_id
# Extract UUID
match = re.search(r"unified_id,([^;]+)", unified_id)
if match:
return match.group(1).strip()
return None
except Exception:
return None
def extract_model_id_from_unified_id(
unified_id: str,
) -> Optional[str]:
"""
Extract model ID from a unified resource ID.
Args:
unified_id: The unified resource ID (decoded or encoded)
Returns:
Model ID string or None
Example:
unified_id = "litellm_proxy:vector_store;...;model_id,gpt-4-model-id;..."
returns: "gpt-4-model-id"
"""
try:
# Ensure unified_id is a string
if not isinstance(unified_id, str):
return None
# Decode if it's base64 encoded
decoded_id = is_base64_encoded_unified_id(unified_id)
if decoded_id:
unified_id = decoded_id
# Extract model ID
match = re.search(r"model_id,([^;]+)", unified_id)
if match:
return match.group(1).strip()
return None
except Exception:
return None
def extract_provider_resource_id_from_unified_id(
unified_id: str,
) -> Optional[str]:
"""
Extract provider resource ID from a unified resource ID.
Args:
unified_id: The unified resource ID (decoded or encoded)
Returns:
Provider resource ID string or None
Example:
unified_id = "litellm_proxy:vector_store;...;resource_id,vs_abc123;..."
returns: "vs_abc123"
"""
try:
# Ensure unified_id is a string
if not isinstance(unified_id, str):
return None
# Decode if it's base64 encoded
decoded_id = is_base64_encoded_unified_id(unified_id)
if decoded_id:
unified_id = decoded_id
# Extract resource ID (try multiple patterns for different resource types)
patterns = [
r"resource_id,([^;]+)",
r"vector_store_id,([^;]+)",
r"file_id,([^;]+)",
]
for pattern in patterns:
match = re.search(pattern, unified_id)
if match:
return match.group(1).strip()
return None
except Exception:
return None
def generate_unified_id_string(
resource_type: str,
unified_uuid: str,
target_model_names: List[str],
provider_resource_id: str,
model_id: str,
additional_fields: Optional[dict] = None,
) -> str:
"""
Generate a unified ID string (before base64 encoding).
Args:
resource_type: Type of resource (e.g., "vector_store", "file")
unified_uuid: UUID for this unified resource
target_model_names: List of target model names
provider_resource_id: Resource ID from the provider
model_id: Model ID from the router
additional_fields: Additional fields to include in the ID
Returns:
Unified ID string (not yet base64 encoded)
Example:
generate_unified_id_string(
resource_type="vector_store",
unified_uuid="abc-123",
target_model_names=["gpt-4", "gemini"],
provider_resource_id="vs_xyz",
model_id="model-id-123",
)
returns: "litellm_proxy:vector_store;unified_id,abc-123;target_model_names,gpt-4,gemini;resource_id,vs_xyz;model_id,model-id-123"
"""
# Build the unified ID string
parts = [
f"litellm_proxy:{resource_type}",
f"unified_id,{unified_uuid}",
f"target_model_names,{','.join(target_model_names)}",
f"resource_id,{provider_resource_id}",
f"model_id,{model_id}",
]
# Add additional fields if provided
if additional_fields:
for key, value in additional_fields.items():
parts.append(f"{key},{value}")
return ";".join(parts)
def encode_unified_id(unified_id_string: str) -> str:
"""
Encode a unified ID string to base64.
Args:
unified_id_string: The unified ID string to encode
Returns:
Base64 encoded unified ID (URL-safe, padding stripped)
"""
return (
base64.urlsafe_b64encode(unified_id_string.encode())
.decode()
.rstrip("=")
)
def decode_unified_id(encoded_unified_id: str) -> Optional[str]:
"""
Decode a base64 encoded unified ID.
Args:
encoded_unified_id: The base64 encoded unified ID
Returns:
Decoded unified ID string or None if invalid
"""
try:
# Add padding back if needed
padded = encoded_unified_id + "=" * (-len(encoded_unified_id) % 4)
# Decode from base64
decoded = base64.urlsafe_b64decode(padded).decode()
# Verify it starts with the expected prefix
if decoded.startswith("litellm_proxy:"):
return decoded
return None
except Exception:
return None
def parse_unified_id(
unified_id: str,
) -> Optional[dict]:
"""
Parse a unified ID into its components.
Args:
unified_id: The unified ID (encoded or decoded)
Returns:
Dictionary with parsed components or None if invalid
Example:
{
"resource_type": "vector_store",
"unified_uuid": "abc-123",
"target_model_names": ["gpt-4", "gemini"],
"provider_resource_id": "vs_xyz",
"model_id": "model-id-123"
}
"""
try:
# Decode if needed
decoded_id = decode_unified_id(unified_id)
if not decoded_id:
# Maybe it's already decoded
if unified_id.startswith("litellm_proxy:"):
decoded_id = unified_id
else:
return None
return {
"resource_type": extract_resource_type_from_unified_id(decoded_id),
"unified_uuid": extract_unified_uuid_from_unified_id(decoded_id),
"target_model_names": extract_target_model_names_from_unified_id(decoded_id),
"provider_resource_id": extract_provider_resource_id_from_unified_id(decoded_id),
"model_id": extract_model_id_from_unified_id(decoded_id),
}
except Exception:
return None