diff --git a/litellm/llms/base_llm/managed_resources/utils.py b/litellm/llms/base_llm/managed_resources/utils.py new file mode 100644 index 00000000000..0d843b6d128 --- /dev/null +++ b/litellm/llms/base_llm/managed_resources/utils.py @@ -0,0 +1,364 @@ +""" +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