Save zeroshot logs in elasticsearch

contrib/gen_env.py

@@ -126,6 +126,7 @@ def generate_env():
         "ZEROSHOT_MODEL_BACKEND": "",
         "ZEROSHOT_BEDROCK_AWS_SECRET": "",
         "ZEROSHOT_BEDROCK_AWS_REGION": "",
+        "ZEROSHOT_LOG_SAVE_ELASTIC": False,
     }
 
     with open(env_path, "w") as configfile:

docker-compose.yml

@@ -58,6 +58,28 @@ services:
     environment:
       - RABBITMQ_DEFAULT_USER=guest
       - RABBITMQ_DEFAULT_PASS=guest
+
+  elasticsearch:
+    image: docker.elastic.co/elasticsearch/elasticsearch:7.17.16
+    ports:
+      - "9200:9200"
+      - "9300:9300"
+    environment:
+      - http.port=9200
+      - discovery.type=single-node
+      - bootstrap.memory_lock=true
+      - ES_HEAP_SIZE=1g
+      - ES_JAVA_OPTS=-Xms1g -Xmx1g
+      - http.cors.enabled=true
+      - http.cors.allow-origin="*"
+      - http.cors.allow-headers=X-Requested-With,X-Auth-Token,Content-Type,Content-Length,Authorization
+      - http.cors.allow-credentials=true
+      - network.publish_host=localhost
+      - xpack.security.enabled=false
+    networks:
+      - default
+      - nexus
+
 networks:
   nexus:
     external: true

nexus/settings.py

@@ -67,6 +67,7 @@
     'rest_framework',
     'drf_yasg',
     'elasticapm.contrib.django',
+    'django_elasticsearch_dsl',
     # apps
     'nexus.users',
     'nexus.db',
@@ -467,3 +468,15 @@
 ZEROSHOT_BEDROCK_AWS_SECRET = env.str("ZEROSHOT_BEDROCK_AWS_SECRET")
 ZEROSHOT_BEDROCK_AWS_REGION = env.str("ZEROSHOT_BEDROCK_AWS_REGION")
 ZEROSHOT_BEDROCK_MODEL_ID = env.str("ZEROSHOT_BEDROCK_MODEL_ID")
+ZEROSHOT_LOG_SAVE_ELASTIC = env.bool("ZEROSHOT_LOG_SAVE_ELASTIC", True)
+
+
+ELASTICSEARCH_HOST = env.str("ELASTICSEARCH_HOST", "http://localhost:9200")
+ELASTICSEARCH_NUMBER_OF_SHARDS = env.int("ELASTICSEARCH_NUMBER_OF_SHARDS", 1)
+ELASTICSEARCH_NUMBER_OF_REPLICAS = env.int("ELASTICSEARCH_NUMBER_OF_REPLICAS", 0)
+ELASTICSEARCH_DSL = {
+    'default': {
+        'hosts': ELASTICSEARCH_HOST,
+        'http_auth': ('username', 'password')
+    }
+}

nexus/task_managers/tasks.py

@@ -22,6 +22,8 @@
 from nexus.usecases.intelligences.update import UpdateContentBaseFileUseCase
 from nexus.usecases.intelligences.get_by_uuid import get_by_contentbase_uuid
 from nexus.usecases.logs.delete import DeleteLogUsecase
+from nexus.usecases.logs.entities import ZeroshotDTO
+from nexus.usecases.logs.create import CreateZeroshotLogsUseCase
 
 from nexus.trulens import wenigpt_evaluation, tru_recorder
 
@@ -209,3 +211,9 @@ def update_healthcheck():
 def update_classification_healthcheck():
     classification_notify = ClassificationHealthCheck()
     classification_notify.check_service_health()
+
+
+@app.task
+def create_zeroshot_logs(zeroshot_dict: Dict):
+    zeroshot_dto = ZeroshotDTO(**zeroshot_dict)
+    CreateZeroshotLogsUseCase(zeroshot_dto).create()

nexus/usecases/logs/create.py

@@ -2,6 +2,8 @@
     Message,
     MessageLog,
 )
+from nexus.zeroshot.models import ZeroshotLogs
+from nexus.usecases.logs.entities import ZeroshotDTO
 
 
 class CreateLogUsecase:
@@ -42,3 +44,21 @@ def update_log_field(self, **kwargs):
             setattr(log, key, kwargs.get(key))
 
         log.save()
+
+
+class CreateZeroshotLogsUseCase:
+    def __init__(
+            self,
+            zeroshot_dto: ZeroshotDTO,
+    ) -> None:
+        self.zeroshot_dto = zeroshot_dto
+
+    def create(self) -> ZeroshotLogs:
+        return ZeroshotLogs.objects.create(
+            text=self.zeroshot_dto.text,
+            classification=self.zeroshot_dto.classification,
+            other=self.zeroshot_dto.other,
+            options=self.zeroshot_dto.options,
+            nlp_log=self.zeroshot_dto.nlp_log,
+            language=self.zeroshot_dto.language,
+        )

nexus/usecases/logs/entities.py

@@ -21,3 +21,13 @@ def dict(self) -> Dict:
             },
             "instructions": self.instructions
         }
+
+
+@dataclass
+class ZeroshotDTO:
+    text: str
+    classification: str
+    other: bool
+    options: Dict
+    nlp_log: str
+    language: str

nexus/usecases/logs/tests/test_create.py

@@ -0,0 +1,30 @@
+from django.test import TestCase
+
+from nexus.zeroshot.models import ZeroshotLogs
+from nexus.usecases.logs.create import CreateZeroshotLogsUseCase
+from nexus.usecases.logs.entities import ZeroshotDTO
+
+
+class CreateZeroshotLogsTestCase(TestCase):
+    def test_create_log(self):
+
+        text = 'oi'
+        classification = "None"
+        other = True
+        options = [
+            {'id': 1, 'class': 'Lorem Ipsum', 'context': 'Dolor sit Amet'},
+            {'id': 2, 'class': 'consectetur adipiscing elit', 'context': 'Vestibulum pharetra erat et nisl pretium viverra'}
+        ]
+        nlp_log = '{"output": {"other": true, "classification": "None"}}'
+        language = 'por'
+
+        dto = ZeroshotDTO(
+            text=text,
+            classification=classification,
+            other=other, options=options,
+            nlp_log=nlp_log,
+            language=language
+        )
+        log = CreateZeroshotLogsUseCase(dto).create()
+
+        self.assertIsInstance(log, ZeroshotLogs)

nexus/zeroshot/api/views.py

@@ -7,7 +7,8 @@
 
 from nexus.zeroshot.client import InvokeModel
 from nexus.zeroshot.api.permissions import ZeroshotTokenPermission
-from nexus.zeroshot.models import ZeroshotLogs
+from nexus.usecases.logs.entities import ZeroshotDTO
+from nexus.task_managers.tasks import create_zeroshot_logs
 
 
 logger = logging.getLogger(__name__)
@@ -23,16 +24,15 @@ def post(self, request):
         try:
             invoke_model = InvokeModel(data)
             response = invoke_model.invoke()
-
-            ZeroshotLogs.objects.create(
+            zeroshot_dto = ZeroshotDTO(
                 text=data.get("text"),
                 classification=response["output"].get("classification"),
                 other=response["output"].get("other", False),
                 options=data.get("options"),
                 nlp_log=str(json.dumps(response)),
                 language=data.get("language")
             )
-
+            create_zeroshot_logs.delay(zeroshot_dto.__dict__)
             return Response(status=200, data=response if response.get("output") else {"error": response})
         except Exception as error:
             traceback.print_exc()

nexus/zeroshot/documents.py

@@ -0,0 +1,34 @@
+import json
+
+from django.conf import settings
+from django_elasticsearch_dsl import Document, fields
+from django_elasticsearch_dsl.registries import registry
+
+from nexus.zeroshot.models import ZeroshotLogs
+
+
+if settings.ZEROSHOT_LOG_SAVE_ELASTIC:
+    @registry.register_document
+    class ZeroshotLogsDocument(Document):
+        class Index:
+            name = "zeroshot_logs_nexus"
+            settings = {
+                'number_of_shards': settings.ELASTICSEARCH_NUMBER_OF_SHARDS,
+                'number_of_replicas': settings.ELASTICSEARCH_NUMBER_OF_REPLICAS
+            }
+
+        class Django:
+            model = ZeroshotLogs
+            fields = [
+                "text",
+                "classification",
+                "other",
+                "nlp_log",
+                "created_at",
+                "language"
+            ]
+
+        options = fields.TextField()
+
+        def prepare_options(self, instance):
+            return json.dumps(instance.options)

pyproject.toml

@@ -37,6 +37,7 @@ uvicorn = "^0.29.0"
 elastic-apm = "^6.22.2"
 django-prometheus = "^2.3.1"
 prometheus-client = "^0.20.0"
+django-elasticsearch-dsl = ">7.0.0,<8.0.0"
 
 
 [tool.poetry.group.dev.dependencies]

router/clients/zeroshot.py

@@ -5,6 +5,8 @@
 
 import requests
 from nexus.zeroshot.client import InvokeModel
+from nexus.usecases.logs.entities import ZeroshotDTO
+from nexus.task_managers.tasks import create_zeroshot_logs
 
 
 class NexusZeroshotClient:
@@ -21,6 +23,15 @@ def fast_predict(self, message: str, actions: List[Dict], language: str = "por")
         }
         zeroshot = InvokeModel(zeroshot_data)
         response = zeroshot.invoke()
+        zeroshot_dto = ZeroshotDTO(
+            text=zeroshot_data.get("text"),
+            classification=response["output"].get("classification"),
+            other=response["output"].get("other", False),
+            options=zeroshot_data.get("options"),
+            nlp_log=str(json.dumps(response)),
+            language=zeroshot_data.get("language")
+        )
+        create_zeroshot_logs.delay(zeroshot_dto.__dict__)
         return response.get("output")