Message Tracking Improvements (#191)

* message init

* log message summary

* deletion by

* improvements logging summaries

* primary key id

* cascade org delete

* PR comments

Author: LennartSchmidtKern

cognition_objects/admin_query_message_summary.py

@@ -0,0 +1,51 @@
+from datetime import date
+from submodules.model.business_objects import general
+from submodules.model.models import AdminQueryMessageSummary
+from ..session import session
+
+
+def get_admin_query_message_summary(
+    org_id: str,
+    project_id: str,
+    day: date,
+) -> AdminQueryMessageSummary | None:
+    return (
+        session.query(AdminQueryMessageSummary)
+        .filter(
+            AdminQueryMessageSummary.organization_id == org_id,
+            AdminQueryMessageSummary.project_id == project_id,
+            AdminQueryMessageSummary.day == day,
+        )
+        .first()
+    )
+
+
+def log_admin_query_message_summary(
+    org_id: str,
+    project_id: str,
+    day: date,
+    counters: dict,  # {"total_messages": 5, "messages_via_api": 2, ....},
+    with_commit: bool = True,
+):
+
+    message_summary = get_admin_query_message_summary(org_id, project_id, day)
+
+    if message_summary:
+        for col, value in counters.items():
+            if not hasattr(message_summary, col):
+                continue
+            current_value = getattr(message_summary, col, 0) or 0
+            setattr(message_summary, col, current_value + value)
+    else:
+        message_summary = AdminQueryMessageSummary(
+            organization_id=org_id, project_id=project_id, day=day
+        )
+        for col, value in counters.items():
+            if hasattr(message_summary, col):
+                setattr(message_summary, col, value)
+        general.add(message_summary, with_commit)
+        return message_summary
+
+    if with_commit:
+        general.flush_or_commit(True)
+    return message_summary

cognition_objects/conversation.py

@@ -364,13 +364,15 @@ def add_message(
     conversation_id: str,
     user_id: str,
     question: str,
+    initiated_via: str,
     with_commit: bool = True,
 ) -> CognitionConversation:
     message_entity = message.create(
         conversation_id=conversation_id,
         project_id=project_id,
         user_id=user_id,
         question=question,
+        initiated_via=initiated_via,
         with_commit=with_commit,
     )
 

cognition_objects/message.py

@@ -273,6 +273,7 @@ def create(
     project_id: str,
     user_id: str,
     question: str,
+    initiated_via: str,
     with_commit: bool = True,
     created_at: Optional[datetime] = None,
 ) -> CognitionMessage:
@@ -288,6 +289,7 @@ def create(
         question=question,
         facts=[],
         version_id=version_id,
+        initiated_via=initiated_via,
     )
 
     general.add(message, with_commit)
@@ -532,3 +534,13 @@ def update_version_id_for_messages(
     )
     if with_commit:
         general.commit()
+
+
+def get_count_by_project_id(project_id: str) -> int:
+    return (
+        session.query(CognitionMessage)
+        .filter(
+            CognitionMessage.project_id == project_id,
+        )
+        .count()
+    )

enums.py

@@ -175,6 +175,7 @@ class Tablenames(Enum):
         "conversation_tag_association"  # association between conversation and tags
     )
     SUMS_TABLE = "sums_table"
+    ADMIN_QUERY_MESSAGE_SUMMARY = "admin_query_message_summary"
 
     def snake_case_to_pascal_case(self):
         # the type name (written in PascalCase) of a table is needed to create backrefs
@@ -955,3 +956,9 @@ class SharepointPropertySyncState(Enum):
     RUNNING = "RUNNING"
     COMPLETED = "COMPLETED"
     FAILED = "FAILED"
+
+
+class MessageInitiationType(Enum):
+    UI = "UI"
+    API = "API"
+    MACRO = "MACRO"

listeners.py

@@ -0,0 +1,31 @@
+from sqlalchemy import event
+from datetime import datetime, timezone
+from submodules.model.models import CognitionConversation, CognitionMessage
+from submodules.model.session_wrapper import with_session
+import traceback
+
+from src.controller.admin_query_message_summary import (
+    manager as admin_query_message_summary_manager,
+)
+
+
+@event.listens_for(CognitionConversation, "after_insert")
+@with_session()
+def after_insert(mapper, connection, conversation_entity: CognitionConversation):
+    try:
+        admin_query_message_summary_manager.log_conversation_summary(
+            conversation_entity, 1
+        )
+    except Exception:
+        print("Error in after_insert listener of CognitionConversation", flush=True)
+        print(traceback.format_exc(), flush=True)
+
+
+@event.listens_for(CognitionMessage, "after_insert")
+@with_session()
+def after_insert(mapper, connection, message_entity: CognitionMessage):
+    try:
+        admin_query_message_summary_manager.log_message_summary(message_entity)
+    except Exception:
+        print("Error in after_insert listener of CognitionMessage", flush=True)
+        print(traceback.format_exc(), flush=True)

models.py

@@ -1286,6 +1286,7 @@ class CognitionMessage(Base):
         ),
         index=True,
     )
+    initiated_via = Column(String)  # of type enums.MessageInitiationType.*.value
 
 
 class CognitionPipelineLogs(Base):
@@ -2453,3 +2454,37 @@ class SumsTable(Base):
     sum_key = Column(String, index=True)  # e.g. enums.AdminQueries
     created_at = Column(DateTime, default=sql.func.now())
     data = Column(JSON)
+
+
+class AdminQueryMessageSummary(Base):
+    __tablename__ = Tablenames.ADMIN_QUERY_MESSAGE_SUMMARY.value
+    __table_args__ = (
+        UniqueConstraint(
+            "organization_id",
+            "project_id",
+            "day",
+            name="unique_admin_query_msg_activity_summary",
+        ),
+        {"schema": "cognition"},
+    )
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    day = Column(Date, nullable=False)
+    organization_id = Column(
+        UUID(as_uuid=True),
+        ForeignKey(f"{Tablenames.ORGANIZATION.value}.id", ondelete="CASCADE"),
+    )
+    project_id = Column(
+        UUID(as_uuid=True),
+        ForeignKey(f"cognition.{Tablenames.PROJECT.value}.id", ondelete="SET NULL"),
+    )
+
+    total_conversations = Column(Integer, default=0)
+    total_messages = Column(Integer, default=0)
+    messages_via_api = Column(Integer, default=0)
+    messages_via_ui = Column(Integer, default=0)
+    messages_via_macro = Column(Integer, default=0)
+    confidential_messages = Column(Integer, default=0)
+    kern_user_messages = Column(Integer, default=0)
+    deleted_messages_by_user = Column(Integer, default=0)
+    deleted_messages_by_system = Column(Integer, default=0)
+    incognito_messages = Column(Integer, default=0)