bakery-ia/shared/clients/training_client.py

# shared/clients/training_client.py
"""
Training Service Client
Handles all API calls to the training service
"""

from typing import Dict, Any, Optional, List
from .base_service_client import BaseServiceClient
from shared.config.base import BaseServiceSettings


class TrainingServiceClient(BaseServiceClient):
    """Client for communicating with the training service"""
    
    def __init__(self, config: BaseServiceSettings, calling_service_name: str = "unknown"):
        super().__init__(calling_service_name, config)
    
    def get_service_base_path(self) -> str:
        return "/api/v1"
    
    # ================================================================
    # TRAINING JOBS
    # ================================================================
    
    async def create_training_job(
        self,
        tenant_id: str,
        include_weather: bool = True,
        include_traffic: bool = False,
        min_data_points: int = 30,
        **kwargs
    ) -> Optional[Dict[str, Any]]:
        """Create a new training job"""
        data = {
            "include_weather": include_weather,
            "include_traffic": include_traffic,
            "min_data_points": min_data_points,
            **kwargs
        }
        return await self.post("jobs", data=data, tenant_id=tenant_id)
    
    async def get_training_job(self, tenant_id: str, job_id: str) -> Optional[Dict[str, Any]]:
        """Get training job details"""
        return await self.get(f"jobs/{job_id}", tenant_id=tenant_id)
    
    async def list_training_jobs(
        self,
        tenant_id: str,
        status: Optional[str] = None,
        limit: int = 50
    ) -> Optional[List[Dict[str, Any]]]:
        """List training jobs for a tenant"""
        params = {"limit": limit}
        if status:
            params["status"] = status
        
        result = await self.get("jobs", tenant_id=tenant_id, params=params)
        return result.get("jobs", []) if result else None
    
    async def cancel_training_job(self, tenant_id: str, job_id: str) -> Optional[Dict[str, Any]]:
        """Cancel a training job"""
        return await self.delete(f"jobs/{job_id}", tenant_id=tenant_id)
    
    # ================================================================
    # MODELS
    # ================================================================
    
    async def get_model(self, tenant_id: str, model_id: str) -> Optional[Dict[str, Any]]:
        """Get model details"""
        return await self.get(f"models/{model_id}", tenant_id=tenant_id)
    
    async def list_models(
        self,
        tenant_id: str,
        status: Optional[str] = None,
        model_type: Optional[str] = None,
        limit: int = 50
    ) -> Optional[List[Dict[str, Any]]]:
        """List models for a tenant"""
        params = {"limit": limit}
        if status:
            params["status"] = status
        if model_type:
            params["model_type"] = model_type
        
        result = await self.get("models", tenant_id=tenant_id, params=params)
        return result.get("models", []) if result else None
    
    async def get_active_model_for_product(
        self,
        tenant_id: str,
        product_name: str
    ) -> Optional[Dict[str, Any]]:
        """
        Get the active model for a specific product
        This is the preferred method since models are stored per product.
        """
        result = await self.get(f"models/{product_name}/active", tenant_id=tenant_id)
        return result
    
    async def deploy_model(self, tenant_id: str, model_id: str) -> Optional[Dict[str, Any]]:
        """Deploy a trained model"""
        return await self.post(f"models/{model_id}/deploy", data={}, tenant_id=tenant_id)
    
    async def delete_model(self, tenant_id: str, model_id: str) -> Optional[Dict[str, Any]]:
        """Delete a model"""
        return await self.delete(f"models/{model_id}", tenant_id=tenant_id)
    
    # ================================================================
    # MODEL METRICS & PERFORMANCE
    # ================================================================
    
    async def get_model_metrics(self, tenant_id: str, model_id: str) -> Optional[Dict[str, Any]]:
        """Get model performance metrics"""
        return await self.get(f"models/{model_id}/metrics", tenant_id=tenant_id)
    
    async def get_model_predictions(
        self,
        tenant_id: str,
        model_id: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
    ) -> Optional[List[Dict[str, Any]]]:
        """Get model predictions for evaluation"""
        params = {}
        if start_date:
            params["start_date"] = start_date
        if end_date:
            params["end_date"] = end_date
        
        result = await self.get(f"models/{model_id}/predictions", tenant_id=tenant_id, params=params)
        return result.get("predictions", []) if result else None