Add the endpoint for entry classification with llm; Add mock response…

… in callback;

analysis_module/mock_templates.py

@@ -607,6 +607,30 @@
     ]
 }
 
+MOCK_ENTRY_CLASSIFICATION_LLM: Dict = {
+    "model_tags": {},
+    "geolocations": [
+        {
+            "entity": "Somalia",
+            "meta": {
+                "offset_start": 88,
+                "offset_end": 94,
+                "latitude": -10,
+                "longitude": -55
+            }
+        },
+        {
+            "entity": "Portugal",
+            "meta": {
+                "offset_start": 183,
+                "offset_end": 191,
+                "latitude": 39.6945,
+                "longitude": -8.13057
+            }
+        }
+    ]
+}
+
 
 """
 it's a huge output (and it can be bigger that this one). Maybe we can truncate it.

analysis_module/mockserver.py

@@ -16,7 +16,12 @@
 
 from core.models import NLPRequest
 from core_server.settings import ENDPOINT_NAME
-from .mock_templates import MOCK_ENTRY_CLASSIFICATION, MOCK_ENTRY_CLASSIFICATION_FORMATTED, MOCK_GEOLOCATION  # noqa
+from .mock_templates import (
+    MOCK_ENTRY_CLASSIFICATION,
+    MOCK_ENTRY_CLASSIFICATION_LLM,
+    MOCK_ENTRY_CLASSIFICATION_FORMATTED,
+    MOCK_GEOLOCATION  # noqa
+)
 from .utils import send_callback_url_request
 
 
@@ -526,6 +531,33 @@ def process_entry_classification_mock(body) -> Any:
     except Exception:
         logger.error("Could not send data to callback url", exc_info=True)
 
+def entry_classification_llm_mock(body) -> Any:
+    process_entry_classification_llm_mock.apply_async(
+        args=(body,), countdown=2
+    )
+    return json.dumps({"status": "Successfully received the request."}), 200
+
+@shared_task
+def process_entry_classification_llm_mock(body) -> Any:
+    callback_payload = MOCK_ENTRY_CLASSIFICATION_LLM
+    callback_payload.update({
+        "client_id": body["entries"][0]["client_id"],
+        "model_info": {
+            "id": "all_tags_model",
+            "version": "1.0.0"
+        },
+        "prediction_status": True
+    })
+    callback_url = body["callback_url"]
+    try:
+        requests.post(
+            callback_url,
+            json=callback_payload,
+            timeout=30
+        )
+        logger.info("Successfully send data on callback url for entry classification")
+    except Exception:
+        logger.error("Could not send data to callback url", exc_info=True)
 
 TYPE_ACTIONS_MOCK = {
     "topicmodel": topicmodeling_mock_model,
@@ -535,7 +567,8 @@ def process_entry_classification_mock(body) -> Any:
     "geolocation": geolocation_mock_model,
     "text-extraction": text_extraction_mock,
     "entry-extraction-classification": entry_extraction_mock,
-    "entry-classification": entry_classification_mock
+    "entry-classification": entry_classification_mock,
+    "entry-classification-llm": entry_classification_llm_mock,
 }
 
 

analysis_module/serializers.py

@@ -69,6 +69,20 @@ class PredictionRequestSerializer(serializers.Serializer):
     mock = serializers.BooleanField(default=False)
 
 
+class PredictionEntryWithLLMSerializer(serializers.Serializer):
+    client_id = serializers.CharField()
+    entry = serializers.CharField()
+    project_id = serializers.IntegerField(required=True)
+    af_id = serializers.IntegerField(required=True)
+
+class PredictionRequestWithLLMSerializer(serializers.Serializer):
+    entries = PredictionEntryWithLLMSerializer(many=True)
+    publishing_organization = serializers.CharField()
+    authoring_organization = serializers.ListField()
+    callback_url = serializers.CharField()
+    mock = serializers.BooleanField(default=False)
+
+
 class ExtractionDocumentSerializer(serializers.Serializer):
     url = serializers.CharField()
     client_id = serializers.CharField()

analysis_module/views/predictions.py

@@ -7,7 +7,11 @@
 
 from core_server.settings import IS_MOCKSERVER
 from core.models import NLPRequest
-from analysis_module.serializers import TagsMappingRequestSerializer, PredictionRequestSerializer
+from analysis_module.serializers import (
+    TagsMappingRequestSerializer,
+    PredictionRequestSerializer,
+    PredictionRequestWithLLMSerializer
+)
 # from analysis_module.mockserver import MOCK_ENTRY_CLASSIFICATION
 from analysis_module.utils import send_classification_tags
 from nlp_scripts.model_prediction.tags_mapping import AF2NLPMapping
@@ -120,3 +124,42 @@ def entry_classification(request: Request):
     #     nlp_request.status = NLPRequest.RequestStatus.SUCCESS
     # nlp_request.save(update_fields=['status'])
     # return Response(resp_data)
+
+
+@api_view(["POST"])
+@permission_classes([IsAuthenticated])
+def llm_entry_classification(request: Request):
+    serializer = PredictionRequestWithLLMSerializer(data=request.data)
+    serializer.is_valid(raise_exception=True)
+
+    entries = serializer.validated_data["entries"]
+    req_type = NLPRequest.FeaturesType.ENTRY_CLASSIFICATION_LLM
+
+    if serializer.validated_data.get("mock") or IS_MOCKSERVER:
+        return process_mock_request(
+            request=serializer.validated_data,
+            request_type=req_type
+        )
+    if not entries:
+        return Response({})
+
+    nlp_request = NLPRequest.objects.create(
+        client_id=entries[0]["client_id"],
+        type=NLPRequest.FeaturesType.ENTRY_CLASSIFICATION_LLM,
+        request_params=serializer.validated_data,
+        created_by=request.user,
+    )
+    resp = {
+        "type": req_type,
+        "message": "Request has been successfully queued.",
+        "request_ids": nlp_request.unique_id
+    }
+
+    # transaction.on_commit(
+    #     lambda:
+    # )
+
+    return Response(
+        resp,
+        status=status.HTTP_202_ACCEPTED
+    )

core/models.py

@@ -450,6 +450,7 @@ class FeaturesType(models.TextChoices):
         GEOLOCATION = "geolocation", "Geolocation"
         TAGS_MAPPING = "tags-mapping", "Tags Mapping"
         ENTRY_CLASSIFICATION = "entry-classification", "Entry Classification"
+        ENTRY_CLASSIFICATION_LLM = "entry-classification-llm", "Entry Classification LLM"
         TEXT_EXTRACTION = "text-extraction", "Text Extraction"
         ENTRY_EXTRACTION = "entry-extraction-classification", "Entry Extraction Classification"
 

core_server/urls.py

@@ -27,6 +27,7 @@
 from analysis_module.views.predictions import (
     tags_mapping,
     entry_classification,
+    llm_entry_classification,
     nlp_tags,
 )
 
@@ -51,4 +52,5 @@
     path("api/v1/analysismodule/status/<uuid:unique_id>/", request_status),
     path("api/v1/models-info/", models_detail),
     path("api/v1/test-auth/", token_auth_dummy_view),
+    path("api/v1/llm-entry-classification/", llm_entry_classification)
 ] + static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)