separate celery tasks, beat and web-core

docker-compose.yml

@@ -19,15 +19,15 @@ services:
                         "
     django:
         extends: django_env
-        command: gunicorn stixify.wsgi:application --bind 0.0.0.0:8004 --reload
+        command: gunicorn stixify.wsgi:application --bind 0.0.0.0:8004 --preload -w 4
         ports:
             - 8004:8004
         depends_on:
             redis:
                 condition: service_started
     celery:
         extends: django_env
-        command: celery -A stixify.worker worker -l INFO
+        command: celery -A stixify.worker worker -l INFO --autoscale 8,2 --pool=prefork
         depends_on:
             - django
             - redis
@@ -39,7 +39,7 @@ services:
         extends: django_env
         command: >
                 bash -c "
-                  celery -A stixify.worker beat -l INFO
+                  celery -A stixify.worker.beat beat -l INFO
                   "
         depends_on:
             redis:

stixify/classifier/tasks.py

@@ -3,7 +3,7 @@
 from typing import Any, List
 
 import numpy as np
-import openai
+from .utils import _openai_client
 import hdbscan
 import joblib
 
@@ -17,38 +17,6 @@ class ClusteringCancelled(Exception):
     pass
 
 
-def _openai_client():
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-    return openai.Client()
-
-
-def compute_embedding_for_document(doc: DocumentEmbedding):
-    """Fetch a document by id, compute embedding using OpenAI small-3."""
-    if not doc.text:
-        raise ValueError("Document text is empty, cannot compute embedding")
-
-    client = _openai_client()
-    try:
-        resp = client.embeddings.create(
-            input=doc.text, model="text-embedding-3-small", dimensions=512
-        )
-        vec = resp.data[0].embedding  # list of floats
-        # store as list of floats; `updated_at` is auto-updated by the model
-        doc.embedding = vec
-        doc.save(update_fields=["embedding", "updated_at"])
-        print(f"Saved embedding for doc {doc.pk}")
-    except Exception as e:
-        print(f"Embedding failed for {doc.pk}: {e}")
-        raise
-
-
-def create_embedding_text(*texts: List[str]) -> str:
-    """Create a single string to embed from multiple text fields."""
-    # simple concat with separator, could be improved with field weighting or truncation
-    texts = [t.strip() for t in texts if t and t.strip()]
-    return " | ".join(texts)
-
-
 def run_clustering(
     min_cluster_size: int = settings.CLASSIFIER_MIN_CLUSTER_SIZE,
     force: bool = False,

stixify/classifier/utils.py

@@ -0,0 +1,43 @@
+import os
+from typing import List
+
+import openai
+
+
+from .models import DocumentEmbedding
+
+
+class ClusteringCancelled(Exception):
+    pass
+
+
+def _openai_client():
+    openai.api_key = os.getenv("OPENAI_API_KEY")
+    return openai.Client()
+
+
+def compute_embedding_for_document(doc: DocumentEmbedding):
+    """Fetch a document by id, compute embedding using OpenAI small-3."""
+    if not doc.text:
+        raise ValueError("Document text is empty, cannot compute embedding")
+
+    client = _openai_client()
+    try:
+        resp = client.embeddings.create(
+            input=doc.text, model="text-embedding-3-small", dimensions=512
+        )
+        vec = resp.data[0].embedding  # list of floats
+        # store as list of floats; `updated_at` is auto-updated by the model
+        doc.embedding = vec
+        doc.save(update_fields=["embedding", "updated_at"])
+        print(f"Saved embedding for doc {doc.pk}")
+    except Exception as e:
+        print(f"Embedding failed for {doc.pk}: {e}")
+        raise
+
+
+def create_embedding_text(*texts: List[str]) -> str:
+    """Create a single string to embed from multiple text fields."""
+    # simple concat with separator, could be improved with field weighting or truncation
+    texts = [t.strip() for t in texts if t and t.strip()]
+    return " | ".join(texts)

stixify/web/models.py

@@ -10,8 +10,6 @@
 from django.core.cache import cache
 import uuid, typing
 from stixify.classifier.models import Cluster, DocumentEmbedding
-from stixify.classifier.tasks import compute_embedding_for_document, create_embedding_text
-import txt2stix, txt2stix.extractions
 from django.core.exceptions import ValidationError
 from datetime import UTC, datetime, timezone
 from django.utils import timezone as dj_timezone
@@ -21,9 +19,11 @@
 from dogesec_commons.stixifier.models import Profile
 from dogesec_commons.identity.models import Identity
 
-from sklearn.metrics.pairwise import cosine_similarity
 from pgvector.django import CosineDistance
-
+from stixify.classifier.utils import (
+    compute_embedding_for_document,
+    create_embedding_text,
+)
 
 if typing.TYPE_CHECKING:
     from .. import settings
@@ -35,8 +35,6 @@ def validate_extractor(types, name):
     pass
 
 
-
-
 class TLP_Levels(models.TextChoices):
     RED = "red"
     AMBER_STRICT = "amber+strict"
@@ -102,7 +100,6 @@ class Meta:
         abstract = True
 
 
-
 def upload_to_func(instance: 'File|FileImage', filename):
     if isinstance(instance, FileImage):
         instance = instance.report
@@ -206,7 +203,7 @@ def process_mode(self):
         return self.mode
 
     def set_txt2stix_data(self, txt2stix_data):
-        from txt2stix.txt2stix import Txt2StixData
+        from txt2stix.utils import Txt2StixData
         if txt2stix_data is None:
             return
 
@@ -233,6 +230,7 @@ def set_txt2stix_data(self, txt2stix_data):
         )
 
     def similar_posts(file, visible_to=None):
+
         if not file.embedding:
             return []
 
@@ -253,15 +251,11 @@ def similar_posts(file, visible_to=None):
                 continue
             if len(results) >= 5:
                 break
-            similarity_score = cosine_similarity(
-                file.embedding.embedding.reshape(1, -1),
-                sfile.embedding.embedding.reshape(1, -1),
-            )[0][0]
             results.append(
                 {
                     "id": sfile.id,
                     "name": sfile.name,  # or get from related file
-                    "score": similarity_score,
+                    "score": 1 - sfile.distance,
                     "tlp_level": sfile.tlp_level,
                     "owner": sfile.identity_id,
                     "added": sfile.created,

stixify/web/serializers.py

@@ -14,7 +14,6 @@
 import file2txt.parsers.core as f2t_core
 from rest_framework.exceptions import ValidationError
 from django.utils.translation import gettext_lazy
-from stixify.worker import pdf_converter
 from django.core.files.base import ContentFile
 import tempfile
 from pathlib import Path
@@ -156,6 +155,8 @@ def create(self, validated_data):
 
         # Handle mhtml-pdf conversion
         if mode == 'mhtml-pdf':
+            from stixify.worker import pdf_converter
+
             # Save uploaded file to temporary location
             with tempfile.NamedTemporaryFile(delete=True, suffix='.mhtml') as temp_file:
                 for chunk in uploaded_file.chunks():

stixify/worker/beat.py

@@ -0,0 +1,14 @@
+from datetime import timedelta
+
+from celery import Celery
+
+
+app = Celery("stixify-beat")
+app.config_from_object("os:environ", namespace="CELERY")
+
+app.conf.beat_schedule = {
+    "auto_refresh_statistics_data": {
+        "task": "stixify.worker.tasks.auto_refresh_statistics_data",
+        "schedule": timedelta(minutes=10),
+    }
+}

stixify/worker/celery.py

@@ -1,4 +1,3 @@
-from datetime import timedelta
 import os
 from celery import Celery
 # Set the default Django settings module for the 'celery' program.
@@ -10,12 +9,11 @@
 
 app.config_from_object('os:environ', namespace='CELERY')
 
+app.conf.imports = (
+    "stixify.worker.process_post",
+    "stixify.classifier.tasks",
+    "stixify.worker.tasks",
+)
+
 # Load task modules from all registered Django apps.
 app.autodiscover_tasks()
-
-app.conf.beat_schedule = {
-    "auto_refresh_statistics_data": {
-        "task": "stixify.worker.tasks.auto_refresh_statistics_data",
-        "schedule": timedelta(minutes=10),
-    }
-}