direct message implementation

Author: Dishant1804

backend/apps/slack/MANIFEST.yaml

@@ -133,6 +133,7 @@ settings:
       - app_mention
       - member_joined_channel
       - message.channels
+      - message.im
       - team_join
   interactivity:
     is_enabled: true

backend/apps/slack/admin/__init__.py

@@ -1,5 +1,6 @@
 """Slack app admin."""
 
+from .chat import ChatAdmin
 from .conversation import ConversationAdmin
 from .event import EventAdmin
 from .member import MemberAdmin

backend/apps/slack/admin/chat.py

@@ -0,0 +1,16 @@
+"""Chat admin configuration."""
+
+from django.contrib import admin
+
+from apps.slack.models.chat import Chat
+
+
+class ChatAdmin(admin.ModelAdmin):
+    """Admin for Chat model."""
+
+    list_display = ("user", "workspace", "created_at")
+    list_filter = ("user", "workspace")
+    search_fields = ("user__username", "workspace__name")
+
+
+admin.site.register(Chat, ChatAdmin)

backend/apps/slack/common/handlers/ai.py

@@ -6,6 +6,7 @@
 
 from apps.ai.agent.tools.rag.rag_tool import RagTool
 from apps.slack.blocks import markdown
+from apps.slack.models import Chat, Member, Workspace
 
 logger = logging.getLogger(__name__)
 
@@ -46,6 +47,59 @@ def process_ai_query(query: str) -> str | None:
     return rag_tool.query(question=query)
 
 
+def get_dm_blocks(query: str, user_id: str, workspace_id: str) -> list[dict]:
+    """Get AI response blocks for DM with conversation context.
+
+    Args:
+        query (str): The user's question.
+        user_id (str): Slack user ID.
+        workspace_id (str): Slack workspace ID.
+
+    Returns:
+        list: A list of Slack blocks representing the AI response.
+
+    """
+    ai_response = process_dm_ai_query(query.strip(), user_id, workspace_id)
+
+    if ai_response:
+        return [markdown(ai_response)]
+    return get_error_blocks()
+
+
+def process_dm_ai_query(query: str, user_id: str, workspace_id: str) -> str | None:
+    """Process the AI query with DM conversation context.
+
+    Args:
+        query (str): The user's question.
+        user_id (str): Slack user ID.
+        workspace_id (str): Slack workspace ID.
+
+    Returns:
+        str | None: The AI response or None if error occurred.
+
+    """
+    user = Member.objects.get(slack_user_id=user_id)
+    workspace = Workspace.objects.get(slack_workspace_id=workspace_id)
+
+    chat = Chat.update_data(user, workspace)
+    context = chat.get_context(limit_exchanges=20)
+
+    rag_tool = RagTool(
+        chat_model="gpt-4o",
+        embedding_model="text-embedding-3-small",
+    )
+
+    if context:
+        enhanced_query = f"Conversation context:\n{context}\n\nCurrent question: {query}"
+    else:
+        enhanced_query = query
+
+    response = rag_tool.query(question=enhanced_query)
+    chat.add_to_context(query, response)
+
+    return response
+
+
 def get_error_blocks() -> list[dict]:
     """Get error response blocks.
 

backend/apps/slack/events/message_posted.py

@@ -1,21 +1,22 @@
-"""Slack message event template."""
+"""Slack message event handler for OWASP NestBot."""
 
 import logging
 from datetime import timedelta
 
 import django_rq
 
 from apps.ai.common.constants import QUEUE_RESPONSE_TIME_MINUTES
+from apps.slack.common.handlers.ai import get_dm_blocks
 from apps.slack.common.question_detector import QuestionDetector
 from apps.slack.events.event import EventBase
-from apps.slack.models import Conversation, Member, Message
+from apps.slack.models import Conversation, Member, Message, Workspace
 from apps.slack.services.message_auto_reply import generate_ai_reply_if_unanswered
 
 logger = logging.getLogger(__name__)
 
 
 class MessagePosted(EventBase):
-    """Handles new messages posted in channels."""
+    """Handles new messages posted in channels or direct messages."""
 
     event_type = "message"
 
@@ -24,25 +25,30 @@ def __init__(self):
         self.question_detector = QuestionDetector()
 
     def handle_event(self, event, client):
-        """Handle an incoming message event."""
+        """Handle incoming Slack message events."""
         if event.get("subtype") or event.get("bot_id"):
-            logger.info("Ignored message due to subtype, bot_id, or thread_ts.")
+            logger.info("Ignored message due to subtype or bot_id.")
+            return
+
+        channel_id = event.get("channel")
+        user_id = event.get("user")
+        text = event.get("text", "")
+        channel_type = event.get("channel_type")
+
+        if channel_type == "im":
+            self.handle_dm(event, client, channel_id, user_id, text)
             return
 
         if event.get("thread_ts"):
             try:
                 Message.objects.filter(
                     slack_message_id=event.get("thread_ts"),
-                    conversation__slack_channel_id=event.get("channel"),
+                    conversation__slack_channel_id=channel_id,
                 ).update(has_replies=True)
             except Message.DoesNotExist:
                 logger.warning("Thread message not found.")
             return
 
-        channel_id = event.get("channel")
-        user_id = event.get("user")
-        text = event.get("text", "")
-
         try:
             conversation = Conversation.objects.get(
                 slack_channel_id=channel_id,
@@ -71,3 +77,50 @@ def handle_event(self, event, client):
             generate_ai_reply_if_unanswered,
             message.id,
         )
+
+    def handle_dm(self, event, client, channel_id, user_id, text):
+        """Handle direct messages with NestBot (DMs)."""
+        workspace_id = event.get("team")
+
+        if not workspace_id:
+            try:
+                channel_info = client.conversations_info(channel=channel_id)
+                workspace_id = channel_info["channel"]["team"]
+            except Exception:
+                logger.exception("Failed to fetch workspace ID for DM.")
+                return
+
+        try:
+            Member.objects.get(slack_user_id=user_id, workspace__slack_workspace_id=workspace_id)
+        except Member.DoesNotExist:
+            try:
+                user_info = client.users_info(user=user_id)
+                workspace = Workspace.objects.get(slack_workspace_id=workspace_id)
+                Member.update_data(user_info["user"], workspace, save=True)
+                logger.info("Created new member for DM")
+            except Exception:
+                logger.exception("Failed to create member for DM.")
+                return
+
+        thread_ts = event.get("thread_ts")
+
+        try:
+            response_blocks = get_dm_blocks(text, user_id, workspace_id)
+            if response_blocks:
+                client.chat_postMessage(
+                    channel=channel_id,
+                    blocks=response_blocks,
+                    text=text,
+                    thread_ts=thread_ts,
+                )
+
+        except Exception:
+            logger.exception("Error processing DM")
+            client.chat_postMessage(
+                channel=channel_id,
+                text=(
+                    "I'm sorry, I'm having trouble processing your message right now. "
+                    "Please try again later."
+                ),
+                thread_ts=thread_ts,
+            )

backend/apps/slack/migrations/0020_chat.py

@@ -0,0 +1,56 @@
+# Generated by Django 5.2.6 on 2025-09-26 19:24
+
+import django.db.models.deletion
+import django.utils.timezone
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("slack", "0019_conversation_is_nest_bot_assistant_enabled"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="Chat",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                ("nest_created_at", models.DateTimeField(auto_now_add=True)),
+                ("nest_updated_at", models.DateTimeField(auto_now=True)),
+                ("context", models.TextField(blank=True)),
+                (
+                    "created_at",
+                    models.DateTimeField(
+                        default=django.utils.timezone.now, verbose_name="Created at"
+                    ),
+                ),
+                ("is_active", models.BooleanField(default=True)),
+                (
+                    "user",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="chats",
+                        to="slack.member",
+                    ),
+                ),
+                (
+                    "workspace",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="chats",
+                        to="slack.workspace",
+                    ),
+                ),
+            ],
+            options={
+                "db_table": "slack_chat",
+                "ordering": ["-created_at"],
+                "unique_together": {("user", "workspace")},
+            },
+        ),
+    ]

backend/apps/slack/models/__init__.py

@@ -1,3 +1,4 @@
+from .chat import Chat
 from .conversation import Conversation
 from .event import Event
 from .member import Member

backend/apps/slack/models/chat.py

@@ -0,0 +1,90 @@
+"""Chat model for storing conversation context."""
+
+from django.db import models
+from django.utils import timezone
+
+from apps.common.models import TimestampedModel
+from apps.slack.models.member import Member
+from apps.slack.models.workspace import Workspace
+
+
+class Chat(TimestampedModel):
+    """Store chat conversation context for DMs."""
+
+    context = models.TextField(blank=True)
+    created_at = models.DateTimeField(verbose_name="Created at", default=timezone.now)
+    is_active = models.BooleanField(default=True)
+    user = models.ForeignKey(Member, on_delete=models.CASCADE, related_name="chats")
+    workspace = models.ForeignKey(Workspace, on_delete=models.CASCADE, related_name="chats")
+
+    class Meta:
+        db_table = "slack_chat"
+        unique_together = [["user", "workspace"]]
+        ordering = ["-created_at"]
+
+    def __str__(self):
+        """Return a concise, human-readable identifier for this chat."""
+        return f"Chat with {self.user.real_name or self.user.username} in {self.workspace.name}"
+
+    @staticmethod
+    def update_data(user: Member, workspace: Workspace, *, save: bool = True) -> "Chat":
+        """Update or create chat data for a user in a workspace.
+
+        Args:
+            user: Member instance to associate with the chat.
+            workspace: Workspace instance to associate with the chat.
+            save: Whether to save the chat to the database.
+
+        Returns:
+            Updated or created Chat instance.
+
+        """
+        try:
+            chat = Chat.objects.get(user=user, workspace=workspace)
+        except Chat.DoesNotExist:
+            chat = Chat(user=user, workspace=workspace, is_active=True)
+
+        if save:
+            chat.save()
+
+        return chat
+
+    def add_to_context(self, user_message: str, bot_response: str | None = None) -> None:
+        """Add messages to the conversation context.
+
+        Args:
+            user_message: The user's message to add to context.
+            bot_response: The bot's response to add to context.
+
+        """
+        if not self.context:
+            self.context = ""
+
+        self.context += f"User: {user_message}\n"
+
+        if bot_response:
+            self.context += f"Bot: {bot_response}\n"
+
+        self.save(update_fields=["context"])
+
+    def get_context(self, limit_exchanges: int | None = None) -> str:
+        """Get the conversation context.
+
+        Args:
+            limit_exchanges: Optional limit on number of exchanges to return.
+
+        Returns:
+            The conversation context, potentially limited to recent exchanges.
+
+        """
+        if not self.context:
+            return ""
+
+        if limit_exchanges is None:
+            return self.context
+
+        lines = self.context.strip().split("\n")
+        if len(lines) <= limit_exchanges * 2:
+            return self.context
+
+        return "\n".join(lines[-(limit_exchanges * 2) :])