custom_types.py

"""
Strict type definitions using newtype pattern for maximum type safety.
This module defines all the domain-specific types used throughout the application.
"""

from typing import NewType, List, Tuple, Dict, Any, Union, Literal
from dataclasses import dataclass
from abc import ABC, abstractmethod


# === Base Protocol for Newtype Pattern ===

class Unwrappable(ABC):
    """Protocol for types that can be unwrapped."""

    @abstractmethod
    def unwrap(self) -> Any:
        """Unwrap the newtype to get the underlying value."""
        pass


# === Core Domain Types ===

# Model-related types
ModelId = NewType('ModelId', str)
ModelName = NewType('ModelName', str)
ModelMemoryUsage = NewType('ModelMemoryUsage', float)  # GB
ModelPriority = NewType('ModelPriority', int)

# Conversation-related types
ConversationId = NewType('ConversationId', str)
UserMessage = NewType('UserMessage', str)
AssistantResponse = NewType('AssistantResponse', str)
Prompt = NewType('Prompt', str)
MessagePair = NewType('MessagePair', Tuple[UserMessage, AssistantResponse])
ConversationHistory = NewType('ConversationHistory', List[MessagePair])

# Generation parameters
Temperature = NewType('Temperature', float)
TopP = NewType('TopP', float)
RepetitionPenalty = NewType('RepetitionPenalty', float)
NoRepeatNGramSize = NewType('NoRepeatNGramSize', int)
MaxNewTokens = NewType('MaxNewTokens', int)
TokensUsed = NewType('TokensUsed', int)
ContextLength = NewType('ContextLength', int)

# System resource types
GPUMemoryUsage = NewType('GPUMemoryUsage', float)  # Fraction 0.0-1.0
GPUMemoryTotal = NewType('GPUMemoryTotal', float)  # GB
CPUMemoryUsage = NewType('CPUMemoryUsage', float)  # Fraction 0.0-1.0
Timestamp = NewType('Timestamp', float)

# API types
APIHost = NewType('APIHost', str)
APIPort = NewType('APIPort', int)
HFToken = NewType('HFToken', str)
LogLevel = NewType('LogLevel', str)

# Configuration types
MaxConversations = NewType('MaxConversations', int)
MaxMessagesPerConversation = NewType('MaxMessagesPerConversation', int)
MaxModelsInMemory = NewType('MaxModelsInMemory', int)


# === Wrapped Types with Newtype Pattern ===

@dataclass(frozen=True)
class ModelIdWrapper:
    """Wrapper for ModelId with validation."""
    _value: ModelId

    def __post_init__(self) -> None:
        if not self._value.strip():
            raise ValueError("ModelId cannot be empty")

        # Allow any thoughtcast/* model (wildcard support)
        if self._value.startswith('thoughtcast/'):
            if len(self._value.split('/')) != 2 or not self._value.split('/')[1].strip():
                raise ValueError("thoughtcast models must be in format 'thoughtcast/model-name'")
            return

        # Allow any doma-dev/* model (wildcard support)
        if self._value.startswith('doma-dev/'):
            if len(self._value.split('/')) != 2 or not self._value.split('/')[1].strip():
                raise ValueError("doma-dev models must be in format 'doma-dev/model-name'")
            return

        # Allow any mistralai/* model (wildcard support)
        if self._value.startswith('mistralai/'):
            if len(self._value.split('/')) != 2 or not self._value.split('/')[1].strip():
                raise ValueError("mistralai models must be in format 'mistralai/model-name'")
            return

        # Allow specific whitelisted models
        allowed_models = {
            'TinyLlama/TinyLlama-1.1B-Chat-v1.0',
            'microsoft/DialoGPT-medium',
            'microsoft/DialoGPT-small',
        }

        if self._value in allowed_models:
            return

        # For other models, require org/model format and explicit allowlisting
        if '/' not in self._value:
            raise ValueError("ModelId must be in format 'org/model' or be a whitelisted model")

        org, model_name = self._value.split('/', 1)
        if not org.strip() or not model_name.strip():
            raise ValueError("ModelId must have valid organization and model name")

        raise ValueError(
            f"Model '{self._value}' is not allowed. Only thoughtcast/*, doma-dev/*, mistralai/*, or explicitly whitelisted models are supported."
        )

    def unwrap(self) -> ModelId:
        return self._value

    def __str__(self) -> str:  # type: ignore[override]
        return str(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"ModelIdWrapper({self._value!r})"


@dataclass(frozen=True)
class ConversationIdWrapper:
    """Wrapper for ConversationId with validation."""
    _value: ConversationId

    def __post_init__(self) -> None:
        if not self._value.strip():
            raise ValueError("ConversationId cannot be empty")
        if len(self._value) > 128:
            raise ValueError("ConversationId cannot be longer than 128 characters")

    def unwrap(self) -> ConversationId:
        return self._value

    def __str__(self) -> str:  # type: ignore[override]
        return str(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"ConversationIdWrapper({self._value!r})"


@dataclass(frozen=True)
class TemperatureWrapper:
    """Wrapper for Temperature with validation."""
    _value: Temperature

    def __post_init__(self) -> None:
        if not (0.0 <= self._value <= 2.0):
            raise ValueError("Temperature must be between 0.0 and 2.0")

    def unwrap(self) -> Temperature:
        return self._value

    def __float__(self) -> float:
        return float(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"TemperatureWrapper({self._value!r})"


@dataclass(frozen=True)
class TopPWrapper:
    """Wrapper for TopP with validation."""
    _value: TopP

    def __post_init__(self) -> None:
        if not (0.0 <= self._value <= 1.0):
            raise ValueError("TopP must be between 0.0 and 1.0")

    def unwrap(self) -> TopP:
        return self._value

    def __float__(self) -> float:
        return float(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"TopPWrapper({self._value!r})"


@dataclass(frozen=True)
class RepetitionPenaltyWrapper:
    """Wrapper for RepetitionPenalty with validation."""
    _value: RepetitionPenalty

    def __post_init__(self) -> None:
        if not (0.1 <= self._value <= 2.0):
            raise ValueError("RepetitionPenalty must be between 0.1 and 2.0")

    def unwrap(self) -> RepetitionPenalty:
        return self._value

    def __float__(self) -> float:
        return float(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"RepetitionPenaltyWrapper({self._value!r})"


@dataclass(frozen=True)
class MaxNewTokensWrapper:
    """Wrapper for MaxNewTokens with validation."""
    _value: MaxNewTokens

    def __post_init__(self) -> None:
        if not (1 <= self._value <= 2048):
            raise ValueError("MaxNewTokens must be between 1 and 2048")

    def unwrap(self) -> MaxNewTokens:
        return self._value

    def __int__(self) -> int:
        return int(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"MaxNewTokensWrapper({self._value!r})"


@dataclass(frozen=True)
class NoRepeatNGramSizeWrapper:
    """Wrapper for NoRepeatNGramSize with validation."""
    _value: NoRepeatNGramSize

    def __post_init__(self) -> None:
        if not (0 <= self._value <= 10):
            raise ValueError("NoRepeatNGramSize must be between 0 and 10")

    def unwrap(self) -> NoRepeatNGramSize:
        return self._value

    def __int__(self) -> int:
        return int(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"NoRepeatNGramSizeWrapper({self._value!r})"


@dataclass(frozen=True)
class GPUMemoryUsageWrapper:
    """Wrapper for GPUMemoryUsage with validation."""
    _value: GPUMemoryUsage

    def __post_init__(self) -> None:
        if not (0.0 <= self._value <= 1.0):
            raise ValueError("GPUMemoryUsage must be between 0.0 and 1.0")

    def unwrap(self) -> GPUMemoryUsage:
        return self._value

    def __float__(self) -> float:
        return float(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"GPUMemoryUsageWrapper({self._value!r})"


@dataclass(frozen=True)
class UserMessageWrapper:
    """Wrapper for UserMessage with validation."""
    _value: UserMessage

    def __post_init__(self) -> None:
        if not self._value.strip():
            raise ValueError("UserMessage cannot be empty")
        if len(self._value) > 4096:
            raise ValueError("UserMessage cannot be longer than 4096 characters")

    def unwrap(self) -> UserMessage:
        return self._value

    def __str__(self) -> str:  # type: ignore[override]
        return str(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"UserMessageWrapper({self._value!r})"


@dataclass(frozen=True)
class AssistantResponseWrapper:
    """Wrapper for AssistantResponse with validation."""
    _value: AssistantResponse

    def __post_init__(self) -> None:
        if len(self._value) > 4096:
            raise ValueError("AssistantResponse cannot be longer than 4096 characters")

    def unwrap(self) -> AssistantResponse:
        return self._value

    def __str__(self) -> str:  # type: ignore[override]
        return str(self._value)

    def __repr__(self) -> str:  # type: ignore[override]
        return f"AssistantResponseWrapper({self._value!r})"


# === External Library Type Wrappers ===

@dataclass(frozen=True)
class HuggingFaceModel:
    """Wrapper for HuggingFace model objects."""
    _value: Any

    def unwrap(self) -> Any:
        """Get the underlying HuggingFace model."""
        return self._value

@dataclass(frozen=True)
class HuggingFaceTokenizer:
    """Wrapper for HuggingFace tokenizer objects."""
    _value: Any

    def unwrap(self) -> Any:
        """Get the underlying HuggingFace tokenizer."""
        return self._value

@dataclass(frozen=True)
class HuggingFaceGenerationConfig:
    """Wrapper for HuggingFace generation config objects."""
    _value: Any

    def unwrap(self) -> Any:
        """Get the underlying HuggingFace generation config."""
        return self._value

@dataclass(frozen=True)
class TorchTensor:
    """Wrapper for PyTorch tensor objects."""
    _value: Any

    def unwrap(self) -> Any:
        """Get the underlying PyTorch tensor."""
        return self._value

@dataclass(frozen=True)
class TorchDevice:
    """Wrapper for PyTorch device objects."""
    _value: Any

    def unwrap(self) -> Any:
        """Get the underlying PyTorch device."""
        return self._value

# === Type Aliases for Complex Types ===

ModelDict = Dict[ModelId, Any]
ConversationDict = Dict[ConversationId, Any]
GenerationConfig = Dict[str, Union[int, float, bool]]
ModelKwargs = Dict[str, Any]


# === Factory Functions ===

def create_model_id(value: str) -> ModelIdWrapper:
    """Create a validated ModelId."""
    return ModelIdWrapper(ModelId(value))


def create_conversation_id(value: str) -> ConversationIdWrapper:
    """Create a validated ConversationId."""
    return ConversationIdWrapper(ConversationId(value))


def create_temperature(value: float) -> TemperatureWrapper:
    """Create a validated Temperature."""
    return TemperatureWrapper(Temperature(value))


def create_top_p(value: float) -> TopPWrapper:
    """Create a validated TopP."""
    return TopPWrapper(TopP(value))


def create_repetition_penalty(value: float) -> RepetitionPenaltyWrapper:
    """Create a validated RepetitionPenalty."""
    return RepetitionPenaltyWrapper(RepetitionPenalty(value))


def create_max_new_tokens(value: int) -> MaxNewTokensWrapper:
    """Create a validated MaxNewTokens."""
    return MaxNewTokensWrapper(MaxNewTokens(value))


def create_no_repeat_ngram_size(value: int) -> NoRepeatNGramSizeWrapper:
    """Create a validated NoRepeatNGramSize."""
    return NoRepeatNGramSizeWrapper(NoRepeatNGramSize(value))


def create_gpu_memory_usage(value: float) -> GPUMemoryUsageWrapper:
    """Create a validated GPUMemoryUsage."""
    return GPUMemoryUsageWrapper(GPUMemoryUsage(value))


def create_user_message(value: str) -> UserMessageWrapper:
    """Create a validated UserMessage."""
    return UserMessageWrapper(UserMessage(value))


def create_assistant_response(value: str) -> AssistantResponseWrapper:
    """Create a validated AssistantResponse."""
    return AssistantResponseWrapper(AssistantResponse(value))


# === External Library Factory Functions ===

def create_huggingface_model(value: Any) -> HuggingFaceModel:
    """Create a wrapped HuggingFace model."""
    return HuggingFaceModel(value)

def create_huggingface_tokenizer(value: Any) -> HuggingFaceTokenizer:
    """Create a wrapped HuggingFace tokenizer."""
    return HuggingFaceTokenizer(value)

def create_huggingface_generation_config(value: Any) -> HuggingFaceGenerationConfig:
    """Create a wrapped HuggingFace generation config."""
    return HuggingFaceGenerationConfig(value)

def create_torch_tensor(value: Any) -> TorchTensor:
    """Create a wrapped PyTorch tensor."""
    return TorchTensor(value)

def create_torch_device(value: Any) -> TorchDevice:
    """Create a wrapped PyTorch device."""
    return TorchDevice(value)


# === Status Types ===

HealthStatus = Literal["healthy", "unhealthy", "degraded"]
ModelStatus = Literal["loaded", "unloaded", "loading", "failed"]
ConversationStatus = Literal["active", "inactive", "archived"]


# === Error Types ===

@dataclass(frozen=True)
class ModelError:
    """Represents a model-related error."""
    model_id: ModelIdWrapper
    message: str
    error_type: Literal["loading", "generation", "memory", "validation"]


@dataclass(frozen=True)
class ConversationError:
    """Represents a conversation-related error."""
    conversation_id: ConversationIdWrapper
    message: str
    error_type: Literal["not_found", "validation", "storage"]


# === Result Types ===

@dataclass(frozen=True)
class GenerationResult:
    """Result of text generation."""
    response: AssistantResponseWrapper
    tokens_used: TokensUsed
    model_id: ModelIdWrapper
    conversation_id: ConversationIdWrapper


@dataclass(frozen=True)
class HealthResult:
    """Result of health check."""
    status: HealthStatus
    models_loaded: List[ModelIdWrapper]
    gpu_memory_used: GPUMemoryUsageWrapper
    active_conversations: int
    timestamp: Timestamp


# === Model Information Types ===

@dataclass(frozen=True)
class ModelInformation:
    """Complete information about a model."""
    model_id: ModelIdWrapper
    memory_used: ModelMemoryUsage
    last_used: Timestamp
    status: ModelStatus
    priority: ModelPriority
    context_length: ContextLength