Improve AI logic

services/training/app/models/__init__.py

@@ -18,6 +18,7 @@ from .training import (
     ModelPerformanceMetric,
     TrainingJobQueue,
     ModelArtifact,
+    TrainingPerformanceMetrics,
 )
 
 # List all models for easier access
@@ -27,5 +28,6 @@ __all__ = [
     "ModelPerformanceMetric",
     "TrainingJobQueue",
     "ModelArtifact",
+    "TrainingPerformanceMetrics",
     "AuditLog",
 ]

services/training/app/models/training.py

@@ -150,7 +150,8 @@ class TrainedModel(Base):
     hyperparameters = Column(JSON)  # Store optimized parameters
     features_used = Column(JSON)    # List of regressor columns
     normalization_params = Column(JSON)  # Store feature normalization parameters for consistent predictions
-    
+    product_category = Column(String, nullable=True)  # Product category for category-specific forecasting
+
     # Model status
     is_active = Column(Boolean, default=True)
     is_production = Column(Boolean, default=False)
@@ -185,6 +186,7 @@ class TrainedModel(Base):
             "training_samples": self.training_samples,
             "hyperparameters": self.hyperparameters,
             "features_used": self.features_used,
+            "product_category": self.product_category,
             "is_active": self.is_active,
             "is_production": self.is_production,
             "created_at": self.created_at.isoformat() if self.created_at else None,