add lifecycle stage to aimodel

Author: Deepthi-Chand

api/models/AIModelVersion.py

@@ -6,7 +6,7 @@
 
 from django.db import models
 
-from api.utils.enums import AIModelStatus
+from api.utils.enums import AIModelLifecycleStage, AIModelStatus
 
 if TYPE_CHECKING:
     from django.db.models import QuerySet
@@ -36,12 +36,18 @@ class AIModelVersion(models.Model):
     output_schema = models.JSONField(default=dict, help_text="Expected output format")
     metadata = models.JSONField(default=dict, help_text="Additional version-specific metadata")
 
-    # Status
+    # Status & Lifecycle
     status = models.CharField(
         max_length=20,
         choices=AIModelStatus.choices,
         default=AIModelStatus.REGISTERED,
     )
+    lifecycle_stage = models.CharField(
+        max_length=20,
+        choices=AIModelLifecycleStage.choices,
+        default=AIModelLifecycleStage.DEVELOPMENT,
+        help_text="Current lifecycle stage of this version",
+    )
     is_latest = models.BooleanField(default=False, help_text="Whether this is the latest version")
 
     # Timestamps

api/schema/aimodel_schema.py

@@ -18,6 +18,7 @@
 from api.schema.extensions import TrackActivity, TrackModelActivity
 from api.types.type_aimodel import (
     AIModelFilter,
+    AIModelLifecycleStageEnum,
     AIModelOrder,
     AIModelProviderEnum,
     AIModelStatusEnum,
@@ -169,13 +170,15 @@ class CreateAIModelVersionInput:
     model_id: int
     version: str
     version_notes: Optional[str] = ""
+    lifecycle_stage: Optional[AIModelLifecycleStageEnum] = None
     supports_streaming: bool = False
     max_tokens: Optional[int] = None
     supported_languages: Optional[List[str]] = None
     input_schema: Optional[strawberry.scalars.JSON] = None
     output_schema: Optional[strawberry.scalars.JSON] = None
     metadata: Optional[strawberry.scalars.JSON] = None
     copy_from_version_id: Optional[int] = None
+    is_latest: Optional[bool] = None
 
 
 @strawberry.input
@@ -185,6 +188,7 @@ class UpdateAIModelVersionInput:
     id: int
     version: Optional[str] = None
     version_notes: Optional[str] = None
+    lifecycle_stage: Optional[AIModelLifecycleStageEnum] = None
     supports_streaming: Optional[bool] = None
     max_tokens: Optional[int] = None
     supported_languages: Optional[List[str]] = None
@@ -748,14 +752,15 @@ def create_ai_model_version(
             ai_model=model,
             version=input.version,
             version_notes=input.version_notes or "",
+            lifecycle_stage=input.lifecycle_stage.value if input.lifecycle_stage else "DEVELOPMENT",  # type: ignore[misc]
             supports_streaming=input.supports_streaming,
             max_tokens=input.max_tokens,
             supported_languages=input.supported_languages or [],
             input_schema=input.input_schema or {},
             output_schema=input.output_schema or {},
             metadata=input.metadata or {},
             status="DRAFT",
-            is_latest=True,
+            is_latest=input.is_latest if input.is_latest is not None else True,
         )
 
         # If copy_from_version_id is provided, copy all providers
@@ -803,6 +808,8 @@ def update_ai_model_version(
             version.version = input.version
         if input.version_notes is not None:
             version.version_notes = input.version_notes
+        if input.lifecycle_stage is not None:
+            version.lifecycle_stage = input.lifecycle_stage.value  # type: ignore[misc]
         if input.supports_streaming is not None:
             version.supports_streaming = input.supports_streaming
         if input.max_tokens is not None:

api/types/type_aimodel.py

@@ -19,6 +19,7 @@
 from api.types.type_sector import TypeSector
 from api.utils.enums import (
     AIModelFramework,
+    AIModelLifecycleStage,
     AIModelProvider,
     AIModelStatus,
     AIModelType,
@@ -39,6 +40,7 @@
 EndpointHTTPMethodEnum = strawberry.enum(EndpointHTTPMethod)  # type: ignore
 AIModelFrameworkEnum = strawberry.enum(AIModelFramework)  # type: ignore
 HFModelClassEnum = strawberry.enum(HFModelClass)  # type: ignore
+AIModelLifecycleStageEnum = strawberry.enum(AIModelLifecycleStage)  # type: ignore
 
 
 @strawberry.type
@@ -195,7 +197,8 @@ def primary_endpoint(self) -> Optional[TypeModelEndpoint]:
     def versions(self) -> List["TypeAIModelVersion"]:
         """Get all versions of this AI model."""
         try:
-            queryset = self.versions.all()  # type: ignore
+            django_instance = cast(AIModel, self)
+            queryset = django_instance.versions.all()
             return TypeAIModelVersion.from_django_list(list(queryset))
         except Exception:
             return []
@@ -204,9 +207,10 @@ def versions(self) -> List["TypeAIModelVersion"]:
     def latest_version(self) -> Optional["TypeAIModelVersion"]:
         """Get the latest version of this AI model."""
         try:
-            version = self.versions.filter(is_latest=True).first()  # type: ignore
+            django_instance = cast(AIModel, self)
+            version = django_instance.versions.filter(is_latest=True).first()
             if not version:
-                version = self.versions.order_by("-created_at").first()  # type: ignore
+                version = django_instance.versions.order_by("-created_at").first()
             if version:
                 return TypeAIModelVersion.from_django(version)
             return None
@@ -265,6 +269,7 @@ class TypeAIModelVersion(BaseType):
     output_schema: strawberry.scalars.JSON
     metadata: strawberry.scalars.JSON
     status: AIModelStatusEnum
+    lifecycle_stage: AIModelLifecycleStageEnum
     is_latest: bool
     created_at: datetime
     updated_at: datetime

api/utils/enums.py

@@ -189,6 +189,16 @@ class AIModelProvider(models.TextChoices):
     HUGGINGFACE = "HUGGINGFACE"
 
 
+class AIModelLifecycleStage(models.TextChoices):
+    DEVELOPMENT = "DEVELOPMENT", "Development"
+    TESTING = "TESTING", "Testing"
+    BETA = "BETA", "Beta Testing"
+    STAGING = "STAGING", "Staging"
+    PRODUCTION = "PRODUCTION", "Production"
+    DEPRECATED = "DEPRECATED", "Deprecated"
+    RETIRED = "RETIRED", "Retired"
+
+
 class AIModelFramework(models.TextChoices):
     PYTORCH = "pt", "PyTorch"
     TENSORFLOW = "tf", "TensorFlow"