From 10cda990064ec6b853c937b7e29f2ef35cfb7afd Mon Sep 17 00:00:00 2001 From: Dishant1804 Date: Tue, 7 Oct 2025 02:26:10 +0530 Subject: [PATCH 1/3] direct message implementation --- backend/apps/slack/MANIFEST.yaml | 1 + backend/apps/slack/admin/__init__.py | 1 + backend/apps/slack/admin/chat.py | 16 ++++ backend/apps/slack/common/handlers/ai.py | 54 +++++++++++ backend/apps/slack/events/message_posted.py | 73 ++++++++++++--- backend/apps/slack/migrations/0020_chat.py | 56 ++++++++++++ backend/apps/slack/models/__init__.py | 1 + backend/apps/slack/models/chat.py | 90 +++++++++++++++++++ .../apps/slack/events/message_posted_test.py | 29 +++--- 9 files changed, 299 insertions(+), 22 deletions(-) create mode 100644 backend/apps/slack/admin/chat.py create mode 100644 backend/apps/slack/migrations/0020_chat.py create mode 100644 backend/apps/slack/models/chat.py diff --git a/backend/apps/slack/MANIFEST.yaml b/backend/apps/slack/MANIFEST.yaml index d6617e94b3..87c08a2a15 100644 --- a/backend/apps/slack/MANIFEST.yaml +++ b/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 diff --git a/backend/apps/slack/admin/__init__.py b/backend/apps/slack/admin/__init__.py index b3193ffb04..ef678fbb00 100644 --- a/backend/apps/slack/admin/__init__.py +++ b/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 diff --git a/backend/apps/slack/admin/chat.py b/backend/apps/slack/admin/chat.py new file mode 100644 index 0000000000..d51a536a6e --- /dev/null +++ b/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) diff --git a/backend/apps/slack/common/handlers/ai.py b/backend/apps/slack/common/handlers/ai.py index ef0452e7b8..c39da35248 100644 --- a/backend/apps/slack/common/handlers/ai.py +++ b/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. diff --git a/backend/apps/slack/events/message_posted.py b/backend/apps/slack/events/message_posted.py index 5b38c3077b..75aadb3d15 100644 --- a/backend/apps/slack/events/message_posted.py +++ b/backend/apps/slack/events/message_posted.py @@ -1,4 +1,4 @@ -"""Slack message event template.""" +"""Slack message event handler for OWASP NestBot.""" import logging from datetime import timedelta @@ -6,16 +6,17 @@ 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, + ) diff --git a/backend/apps/slack/migrations/0020_chat.py b/backend/apps/slack/migrations/0020_chat.py new file mode 100644 index 0000000000..06f48062dc --- /dev/null +++ b/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")}, + }, + ), + ] diff --git a/backend/apps/slack/models/__init__.py b/backend/apps/slack/models/__init__.py index 3bbe0878de..10a0a0473f 100644 --- a/backend/apps/slack/models/__init__.py +++ b/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 diff --git a/backend/apps/slack/models/chat.py b/backend/apps/slack/models/chat.py new file mode 100644 index 0000000000..a2ce4f5901 --- /dev/null +++ b/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) :]) diff --git a/backend/tests/apps/slack/events/message_posted_test.py b/backend/tests/apps/slack/events/message_posted_test.py index 2f431416f9..2df5bf5e0a 100644 --- a/backend/tests/apps/slack/events/message_posted_test.py +++ b/backend/tests/apps/slack/events/message_posted_test.py @@ -88,11 +88,17 @@ def test_handle_event_ignores_thread_messages(self, message_handler): client = Mock() with patch("apps.slack.events.message_posted.Message") as mock_message: - mock_message.DoesNotExist = Exception - mock_message.objects.get.side_effect = Exception("Message not found") + mock_filter = Mock() + mock_message.objects.filter.return_value = mock_filter message_handler.handle_event(event, client) + mock_message.objects.filter.assert_called_once_with( + slack_message_id=event.get("thread_ts"), + conversation__slack_channel_id=event.get("channel"), + ) + mock_filter.update.assert_called_once_with(has_replies=True) + client.chat_postMessage.assert_not_called() def test_handle_event_conversation_not_found(self, message_handler): @@ -260,6 +266,7 @@ def test_handle_event_member_not_found(self, message_handler, conversation_mock) patch("apps.slack.events.message_posted.Conversation") as mock_conversation, patch("apps.slack.events.message_posted.Member") as mock_member, patch("apps.slack.events.message_posted.Message") as mock_message_model, + patch("apps.slack.events.message_posted.django_rq") as mock_django_rq, ): mock_conversation.objects.get.return_value = conversation_mock @@ -273,19 +280,17 @@ def test_handle_event_member_not_found(self, message_handler, conversation_mock) mock_message.id = 1 mock_message_model.update_data.return_value = mock_message - with ( - patch.object( - message_handler.question_detector, - "is_owasp_question", - return_value=True, - ), - patch("apps.slack.events.message_posted.django_rq") as mock_django_rq, - ): - mock_queue = Mock() - mock_django_rq.get_queue.return_value = mock_queue + mock_queue = Mock() + mock_django_rq.get_queue.return_value = mock_queue + with patch.object( + message_handler.question_detector, + "is_owasp_question", + return_value=True, + ): message_handler.handle_event(event, client) + mock_member.update_data.assert_called_once() mock_django_rq.get_queue.assert_called_once() def test_handle_event_empty_text(self, message_handler): From 08afea27bfd5c796614e20e90075d8f372fd2d62 Mon Sep 17 00:00:00 2001 From: Dishant1804 Date: Wed, 8 Oct 2025 15:25:44 +0530 Subject: [PATCH 2/3] suggestions implemented --- backend/apps/slack/admin/__init__.py | 1 - backend/apps/slack/admin/chat.py | 16 --- backend/apps/slack/common/handlers/ai.py | 20 +-- backend/apps/slack/constants.py | 1 + backend/apps/slack/events/message_posted.py | 30 ++--- backend/apps/slack/migrations/0020_chat.py | 56 --------- .../0020_conversation_conversation_context.py | 17 +++ backend/apps/slack/models/__init__.py | 1 - backend/apps/slack/models/chat.py | 90 ------------- backend/apps/slack/models/conversation.py | 41 ++++++ .../apps/slack/models/conversation_test.py | 119 ++++++++++++++++++ 11 files changed, 201 insertions(+), 191 deletions(-) delete mode 100644 backend/apps/slack/admin/chat.py delete mode 100644 backend/apps/slack/migrations/0020_chat.py create mode 100644 backend/apps/slack/migrations/0020_conversation_conversation_context.py delete mode 100644 backend/apps/slack/models/chat.py diff --git a/backend/apps/slack/admin/__init__.py b/backend/apps/slack/admin/__init__.py index ef678fbb00..b3193ffb04 100644 --- a/backend/apps/slack/admin/__init__.py +++ b/backend/apps/slack/admin/__init__.py @@ -1,6 +1,5 @@ """Slack app admin.""" -from .chat import ChatAdmin from .conversation import ConversationAdmin from .event import EventAdmin from .member import MemberAdmin diff --git a/backend/apps/slack/admin/chat.py b/backend/apps/slack/admin/chat.py deleted file mode 100644 index d51a536a6e..0000000000 --- a/backend/apps/slack/admin/chat.py +++ /dev/null @@ -1,16 +0,0 @@ -"""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) diff --git a/backend/apps/slack/common/handlers/ai.py b/backend/apps/slack/common/handlers/ai.py index c39da35248..6681e8611e 100644 --- a/backend/apps/slack/common/handlers/ai.py +++ b/backend/apps/slack/common/handlers/ai.py @@ -6,7 +6,8 @@ from apps.ai.agent.tools.rag.rag_tool import RagTool from apps.slack.blocks import markdown -from apps.slack.models import Chat, Member, Workspace +from apps.slack.constants import CONVERSATION_CONTEXT_LIMIT +from apps.slack.models import Conversation, Workspace logger = logging.getLogger(__name__) @@ -47,42 +48,41 @@ 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]: +def get_dm_blocks(query: str, workspace_id: str, channel_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. + channel_id (str): Slack channel ID for the DM. Returns: list: A list of Slack blocks representing the AI response. """ - ai_response = process_dm_ai_query(query.strip(), user_id, workspace_id) + ai_response = process_dm_ai_query(query.strip(), workspace_id, channel_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: +def process_dm_ai_query(query: str, workspace_id: str, channel_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. + channel_id (str): Slack channel ID for the DM. 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) + conversation = Conversation.objects.get(slack_channel_id=channel_id, workspace=workspace) - chat = Chat.update_data(user, workspace) - context = chat.get_context(limit_exchanges=20) + context = conversation.get_context(conversation_context_limit=CONVERSATION_CONTEXT_LIMIT) rag_tool = RagTool( chat_model="gpt-4o", @@ -95,7 +95,7 @@ def process_dm_ai_query(query: str, user_id: str, workspace_id: str) -> str | No enhanced_query = query response = rag_tool.query(question=enhanced_query) - chat.add_to_context(query, response) + conversation.add_to_context(query, response) return response diff --git a/backend/apps/slack/constants.py b/backend/apps/slack/constants.py index ef9bb7b9bb..3bed72b5a4 100644 --- a/backend/apps/slack/constants.py +++ b/backend/apps/slack/constants.py @@ -2,6 +2,7 @@ from apps.common.constants import NL +CONVERSATION_CONTEXT_LIMIT = 20 NEST_BOT_NAME = "NestBot" OWASP_APPSEC_CHANNEL_ID = "#C0F7D6DFH" diff --git a/backend/apps/slack/events/message_posted.py b/backend/apps/slack/events/message_posted.py index 75aadb3d15..4558777350 100644 --- a/backend/apps/slack/events/message_posted.py +++ b/backend/apps/slack/events/message_posted.py @@ -81,31 +81,27 @@ def handle_event(self, event, client): def handle_dm(self, event, client, channel_id, user_id, text): """Handle direct messages with NestBot (DMs).""" workspace_id = event.get("team") + channel_info = client.conversations_info(channel=channel_id) - 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: + workspace = Workspace.objects.get(slack_workspace_id=workspace_id) + except Workspace.DoesNotExist: + logger.exception("Workspace not found for DM.") + return + + Conversation.update_data(channel_info["channel"], workspace) try: - Member.objects.get(slack_user_id=user_id, workspace__slack_workspace_id=workspace_id) + Member.objects.get(slack_user_id=user_id, workspace=workspace) 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 + user_info = client.users_info(user=user_id) + Member.update_data(user_info["user"], workspace, save=True) + logger.info("Created new member for DM") thread_ts = event.get("thread_ts") try: - response_blocks = get_dm_blocks(text, user_id, workspace_id) + response_blocks = get_dm_blocks(text, workspace_id, channel_id) if response_blocks: client.chat_postMessage( channel=channel_id, diff --git a/backend/apps/slack/migrations/0020_chat.py b/backend/apps/slack/migrations/0020_chat.py deleted file mode 100644 index 06f48062dc..0000000000 --- a/backend/apps/slack/migrations/0020_chat.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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")}, - }, - ), - ] diff --git a/backend/apps/slack/migrations/0020_conversation_conversation_context.py b/backend/apps/slack/migrations/0020_conversation_conversation_context.py new file mode 100644 index 0000000000..b39fe8e246 --- /dev/null +++ b/backend/apps/slack/migrations/0020_conversation_conversation_context.py @@ -0,0 +1,17 @@ +# Generated by Django 5.2.6 on 2025-10-08 07:21 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("slack", "0019_conversation_is_nest_bot_assistant_enabled"), + ] + + operations = [ + migrations.AddField( + model_name="conversation", + name="conversation_context", + field=models.TextField(blank=True, verbose_name="Conversation context"), + ), + ] diff --git a/backend/apps/slack/models/__init__.py b/backend/apps/slack/models/__init__.py index 10a0a0473f..3bbe0878de 100644 --- a/backend/apps/slack/models/__init__.py +++ b/backend/apps/slack/models/__init__.py @@ -1,4 +1,3 @@ -from .chat import Chat from .conversation import Conversation from .event import Event from .member import Member diff --git a/backend/apps/slack/models/chat.py b/backend/apps/slack/models/chat.py deleted file mode 100644 index a2ce4f5901..0000000000 --- a/backend/apps/slack/models/chat.py +++ /dev/null @@ -1,90 +0,0 @@ -"""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) :]) diff --git a/backend/apps/slack/models/conversation.py b/backend/apps/slack/models/conversation.py index 9735786c24..c0899c079f 100644 --- a/backend/apps/slack/models/conversation.py +++ b/backend/apps/slack/models/conversation.py @@ -42,6 +42,7 @@ class Meta: # Additional attributes. sync_messages = models.BooleanField(verbose_name="Sync messages", default=False) + conversation_context = models.TextField(blank=True, verbose_name="Conversation context") def __str__(self): """Channel human readable representation.""" @@ -105,3 +106,43 @@ def update_data(conversation_data, workspace, *, save=True): conversation.save() return conversation + + 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.conversation_context: + self.conversation_context = "" + + self.conversation_context = f"{self.conversation_context}{f'User: {user_message}\n'}" + + if bot_response: + self.conversation_context = f"{self.conversation_context}{f'Bot: {bot_response}\n'}" + + self.save(update_fields=["conversation_context"]) + + def get_context(self, conversation_context_limit: int | None = None) -> str: + """Get the conversation context. + + Args: + conversation_context_limit: Optional limit on number of exchanges to return. + + Returns: + The conversation context, potentially limited to recent exchanges. + + """ + if not self.conversation_context: + return "" + + if conversation_context_limit is None: + return self.conversation_context + + lines = self.conversation_context.strip().split("\n") + if len(lines) <= conversation_context_limit * 2: + return self.conversation_context + + return "\n".join(lines[-(conversation_context_limit * 2) :]) diff --git a/backend/tests/apps/slack/models/conversation_test.py b/backend/tests/apps/slack/models/conversation_test.py index cdf9732cd0..324ba0dca3 100644 --- a/backend/tests/apps/slack/models/conversation_test.py +++ b/backend/tests/apps/slack/models/conversation_test.py @@ -137,3 +137,122 @@ def test_str_method(self): # Check __str__ returns the name assert str(conversation) == "test-workspace #test-channel" + + def test_add_to_context_first_message(self, mocker): + """Test adding the first message to an empty conversation context.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = "" + + save_mock = mocker.patch.object(Conversation, "save") + conversation.add_to_context("Hello, how are you?") + + assert conversation.conversation_context == "User: Hello, how are you?\n" + save_mock.assert_called_once_with(update_fields=["conversation_context"]) + + def test_add_to_context_with_bot_response(self, mocker): + """Test adding a user message and bot response to context.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = "User: Hello\nBot: Hi there!\n" + + save_mock = mocker.patch.object(Conversation, "save") + + conversation.add_to_context( + "What is OWASP?", "OWASP stands for Open Web Application Security Project." + ) + + expected_context = ( + "User: Hello\n" + "Bot: Hi there!\n" + "User: What is OWASP?\n" + "Bot: OWASP stands for Open Web Application Security Project.\n" + ) + assert conversation.conversation_context == expected_context + save_mock.assert_called_once_with(update_fields=["conversation_context"]) + + def test_add_to_context_empty_initial_context(self, mocker): + """Test adding context when conversation_context is None.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = None + + save_mock = mocker.patch.object(Conversation, "save") + + conversation.add_to_context("First message", "First response") + + expected_context = "User: First message\nBot: First response\n" + assert conversation.conversation_context == expected_context + save_mock.assert_called_once_with(update_fields=["conversation_context"]) + + def test_get_context_empty(self): + """Test getting context when conversation_context is empty.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = "" + + result = conversation.get_context() + + assert result == "" + + def test_get_context_no_limit(self): + """Test getting full context without limit.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = ( + "User: Message 1\n" + "Bot: Response 1\n" + "User: Message 2\n" + "Bot: Response 2\n" + "User: Message 3\n" + "Bot: Response 3\n" + ) + + result = conversation.get_context() + + assert result == conversation.conversation_context + + def test_get_context_with_limit_below_threshold(self): + """Test getting context with limit when total exchanges are below limit.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = ( + "User: Message 1\nBot: Response 1\nUser: Message 2\nBot: Response 2\n" + ) + + result = conversation.get_context(conversation_context_limit=3) + + assert result == conversation.conversation_context + + def test_get_context_with_limit_above_threshold(self): + """Test getting context with limit when total exchanges exceed limit.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = ( + "User: Message 1\n" + "Bot: Response 1\n" + "User: Message 2\n" + "Bot: Response 2\n" + "User: Message 3\n" + "Bot: Response 3\n" + "User: Message 4\n" + "Bot: Response 4\n" + ) + + result = conversation.get_context(conversation_context_limit=2) + + expected = "User: Message 3\nBot: Response 3\nUser: Message 4\nBot: Response 4" + assert result == expected + + def test_get_context_none_context(self): + """Test getting context when conversation_context is None.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = None + + result = conversation.get_context() + + assert result == "" + + def test_get_context_with_limit_exact_threshold(self): + """Test getting context when exchanges exactly match the limit.""" + conversation = Conversation(slack_channel_id="C12345") + conversation.conversation_context = ( + "User: Message 1\nBot: Response 1\nUser: Message 2\nBot: Response 2\n" + ) + + result = conversation.get_context(conversation_context_limit=2) + + assert result == conversation.conversation_context From cb202a6748f459d277a793c9144ecfacd149050b Mon Sep 17 00:00:00 2001 From: Dishant1804 Date: Wed, 8 Oct 2025 15:58:49 +0530 Subject: [PATCH 3/3] code rabbit suggestions implemented --- backend/apps/slack/common/handlers/ai.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/backend/apps/slack/common/handlers/ai.py b/backend/apps/slack/common/handlers/ai.py index 6681e8611e..2056295f68 100644 --- a/backend/apps/slack/common/handlers/ai.py +++ b/backend/apps/slack/common/handlers/ai.py @@ -79,8 +79,12 @@ def process_dm_ai_query(query: str, workspace_id: str, channel_id: str) -> str | str | None: The AI response or None if error occurred. """ - workspace = Workspace.objects.get(slack_workspace_id=workspace_id) - conversation = Conversation.objects.get(slack_channel_id=channel_id, workspace=workspace) + try: + workspace = Workspace.objects.get(slack_workspace_id=workspace_id) + conversation = Conversation.objects.get(slack_channel_id=channel_id, workspace=workspace) + except (Workspace.DoesNotExist, Conversation.DoesNotExist): + logger.exception("Workspace or conversation not found for DM processing") + return None context = conversation.get_context(conversation_context_limit=CONVERSATION_CONTEXT_LIMIT)