add domain to aimodel

Author: Deepthi-Chand

api/admin.py

@@ -59,6 +59,7 @@ class AIModelAdmin(admin.ModelAdmin):
     list_filter = (
         "provider",
         "model_type",
+        "domain",
         "status",
         "is_public",
         "is_active",
@@ -85,7 +86,7 @@ class AIModelAdmin(admin.ModelAdmin):
             "Schema",
             {"fields": ("input_schema", "output_schema"), "classes": ("collapse",)},
         ),
-        ("Metadata", {"fields": ("tags", "metadata"), "classes": ("collapse",)}),
+        ("Metadata", {"fields": ("tags", "domain", "metadata"), "classes": ("collapse",)}),
         ("Status & Visibility", {"fields": ("status", "is_public", "is_active")}),
         (
             "Performance Metrics",

api/models/AIModel.py

@@ -12,6 +12,7 @@
     EndpointAuthType,
     EndpointHTTPMethod,
     HFModelClass,
+    PromptDomain,
 )
 
 User = get_user_model()
@@ -107,6 +108,13 @@ class AIModel(models.Model):
     tags = models.ManyToManyField("api.Tag", blank=True)
     sectors = models.ManyToManyField("api.Sector", blank=True, related_name="ai_models")
     geographies = models.ManyToManyField("api.Geography", blank=True, related_name="ai_models")
+    domain = models.CharField(
+        max_length=200,
+        choices=PromptDomain.choices,
+        blank=True,
+        null=True,
+        help_text="Domain or category (e.g., healthcare, education, legal)",
+    )
     metadata = models.JSONField(
         default=dict,
         blank=True,

api/schema/aimodel_schema.py

@@ -1,6 +1,6 @@
 """GraphQL schema for AI Model."""
 
-# mypy: disable-error-code=union-attr
+# mypy: disable-error-code="union-attr,misc"
 
 import datetime
 from typing import List, Optional
@@ -27,6 +27,7 @@
     AIModelVersionOrder,
     EndpointAuthTypeEnum,
     EndpointHTTPMethodEnum,
+    PromptDomainEnum,
     TypeAIModel,
     TypeAIModelVersion,
     TypeModelEndpoint,
@@ -95,6 +96,7 @@ class CreateAIModelInput:
     tags: Optional[List[str]] = None
     sectors: Optional[List[str]] = None
     geographies: Optional[List[int]] = None
+    domain: Optional[PromptDomainEnum] = None
     metadata: Optional[strawberry.scalars.JSON] = None
     is_public: bool = False
 
@@ -119,6 +121,7 @@ class UpdateAIModelInput:
     tags: Optional[List[str]] = None
     sectors: Optional[List[str]] = None
     geographies: Optional[List[int]] = None
+    domain: Optional[PromptDomainEnum] = None
     metadata: Optional[strawberry.scalars.JSON] = None
     is_public: Optional[bool] = None
     is_active: Optional[bool] = None
@@ -441,6 +444,7 @@ def create_ai_model(
                 supported_languages=supported_languages,
                 input_schema=input_schema,
                 output_schema=output_schema,
+                domain=input.domain if input.domain else None,
                 metadata=metadata,
                 is_public=input.is_public,
                 status="REGISTERED",
@@ -518,6 +522,8 @@ def update_ai_model(
             model.input_schema = input.input_schema
         if input.output_schema is not None:
             model.output_schema = input.output_schema
+        if input.domain is not None:
+            model.domain = input.domain
         if input.metadata is not None:
             model.metadata = input.metadata
         if input.is_public is not None:

api/types/type_aimodel.py

@@ -26,6 +26,7 @@
     EndpointAuthType,
     EndpointHTTPMethod,
     HFModelClass,
+    PromptDomain,
 )
 from authorization.types import TypeUser
 
@@ -41,6 +42,7 @@
 AIModelFrameworkEnum = strawberry.enum(AIModelFramework)  # type: ignore
 HFModelClassEnum = strawberry.enum(HFModelClass)  # type: ignore
 AIModelLifecycleStageEnum = strawberry.enum(AIModelLifecycleStage)  # type: ignore
+PromptDomainEnum = strawberry.enum(PromptDomain)  # type: ignore
 
 
 @strawberry.type
@@ -83,6 +85,7 @@ class AIModelFilter:
     status: Optional[AIModelStatusEnum]
     model_type: Optional[AIModelTypeEnum]
     provider: Optional[AIModelProviderEnum]
+    domain: Optional[PromptDomainEnum]
     is_public: Optional[bool]
     is_active: Optional[bool]
 
@@ -123,6 +126,7 @@ class TypeAIModel(BaseType):
     supported_languages: Optional[strawberry.scalars.JSON]
     input_schema: Optional[strawberry.scalars.JSON]
     output_schema: Optional[strawberry.scalars.JSON]
+    domain: Optional[PromptDomainEnum]
     metadata: Optional[strawberry.scalars.JSON]
     status: AIModelStatusEnum
     is_public: bool

api/views/aimodel_detail.py

@@ -1,8 +1,8 @@
 """API view for AI Model detail."""
 
+import logging
 from typing import Any, Dict, List, Optional
 
-import logging
 from rest_framework import serializers, status
 from rest_framework.permissions import AllowAny
 from rest_framework.request import Request
@@ -11,9 +11,9 @@
 
 from api.models.AIModel import AIModel, ModelEndpoint
 
-
 logger = logging.getLogger(__name__)
 
+
 class ModelEndpointSerializer(serializers.ModelSerializer):
     """Serializer for Model Endpoint."""
 
@@ -59,6 +59,7 @@ class Meta:
             "tags",
             "sectors",
             "geographies",
+            "domain",
             "metadata",
             "status",
             "is_public",

api/views/search_aimodel.py

@@ -103,6 +103,7 @@ class Meta:
             "is_individual_model",
             "has_active_endpoints",
             "endpoint_count",
+            "domain",
             "version_count",
             "lifecycle_stage",
             "all_providers",
@@ -154,6 +155,7 @@ def get_searchable_and_aggregations(self) -> Tuple[List[str], Dict[str, str]]:
             "supports_streaming": "terms",
             "lifecycle_stage": "terms",
             "all_providers": "terms",
+            "domain": "terms",
         }
 
         return searchable_fields, aggregations
@@ -208,6 +210,7 @@ def add_filters(self, filters: Dict[str, str], search: Search) -> Search:
                 "supported_languages",
                 "lifecycle_stage",
                 "all_providers",
+                "domain",
             ]:
                 # Handle single or multi-value filters
                 filter_values = filters[filter_key].split(",")

dataspace_sdk/resources/aimodels.py

@@ -17,6 +17,7 @@ def search(
         status: Optional[str] = None,
         model_type: Optional[str] = None,
         provider: Optional[str] = None,
+        domain: Optional[str] = None,
         sort: Optional[str] = None,
         page: int = 1,
         page_size: int = 10,
@@ -32,6 +33,7 @@ def search(
             status: Filter by status (ACTIVE, INACTIVE, etc.)
             model_type: Filter by model type (LLM, VISION, etc.)
             provider: Filter by provider (OPENAI, ANTHROPIC, etc.)
+            domain: Filter by domain (HEALTHCARE, EDUCATION, etc.)
             sort: Sort order (recent, alphabetical)
             page: Page number (1-indexed)
             page_size: Number of results per page
@@ -58,6 +60,8 @@ def search(
             params["model_type"] = model_type
         if provider:
             params["provider"] = provider
+        if domain:
+            params["domain"] = domain
         if sort:
             params["sort"] = sort
 
@@ -94,6 +98,7 @@ def get_by_id_graphql(self, model_id: str) -> Dict[str, Any]:
                 displayName
                 description
                 modelType
+                domain
                 status
                 isPublic
                 createdAt
@@ -207,6 +212,7 @@ def list_all(
                 displayName
                 description
                 modelType
+                domain
                 status
                 isPublic
                 createdAt

search/documents/aimodel_document.py

@@ -54,6 +54,7 @@ class AIModelDocument(Document):
 
     # Model configuration
     model_type = fields.KeywordField()
+    domain = fields.KeywordField()
 
     # Status and visibility
     status = fields.KeywordField()