From 8df8392df9fa0ca040f176f566ee7e77b2b4f03b Mon Sep 17 00:00:00 2001
From: Hironsan <light.tree.1.13@gmail.com>
Date: Sat, 30 Jan 2021 17:02:05 +0900
Subject: [PATCH] Add AutoLabelingTemplateAPI to retrieve base templates

---
 app/api/urls.py  |  6 ++++++
 app/api/views.py | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/app/api/urls.py b/app/api/urls.py
index 2e2d9a33..c07c52a4 100644
--- a/app/api/urls.py
+++ b/app/api/urls.py
@@ -11,6 +11,7 @@ from .views import CommentList, CommentDetail
 from .views import TextUploadAPI, TextDownloadAPI, CloudUploadAPI
 from .views import StatisticsAPI
 from .views import RoleMappingList, RoleMappingDetail, Roles
+from .views import AutoLabelingTemplateAPI
 
 urlpatterns = [
     path('health', Health.as_view(), name='health'),
@@ -52,6 +53,11 @@ urlpatterns = [
          RoleMappingList.as_view(), name='rolemapping_list'),
     path('projects/<int:project_id>/roles/<int:rolemapping_id>',
          RoleMappingDetail.as_view(), name='rolemapping_detail'),
+    path(
+        route='projects/<int:project_id>/auto-labeling-templates',
+        view=AutoLabelingTemplateAPI.as_view(),
+        name='auto_labeling_templates'
+    )
 ]
 
 # urlpatterns = format_suffix_patterns(urlpatterns, allowed=['json', 'xml'])
diff --git a/app/api/views.py b/app/api/views.py
index a4f0707c..71fcd7c3 100644
--- a/app/api/views.py
+++ b/app/api/views.py
@@ -2,6 +2,7 @@ import collections
 import json
 import random
 
+from auto_labeling_pipeline.menu import Options
 from django.conf import settings
 from django.contrib.auth.models import User
 from django.db import transaction
@@ -481,3 +482,15 @@ class LabelUploadAPI(APIView):
         except IntegrityError:
             content = {'error': 'IntegrityError: you cannot create a label with same name or shortkey.'}
             return Response(content, status=status.HTTP_400_BAD_REQUEST)
+
+
+class AutoLabelingTemplateAPI(APIView):
+    task_mapping = {'DocumentClassification': 'TextClassification'}
+    permission_classes = [IsAuthenticated & IsProjectAdmin]
+
+    def get(self, request, *args, **kwargs):
+        project = get_object_or_404(Project, pk=self.kwargs['project_id'])
+        task = self.task_mapping.get(project.project_type, project.project_type)
+        options = Options.filter_by_task(task_name=task)
+        option_names = [o.name for o in options]
+        return Response(option_names, status=status.HTTP_200_OK)