from synapse_sdk.plugins.schemas import (
    TrainResult,
    InferenceResult,
    ExportResult,
    UploadResult,
    WeightsResult,
    MetricsResult,
)

from synapse_sdk.plugins.actions.train import BaseTrainAction
from synapse_sdk.plugins.schemas import TrainResult

class MyTrainAction(BaseTrainAction[TrainParams]):
    result_model = TrainResult  # Enable result validation

    def execute(self) -> TrainResult:
        # ... training ...
        return TrainResult(
            weights_path='/models/best.pt',
            final_epoch=100,
            train_metrics={'loss': 0.05},
            val_metrics={'mAP50': 0.85},
        )

from synapse_sdk.plugins.decorators import action
from synapse_sdk.plugins.schemas import InferenceResult

@action(name='infer', result=InferenceResult)
def infer(params: InferParams, ctx: RuntimeContext) -> InferenceResult:
    return InferenceResult(
        predictions=[{'class': 'dog', 'confidence': 0.95}],
        processed_count=100,
    )

class TrainResult(BaseModel):
    weights_path: str            # Path to trained model
    final_epoch: int             # Last completed epoch
    best_epoch: int | None       # Best epoch by val metric
    train_metrics: dict = {}     # Final training metrics
    val_metrics: dict = {}       # Final validation metrics

class InferenceResult(BaseModel):
    predictions: list[dict] = [] # Prediction results
    processed_count: int = 0     # Items processed
    output_path: str | None      # Output file path

class ExportResult(BaseModel):
    output_path: str             # Export path
    exported_count: int          # Items exported
    format: str                  # Export format
    file_size_bytes: int | None  # File size

class UploadResult(BaseModel):
    uploaded_count: int          # Items uploaded
    remote_path: str | None      # Remote URL/path
    status: str = 'completed'    # Upload status

class WeightsResult(BaseModel):
    weights_path: str            # Best/final weights path
    checkpoint_paths: list = []  # Intermediate checkpoints
    format: str = 'pt'           # Weights format

class MetricsResult(BaseModel):
    metrics: dict[str, float]    # Metric values
    category: str = 'default'    # Metrics category