From 1d36ae1fcbf77459b12998800fc532d0681f9240 Mon Sep 17 00:00:00 2001
From: gecBurton <g.e.c.burton@gmail.com>
Date: Thu, 17 Oct 2024 09:32:06 +0100
Subject: [PATCH] added cmd to back populate logs

---
 .../redbox_app/redbox_core/consumers.py       | 21 ++-----------------
 .../commands/back_populate_elastic_logs.py    | 11 ++++++++++
 django_app/redbox_app/redbox_core/models.py   | 19 ++++++++++++++++-
 3 files changed, 31 insertions(+), 20 deletions(-)
 create mode 100644 django_app/redbox_app/redbox_core/management/commands/back_populate_elastic_logs.py

diff --git a/django_app/redbox_app/redbox_core/consumers.py b/django_app/redbox_app/redbox_core/consumers.py
index df38e3378..be6983bab 100644
--- a/django_app/redbox_app/redbox_core/consumers.py
+++ b/django_app/redbox_app/redbox_core/consumers.py
@@ -3,7 +3,7 @@
 from asyncio import CancelledError
 from collections.abc import Mapping, Sequence
 from typing import Any, ClassVar
-from uuid import UUID, uuid4
+from uuid import UUID
 
 from channels.db import database_sync_to_async
 from channels.generic.websocket import AsyncWebsocketConsumer
@@ -222,25 +222,8 @@ def save_message(
                     model_name=model,
                     token_count=token_count,
                 )
-            token_count = sum(metadata.output_tokens.values())
-        else:
-            token_count = 0
-
-        elastic_log_msg = {
-            "@timestamp": chat_message.created_at.isotime(),
-            "id": chat_message.id,
-            "chat_id": chat_message.chat.id,
-            "user_id": chat_message.chat.user.id,
-            "text": chat_message.text,
-            "route": chat_message.route,
-            "role": "ai",
-            "token_count": token_count,
-            "rating": chat_message.rating,
-            "rating_text": chat_message.rating_text,
-            "rating_chips": chat_message.rating_chips,
-        }
-        elasticsearch_client.create(env.elastic_chat_mesage_index, uuid4(), elastic_log_msg)
 
+        chat_message.log()
         return chat_message
 
     @staticmethod
diff --git a/django_app/redbox_app/redbox_core/management/commands/back_populate_elastic_logs.py b/django_app/redbox_app/redbox_core/management/commands/back_populate_elastic_logs.py
new file mode 100644
index 000000000..99a2534f5
--- /dev/null
+++ b/django_app/redbox_app/redbox_core/management/commands/back_populate_elastic_logs.py
@@ -0,0 +1,11 @@
+from django.core.management import BaseCommand
+
+from redbox_app.redbox_core.models import ChatMessage
+
+
+class Command(BaseCommand):
+    help = """This is a one-off command to back populate elastic logs."""
+
+    def handle(self, *args, **kwargs):  # noqa:ARG002
+        for chat_message in ChatMessage.objects.all():
+            chat_message.log()
diff --git a/django_app/redbox_app/redbox_core/models.py b/django_app/redbox_app/redbox_core/models.py
index 8c572bff7..ac454b849 100644
--- a/django_app/redbox_app/redbox_core/models.py
+++ b/django_app/redbox_app/redbox_core/models.py
@@ -13,7 +13,7 @@
 from django.contrib.postgres.fields import ArrayField
 from django.core import validators
 from django.db import models
-from django.db.models import Max, Min, Prefetch, UniqueConstraint
+from django.db.models import Max, Min, Prefetch, Sum, UniqueConstraint
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
 from django_use_email_as_username.models import BaseUser, BaseUserManager
@@ -786,6 +786,23 @@ def get_messages_ordered_by_citation_priority(cls, chat_id: uuid.UUID) -> Sequen
             )
         )
 
+    def log(self):
+        token_count = self.chatmessagetokenuse_set.annotate(token_count=Sum("token_count"))
+        elastic_log_msg = {
+            "@timestamp": self.created_at.isotime(),
+            "id": self.id,
+            "chat_id": self.chat.id,
+            "user_id": self.chat.user.id,
+            "text": self.text,
+            "route": self.route,
+            "role": "ai",
+            "token_count": token_count["token_count"],
+            "rating": self.rating,
+            "rating_text": self.rating_text,
+            "rating_chips": self.rating_chips,
+        }
+        es_client.create(env.elastic_chat_mesage_index, uuid.uuid4(), elastic_log_msg)
+
 
 class ChatMessageTokenUse(UUIDPrimaryKeyBase, TimeStampedModel):
     class UseTypeEnum(models.TextChoices):