""" Type serialization system for Python side of the bridge. Supports both JSON and binary serialization for efficient handling of large numerical data like tensors and embeddings. """ # Try to import orjson for 6x performance boost, fallback to stdlib json try: import orjson _use_orjson = True except ImportError: import json _use_orjson = False import pickle import numpy as np from typing import Any, Dict, Union, Tuple, Optional from google.protobuf import any_pb2 # Size threshold for using binary serialization (10KB) BINARY_THRESHOLD = 10_240 class TypeSerializer: """Unified type serialization for Python side.""" @staticmethod def encode_any(value: Any, var_type: str) -> Tuple[any_pb2.Any, Optional[bytes]]: """ Encode a Python value to protobuf Any with optional binary data. Returns: Tuple of (Any message, optional binary data) """ # Normalize value based on type normalized = TypeSerializer._normalize_value(value, var_type) # Check if we should use binary serialization if TypeSerializer._should_use_binary(normalized, var_type): return TypeSerializer._encode_with_binary(normalized, var_type) else: # Standard JSON serialization json_str = TypeSerializer._serialize_value(normalized, var_type) # Create Any message any_msg = any_pb2.Any() any_msg.type_url = f"type.googleapis.com/snakepit.{var_type}" any_msg.value = json_str.encode('utf-8') return any_msg, None @staticmethod def decode_any(any_msg: any_pb2.Any, binary_data: Optional[bytes] = None) -> Any: """ Decode protobuf Any to Python value with optional binary data. Args: any_msg: Protobuf Any message binary_data: Optional binary data for large values Returns: Decoded Python value """ # Check if this is a binary-encoded value if any_msg.type_url.endswith('.binary') and binary_data is not None: return TypeSerializer._decode_with_binary(any_msg, binary_data) else: # Standard JSON decoding # Extract type from URL - handle various formats: # - "type.googleapis.com/snakepit.float" -> "float" # - "dspex.variables/float" -> "float" # - "type.googleapis.com/google.protobuf.StringValue" -> "StringValue" type_url = any_msg.type_url # First split by / to get the type part if '/' in type_url: type_part = type_url.split('/')[-1] # e.g., "snakepit.float" or "float" else: type_part = type_url # No slash, use as-is # Then split by . to get the final type var_type = type_part.split('.')[-1] # e.g., "float" # Decode JSON json_str = any_msg.value.decode('utf-8') value = TypeSerializer._deserialize_json(json_str) # Convert to appropriate Python type return TypeSerializer._deserialize_value(value, var_type) @staticmethod def _normalize_value(value: Any, var_type: str) -> Any: """Normalize Python values for consistency.""" if var_type == 'float': if isinstance(value, (int, float)): return float(value) raise ValueError(f"Expected number, got {type(value)}") elif var_type == 'integer': if isinstance(value, (int, float)): if isinstance(value, float) and value.is_integer(): return int(value) elif isinstance(value, int): return value raise ValueError(f"Expected integer, got {value}") elif var_type == 'string': return str(value) elif var_type == 'boolean': if isinstance(value, bool): return value raise ValueError(f"Expected boolean, got {type(value)}") elif var_type == 'choice': return str(value) elif var_type == 'module': return str(value) elif var_type == 'embedding': if isinstance(value, np.ndarray): return value.tolist() elif isinstance(value, list): return [float(x) for x in value] raise ValueError(f"Expected array/list, got {type(value)}") elif var_type == 'tensor': if isinstance(value, np.ndarray): return { 'shape': list(value.shape), 'data': value.tolist() } elif isinstance(value, dict) and 'shape' in value and 'data' in value: return value raise ValueError(f"Expected tensor, got {type(value)}") else: return value @staticmethod def _serialize_value(value: Any, var_type: str) -> str: """ Serialize normalized value to JSON string. Uses orjson for 6x performance boost if available, falls back to stdlib json otherwise. """ # Handle special float values if var_type == 'float': if isinstance(value, float): if np.isnan(value): value_to_serialize = "NaN" elif np.isinf(value): value_to_serialize = "Infinity" if value > 0 else "-Infinity" else: value_to_serialize = value else: value_to_serialize = value else: value_to_serialize = value # Use orjson if available, otherwise stdlib json if _use_orjson: # orjson returns bytes, must decode to str return orjson.dumps(value_to_serialize).decode('utf-8') else: import json return json.dumps(value_to_serialize) @staticmethod def _deserialize_json(json_str: str) -> Any: """ Deserialize JSON string to Python value. Uses orjson for 6x performance boost if available, falls back to stdlib json otherwise. """ if _use_orjson: # orjson.loads accepts str or bytes return orjson.loads(json_str) else: import json return json.loads(json_str) @staticmethod def _deserialize_value(value: Any, var_type: str) -> Any: """Convert JSON-decoded value to appropriate Python type.""" if var_type == 'float': if value == "NaN": return float('nan') elif value == "Infinity": return float('inf') elif value == "-Infinity": return float('-inf') return float(value) elif var_type == 'integer': return int(value) elif var_type == 'embedding': # Could convert back to numpy array return value elif var_type == 'tensor': # Could reconstruct numpy array if isinstance(value, dict) and 'data' in value and 'shape' in value: data = np.array(value['data']) return data.reshape(value['shape']) return value else: return value @staticmethod def validate_constraints(value: Any, var_type: str, constraints: Dict) -> None: """Validate value against type constraints.""" if var_type == 'float' or var_type == 'integer': min_val = constraints.get('min') max_val = constraints.get('max') if min_val is not None and value < min_val: raise ValueError(f"Value {value} is below minimum {min_val}") if max_val is not None and value > max_val: raise ValueError(f"Value {value} is above maximum {max_val}") elif var_type == 'string': min_len = constraints.get('min_length', 0) max_len = constraints.get('max_length') length = len(value) if length < min_len: raise ValueError(f"String too short: {length} < {min_len}") if max_len and length > max_len: raise ValueError(f"String too long: {length} > {max_len}") elif var_type == 'choice': choices = constraints.get('choices', []) if choices and value not in choices: raise ValueError(f"Value {value} not in allowed choices: {choices}") elif var_type == 'module': allowed_modules = constraints.get('allowed_modules', []) if allowed_modules and value not in allowed_modules: raise ValueError(f"Module {value} not in allowed modules: {allowed_modules}") elif var_type == 'embedding': dimensions = constraints.get('dimensions') if dimensions and len(value) != dimensions: raise ValueError(f"Wrong dimensions: {len(value)} != {dimensions}") elif var_type == 'tensor': expected_shape = constraints.get('shape') if expected_shape: actual_shape = value.get('shape') if isinstance(value, dict) else list(value.shape) if actual_shape != expected_shape: raise ValueError(f"Wrong shape: {actual_shape} != {expected_shape}") @staticmethod def _should_use_binary(value: Any, var_type: str) -> bool: """Check if value should use binary serialization.""" if var_type not in ['tensor', 'embedding']: return False # Estimate size if var_type == 'tensor': if isinstance(value, dict) and 'data' in value: data = value['data'] if isinstance(data, list): estimated_size = len(data) * 8 # 8 bytes per float return estimated_size > BINARY_THRESHOLD elif var_type == 'embedding': if isinstance(value, list): estimated_size = len(value) * 8 return estimated_size > BINARY_THRESHOLD return False @staticmethod def _encode_with_binary(value: Any, var_type: str) -> Tuple[any_pb2.Any, bytes]: """Encode large data with binary serialization.""" if var_type == 'tensor': shape = value.get('shape', []) data = value.get('data', []) # Create metadata metadata = { 'shape': shape, 'dtype': 'float32', 'binary_format': 'pickle', 'type': var_type } # Create Any message with metadata any_msg = any_pb2.Any() any_msg.type_url = f"type.googleapis.com/snakepit.{var_type}.binary" any_msg.value = TypeSerializer._serialize_value(metadata, 'string').encode('utf-8') # Serialize data as binary binary_data = pickle.dumps(data, protocol=pickle.HIGHEST_PROTOCOL) return any_msg, binary_data elif var_type == 'embedding': # Create metadata metadata = { 'shape': [len(value)], 'dtype': 'float32', 'binary_format': 'pickle', 'type': var_type } # Create Any message with metadata any_msg = any_pb2.Any() any_msg.type_url = f"type.googleapis.com/snakepit.{var_type}.binary" any_msg.value = TypeSerializer._serialize_value(metadata, 'string').encode('utf-8') # Serialize data as binary binary_data = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL) return any_msg, binary_data else: raise ValueError(f"Binary encoding not supported for type: {var_type}") @staticmethod def _decode_with_binary(any_msg: any_pb2.Any, binary_data: bytes) -> Any: """Decode binary-encoded data.""" # Extract base type (remove .binary suffix) type_parts = any_msg.type_url.split('.') var_type = type_parts[-2] # Get type before .binary # Decode metadata metadata = TypeSerializer._deserialize_json(any_msg.value.decode('utf-8')) # Deserialize binary data data = pickle.loads(binary_data) # Reconstruct value based on type if var_type == 'tensor': return { 'shape': metadata.get('shape', []), 'data': data } elif var_type == 'embedding': return data else: raise ValueError(f"Binary decoding not supported for type: {var_type}")