Add dataloader of llm assisted tagging

Author: sudan45

apps/assisted_tagging/admin.py

@@ -2,9 +2,10 @@
 from admin_auto_filters.filters import AutocompleteFilterFactory
 from django.contrib import admin
 
-from assisted_tagging.models import AssistedTaggingModelPredictionTag, AssistedTaggingPrediction, DraftEntry
+from assisted_tagging.models import AssistedTaggingModelPredictionTag, AssistedTaggingPrediction, DraftEntry, LLMAssistedTaggingPredication
 from deep.admin import VersionAdmin
 
+admin.site.register(LLMAssistedTaggingPredication)
 
 @admin.register(DraftEntry)
 class DraftEntryAdmin(VersionAdmin):

apps/assisted_tagging/dataloaders.py

@@ -3,7 +3,7 @@
 
 from django.utils.functional import cached_property
 
-from assisted_tagging.models import AssistedTaggingPrediction
+from assisted_tagging.models import AssistedTaggingPrediction, LLMAssistedTaggingPredication
 
 from utils.graphene.dataloaders import DataLoaderWithContext, WithContextMixin
 
@@ -18,7 +18,22 @@ def batch_load_fn(self, keys):
         return Promise.resolve([_map.get(key, []) for key in keys])
 
 
+class LLMDraftEntryPredicationsLoader(DataLoaderWithContext):
+    def batch_load_fn(self, keys):
+        llm_assisted_tagging_qs = LLMAssistedTaggingPredication.objects\
+            .filter(draft_entry_id__in=keys)
+        _map = {
+            assisted_tagging.draft_entry: assisted_tagging.draft_entry
+            for assisted_tagging in llm_assisted_tagging_qs
+        }
+        return Promise.resolve([_map.get(key) for key in keys])
+
+
 class DataLoaders(WithContextMixin):
     @cached_property
     def draft_entry_predications(self):
         return DraftEntryPredicationsLoader(context=self.context)
+
+    @cached_property
+    def llm_draft_entry_predications(self):
+        return LLMDraftEntryPredicationsLoader(context=self.context)

apps/assisted_tagging/schema.py

@@ -178,9 +178,7 @@ class Meta:
 class DraftEntryType(DjangoObjectType):
     prediction_status = graphene.Field(DraftEntryPredictionStatusEnum, required=True)
     prediction_status_display = EnumDescription(source='get_prediction_status_display', required=True)
-    tags = graphene.Field(
-        graphene.List(LLMAssistedTaggingPredictionType),
-    )
+    tags = graphene.Field(LLMAssistedTaggingPredictionType)
     geo_areas = graphene.List(
         graphene.NonNull(ProjectGeoAreaType)
     )
@@ -206,8 +204,10 @@ def resolve_geo_areas(root, info, **_):
         return info.context.dl.geo.draft_entry_geo_area.load(root.pk)
 
     # TODO use dataloader
+    @staticmethod
     def resolve_tags(root, info, **_):
-        return LLMAssistedTaggingPredication.objects.filter(draft_entry=root)
+        # return info.context.dl.assisted_tagging.llm_draft_entry_predications.load(root.pk)
+        return LLMAssistedTaggingPredication.objects.get(draft_entry=root)
 
 
 class DraftEntryListType(CustomDjangoListObjectType):

apps/assisted_tagging/serializers.py

@@ -23,7 +23,7 @@ def validate_lead(self, lead):
         if lead.project != self.project:
             raise serializers.ValidationError('Only lead from current project are allowed.')
         af = lead.project.analysis_framework
-        if af is None or not af.assisted_tagging_enabled:
+        if af is None:
             raise serializers.ValidationError('Assisted tagging is disabled for the Framework used by this project.')
         return lead
 

deep/deepl.py

@@ -19,4 +19,4 @@ class DeeplServiceEndpoint():
     ANALYSIS_GEO = f'{DEEPL_SERVER_DOMAIN}/api/v1/geolocation/'
     ASSISTED_TAGGING_ENTRY_PREDICT_ENDPOINT = f'{DEEPL_SERVER_DOMAIN}/api/v1/entry-classification/'
     ENTRY_EXTRACTION_CLASSIFICATION = f'{DEEPL_SERVER_DOMAIN}/api/v1/entry-extraction-classification/'
-    LLM_ASSISTED_TAGGING_ENTRY_PREDICT_ENDPOINT = f'{DEEPL_SERVER_DOMAIN}/api/v1/llm-entry-classification/'
+    LLM_ASSISTED_TAGGING_ENTRY_PREDICT_ENDPOINT = f'{DEEPL_SERVER_DOMAIN}/api/v1/entry-classification-llm/'

schema.graphql

@@ -3238,20 +3238,6 @@ enum AssistedTaggingPredictionDataTypeEnum {
   TAG
 }
 
-type AssistedTaggingPredictionType {
-  id: ID!
-  value: String!
-  prediction: Decimal
-  threshold: Decimal
-  isSelected: Boolean!
-  modelVersion: ID!
-  draftEntry: ID!
-  dataType: AssistedTaggingPredictionDataTypeEnum!
-  dataTypeDisplay: EnumDescription!
-  category: ID
-  tag: ID
-}
-
 type AssistedTaggingQueryType {
   draftEntry(id: ID!): DraftEntryType
   draftEntries(lead: ID, draftEntryTypes: [DraftEntryTypeEnum!], ignoreIds: [ID!], isDiscarded: Boolean, page: Int = 1, pageSize: Int): DraftEntryListType
@@ -3725,8 +3711,6 @@ type DateCountType {
 
 scalar DateTime
 
-scalar Decimal
-
 type DeleteAdminLevel {
   errors: [GenericScalar!]
   ok: Boolean
@@ -3907,7 +3891,7 @@ type DraftEntryType {
   predictionReceivedAt: DateTime
   predictionStatus: DraftEntryPredictionStatusEnum!
   predictionStatusDisplay: EnumDescription!
-  predictionTags: [AssistedTaggingPredictionType!]
+  tags: LLMAssistedTaggingPredictionType
   geoAreas: [ProjectGeoAreaType!]
 }
 
@@ -4628,6 +4612,13 @@ type JwtTokenType {
   expiresIn: String
 }
 
+type LLMAssistedTaggingPredictionType {
+  id: ID!
+  modelVersion: ID!
+  draftEntry: ID!
+  modelTags: GenericScalar
+}
+
 enum LeadAutoEntryExtractionTypeEnum {
   NONE
   STARTED