✨ Multi-modal data sdk

Author: Zhi-a

sdk/nexent/multi_modal/__init__.py

@@ -0,0 +1,383 @@
+import functools
+import inspect
+import logging
+from io import BytesIO
+from typing import Any, Callable, List, Optional, Tuple
+import requests
+
+from .utils import (
+    is_url,
+    generate_object_name,
+    detect_content_type_from_bytes,
+    guess_extension_from_content_type,
+)
+
+logger = logging.getLogger("multi_modal")
+
+
+# Global storage client instance
+_storage_client: Optional[Any] = None
+
+
+def set_storage_client(storage_client: Any) -> None:
+    """
+    Initialize the global storage client for load_object and save_object decorators
+    
+    This should be called once at application startup to set up the storage client
+    that will be used by all decorators.
+    
+    Args:
+        storage_client: Storage client instance with get_file_stream() and upload_fileobj() methods
+    """
+    global _storage_client
+    _storage_client = storage_client
+
+
+def get_storage_client() -> Optional[Any]:
+    """
+    Get the global storage client instance
+    
+    Returns:
+        Storage client instance or None if not initialized
+    """
+    return _storage_client
+
+
+def parse_s3_url(s3_url: str) -> Tuple[str, str]:
+    """
+    Parse S3 URL to extract bucket and object_name
+
+    Supports formats:
+    - s3://bucket/key
+    - /bucket/key (MinIO path format)
+
+    Args:
+        s3_url: S3 URL string
+
+    Returns:
+        Tuple[bucket, object_name]
+
+    Raises:
+        ValueError: If URL format is not recognized
+    """
+    if not s3_url:
+        raise ValueError("S3 URL cannot be empty")
+
+    # s3:// protocol
+    if s3_url.startswith('s3://'):
+        parts = s3_url.replace('s3://', '').split('/', 1)
+        if len(parts) == 2:
+            bucket, object_name = parts
+            if not bucket or not object_name:
+                raise ValueError(f"Invalid s3:// URL format: {s3_url}")
+            return bucket, object_name
+        else:
+            raise ValueError(f"Invalid s3:// URL format: {s3_url}")
+
+    # MinIO path format: /bucket/key
+    elif s3_url.startswith('/'):
+        parts = s3_url.lstrip('/').split('/', 1)
+        if len(parts) == 2:
+            bucket, object_name = parts
+            return bucket, object_name
+        else:
+            raise ValueError(f"Invalid path format: {s3_url}")
+
+    else:
+        raise ValueError(f"Unrecognized S3 URL format: {s3_url[:50]}...")
+
+
+def download_file_from_url(url: str, timeout: int = 30) -> Optional[bytes]:
+    """
+    Download file content from S3 URL or HTTP/HTTPS URL as bytes
+
+    Uses the global storage client for S3 URLs. Call set_storage_client() to initialize.
+
+    Args:
+        url: S3 URL (s3://bucket/key or /bucket/key) or HTTP/HTTPS URL
+        timeout: Request timeout in seconds (for HTTP URLs)
+
+    Returns:
+        File content as bytes, or None if failed
+
+    Raises:
+        ValueError: If storage client is not initialized for S3 URLs
+    """
+    if not url:
+        return None
+
+    try:
+        # Handle HTTP/HTTPS URLs - direct download
+        if url.startswith(('http://', 'https://')):
+            response = requests.get(url, timeout=timeout)
+            response.raise_for_status()
+            return response.content
+
+        # Handle S3 URLs - use global storage_client
+        elif url.startswith('s3://') or url.startswith('/'):
+            if _storage_client is None:
+                raise ValueError(
+                    f"Storage client required for S3 URL, but not initialized. "
+                    f"Call set_storage_client() first. URL: {url[:50]}..."
+                )
+
+            bucket, object_name = parse_s3_url(url)
+
+            # Get file stream from storage client
+            if hasattr(_storage_client, 'get_file_stream'):
+                success, stream = _storage_client.get_file_stream(object_name, bucket)
+                if not success:
+                    raise ValueError(f"Failed to get file stream from storage: {stream}")
+
+                # Read stream content to bytes
+                try:
+                    bytes_data = stream.read()
+                    if hasattr(stream, 'close'):
+                        stream.close()
+                    return bytes_data
+                except Exception as e:
+                    raise ValueError(f"Failed to read stream content: {e}")
+            else:
+                raise ValueError(f"Storage client does not have get_file_stream method")
+
+        else:
+            raise ValueError(f"Unsupported URL format: {url[:50]}...")
+
+    except Exception as e:
+        logger.error(f"Failed to download file from URL: {e}")
+        return None
+
+
+def _upload_bytes_to_minio(
+    bytes_data: bytes,
+    object_name: Optional[str] = None,
+    bucket: str = "multi-modal",
+    content_type: str = "application/octet-stream"
+) -> str:
+    """
+    Upload bytes to MinIO and return file URL
+    
+    Uses the global storage client. Call set_storage_client() to initialize.
+    
+    Args:
+        bytes_data: File content as bytes
+        object_name: Object name, if not specified will be auto-generated
+        bucket: Bucket name, if not specified use default bucket
+        content_type: MIME type for file extension guessing
+        
+    Returns:
+        File URL (path format: /bucket/object_name)
+        
+    Raises:
+        ValueError: If upload fails or storage_client is not initialized
+    """
+    if _storage_client is None:
+        raise ValueError(
+            "Storage client is required for uploading to MinIO. "
+            "Call set_storage_client() first."
+        )
+    
+    if not hasattr(_storage_client, 'upload_fileobj'):
+        raise ValueError("Storage client must have upload_fileobj method")
+    
+    # Generate object name if not provided
+    if object_name is None:
+        file_ext = guess_extension_from_content_type(content_type)
+        object_name = generate_object_name(file_ext)
+    
+    # Create BytesIO object from bytes
+    file_obj = BytesIO(bytes_data)
+    
+    # Upload to MinIO
+    success, result = _storage_client.upload_fileobj(file_obj, object_name, bucket)
+    
+    if not success:
+        raise ValueError(f"Failed to upload file to MinIO: {result}")
+    
+    return result
+
+
+def load_object(
+        input_names: List[str],
+        input_data_transformer: Optional[List[Callable[[bytes], Any]]] = None,
+):
+    """
+    Decorator: Automatically download files from S3/HTTP URLs and convert to specified format
+    
+    Behavior:
+    - Input should be S3 URL (s3://bucket/key or /bucket/key) or HTTP/HTTPS URL
+    - Download file and convert to binary stream (bytes)
+    - If input_data_transformer is provided, use custom function to convert bytes to corresponding object
+    - If input_data_transformer is not provided, pass bytes directly to function
+    
+    Uses the global storage client for S3 URLs. Call set_storage_client() to initialize.
+    
+    Args:
+        input_names: List of input parameter names to transform
+        input_data_transformer: Optional list of custom transformer functions.
+                               Each function should accept bytes (binary stream) and return transformed object.
+                               If provided, bytes will be passed to transformer to convert to custom format.
+    """
+
+    def decorator(func: Callable):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Get function signature and bound arguments
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+
+            # Iterate through all parameters that need transformation
+            for i, param_name in enumerate(input_names):
+                if param_name in bound_args.arguments:
+                    original_data = bound_args.arguments[param_name]
+                    
+                    # Skip if data is None or empty
+                    if original_data is None:
+                        continue
+
+                    # Check if it's a URL - download and convert to bytes
+                    if isinstance(original_data, str) and is_url(original_data):
+                        try:
+                            # Download file as bytes (binary stream)
+                            bytes_data = download_file_from_url(original_data)
+                            
+                            if bytes_data is None:
+                                raise ValueError(f"Failed to download file from URL: {original_data}")
+                            
+                            # If custom transformer is provided, use it to convert bytes to object
+                            if input_data_transformer and i < len(input_data_transformer):
+                                transformer_func = input_data_transformer[i]
+                                try:
+                                    transformed_data = transformer_func(bytes_data)
+                                    logger.info(f"Downloaded {param_name} from URL and transformed using {transformer_func.__name__}")
+                                except Exception as e:
+                                    logger.error(f"Error transforming {param_name} with {transformer_func.__name__}: {e}")
+                                    raise e
+                            else:
+                                # No transformer: pass bytes directly
+                                transformed_data = bytes_data
+                                logger.info(f"Downloaded {param_name} from URL as bytes (binary stream)")
+                        except Exception as e:
+                            logger.error(f"Error downloading {param_name} from URL: {e}")
+                            raise e
+                    else:
+                        # Not a URL: raise error or pass through?
+                        # For now, we'll raise an error since load_object is specifically for URLs
+                        raise ValueError(
+                            f"Parameter '{param_name}' is not a URL. "
+                            f"load_object decorator expects S3 or HTTP/HTTPS URLs. "
+                            f"Got: {type(original_data).__name__}"
+                        )
+                    
+                    # Update parameter value with transformed result
+                    bound_args.arguments[param_name] = transformed_data
+
+            # Call original function with transformed parameters
+            return func(*bound_args.args, **bound_args.kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def save_object(
+        output_names: List[str],
+        output_transformers: Optional[List[Callable[[Any], bytes]]] = None,
+        bucket: str = "multi-modal"
+):
+    """
+    Decorator: Automatically upload function return values to MinIO and return file URLs
+    
+    Behavior:
+    - Function return values are expected to be bytes by default, which will be uploaded directly to MinIO
+    - If output_transformers is provided, use custom functions to convert return values to bytes first
+    - Upload bytes to MinIO and return file URLs instead of original return values
+    
+    Uses the global storage client. Call set_storage_client() to initialize.
+    
+    Args:
+        output_names: List of output result names (for logging/debugging)
+        output_transformers: Optional list of custom transformer functions.
+                           Each function should accept the original return value and return bytes.
+                           If provided, will be used to convert return values to bytes before uploading.
+        bucket: Bucket name, if not specified use default bucket from storage_client
+    """
+
+    def decorator(func: Callable) -> Callable:
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # Call original function
+            results = func(*args, **kwargs)
+
+            # Normalize results to tuple
+            if not isinstance(results, tuple):
+                results_tuple = (results,)
+            else:
+                results_tuple = results
+
+            if len(results_tuple) != len(output_names):
+                raise ValueError(
+                    f"Function returned {len(results_tuple)} values, "
+                    f"but expected {len(output_names)} outputs"
+                )
+
+            # Process each return value
+            uploaded_urls = []
+            for i, (result, name) in enumerate(zip(results_tuple, output_names)):
+                bytes_data = None
+                
+                # If transformer is provided, use it to convert to bytes
+                if output_transformers and i < len(output_transformers):
+                    transformer_func = output_transformers[i]
+                    try:
+                        bytes_data = transformer_func(result)
+                        if not isinstance(bytes_data, bytes):
+                            raise ValueError(
+                                f"Transformer {transformer_func.__name__} for {name} must return bytes, "
+                                f"got {type(bytes_data).__name__}"
+                            )
+                        logger.info(f"Transformed {name} using {transformer_func.__name__} to bytes")
+                    except Exception as e:
+                        logger.error(f"Error transforming {name} with {transformer_func.__name__}: {e}")
+                        raise e
+                else:
+                    # No transformer: assume result is already bytes
+                    if not isinstance(result, bytes):
+                        raise ValueError(
+                            f"Return value for {name} must be bytes when no transformer is provided, "
+                            f"got {type(result).__name__}"
+                        )
+                    bytes_data = result
+                    logger.info(f"Using {name} as bytes directly")
+                
+                # Detect content type from binary data (magic bytes)
+                content_type = detect_content_type_from_bytes(bytes_data)
+                logger.info(f"Detected content type for {name}: {content_type}")
+                
+                # Upload to MinIO
+                try:
+                    file_url = _upload_bytes_to_minio(
+                        bytes_data,
+                        object_name=None,  # Auto-generate object name
+                        content_type=content_type,
+                        bucket=bucket
+                    )
+                    logger.info(f"Uploaded {name} to MinIO: {file_url}")
+                    uploaded_urls.append(file_url)
+                except Exception as e:
+                    logger.error(f"Error uploading {name} to MinIO: {e}")
+                    raise e
+
+            # Return tuple if original result was a tuple, otherwise return single value
+            if len(uploaded_urls) == 1:
+                return uploaded_urls[0]
+            return tuple(uploaded_urls)
+
+        return wrapper
+
+    return decorator
+
+
+