Ignore WS heartbeats and reject empty chat messages (#32)
## Summary - Closes #31 - Ignore WebSocket `type: ping` heartbeats so keepalives no longer create conversations or hit title/LLM pipelines - Reject empty/whitespace user messages in both chat consumers, `PromptSerializer`, and REST conversation prompt POST ## Test plan - [x] `UserPromptGuardTestCase`, `PromptSerializerTestCase` blank/whitespace cases - [x] `WebSocketReceiveGuardTestCase` ping ignore + empty message rejection (both WS routes) - [ ] Deploy to beta; leave idle tab open and confirm no new rogue conversations - [ ] Confirm normal chat send still works Related FE: https://git.aimloperations.com/ai_ml_operations/chat_web_app/issues/51Reviewed-on: #32
This commit was merged in pull request #32.
This commit is contained in:
@@ -27,6 +27,7 @@ from .services.title_generator import title_generator
|
|||||||
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
||||||
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
||||||
from .services.data_analysis_service import AsyncDataAnalysisService
|
from .services.data_analysis_service import AsyncDataAnalysisService
|
||||||
|
from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -238,13 +239,29 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
|
|||||||
logger.debug(f"Bytes Data: {bytes_data}")
|
logger.debug(f"Bytes Data: {bytes_data}")
|
||||||
if text_data:
|
if text_data:
|
||||||
data = json.loads(text_data)
|
data = json.loads(text_data)
|
||||||
message = data.get("message", None)
|
# Keepalive frames must not create conversations or hit the LLM.
|
||||||
|
if is_heartbeat_payload(data):
|
||||||
|
return
|
||||||
|
|
||||||
|
message = normalize_user_message(data.get("message", None))
|
||||||
conversation_id = data.get("conversation_id", None)
|
conversation_id = data.get("conversation_id", None)
|
||||||
email = data.get("email", None)
|
email = data.get("email", None)
|
||||||
file = data.get("file", None)
|
file = data.get("file", None)
|
||||||
file_type = data.get("fileType", "")
|
file_type = data.get("fileType", "")
|
||||||
model = data.get("modelName", "Turbo")
|
model = data.get("modelName", "Turbo")
|
||||||
|
|
||||||
|
if not has_usable_user_prompt(message, file):
|
||||||
|
logger.info("Ignoring websocket payload with empty message")
|
||||||
|
await self.send_json_message(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"content": "Message text cannot be empty.",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
if not conversation_id:
|
if not conversation_id:
|
||||||
# we need to create a new conversation
|
# we need to create a new conversation
|
||||||
# we will generate a name for it too
|
# we will generate a name for it too
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ from .services.title_generator import title_generator
|
|||||||
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
||||||
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
||||||
from .services.data_analysis_service import AsyncDataAnalysisService
|
from .services.data_analysis_service import AsyncDataAnalysisService
|
||||||
|
from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -288,13 +289,29 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
|
|||||||
print("Text Data: ", text_data)
|
print("Text Data: ", text_data)
|
||||||
if text_data:
|
if text_data:
|
||||||
data = json.loads(text_data)
|
data = json.loads(text_data)
|
||||||
|
# Keepalive frames must not create conversations or hit the LLM.
|
||||||
|
if is_heartbeat_payload(data):
|
||||||
|
return
|
||||||
|
|
||||||
model = data.get("modelName", "Turbo")
|
model = data.get("modelName", "Turbo")
|
||||||
message = data.get("message", None)
|
message = normalize_user_message(data.get("message", None))
|
||||||
conversation_id = data.get("conversation_id", None)
|
conversation_id = data.get("conversation_id", None)
|
||||||
email = data.get("email", None)
|
email = data.get("email", None)
|
||||||
file = data.get("file", None)
|
file = data.get("file", None)
|
||||||
file_type = data.get("fileType", "")
|
file_type = data.get("fileType", "")
|
||||||
|
|
||||||
|
if not has_usable_user_prompt(message, file):
|
||||||
|
logger.info("Ignoring websocket payload with empty message")
|
||||||
|
await self.send_json_message(
|
||||||
|
json.dumps(
|
||||||
|
{
|
||||||
|
"type": "error",
|
||||||
|
"content": "Message text cannot be empty.",
|
||||||
|
}
|
||||||
|
)
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
if not conversation_id:
|
if not conversation_id:
|
||||||
title = await title_generator.generate_async(message)
|
title = await title_generator.generate_async(message)
|
||||||
conversation_id = await create_conversation(message, email, title)
|
conversation_id = await create_conversation(message, email, title)
|
||||||
|
|||||||
@@ -122,6 +122,11 @@ class PromptSerializer(serializers.ModelSerializer):
|
|||||||
"id",
|
"id",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def validate_message(self, value: str) -> str:
|
||||||
|
if value is None or not str(value).strip():
|
||||||
|
raise serializers.ValidationError("Message text cannot be empty.")
|
||||||
|
return str(value).strip()
|
||||||
|
|
||||||
|
|
||||||
class BasicUserSerializer(serializers.ModelSerializer):
|
class BasicUserSerializer(serializers.ModelSerializer):
|
||||||
class Meta:
|
class Meta:
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
from unittest import mock
|
from unittest import mock
|
||||||
|
import json
|
||||||
|
|
||||||
from asgiref.sync import sync_to_async
|
from asgiref.sync import sync_to_async
|
||||||
from channels.testing import WebsocketCommunicator
|
from channels.testing import WebsocketCommunicator
|
||||||
@@ -412,3 +413,61 @@ class WebSocketRoutingTestCase(TransactionTestCase):
|
|||||||
[str(route.pattern) for route in websocket_urlpatterns],
|
[str(route.pattern) for route in websocket_urlpatterns],
|
||||||
["ws/chat_again/$", "ws/conditional_chat/$"],
|
["ws/chat_again/$", "ws/conditional_chat/$"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class WebSocketReceiveGuardTestCase(TransactionTestCase):
|
||||||
|
"""Heartbeats / empty prompts must not spawn conversations or LLM work."""
|
||||||
|
|
||||||
|
@parameterized.expand(
|
||||||
|
[("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")]
|
||||||
|
)
|
||||||
|
async def test_ping_heartbeat_is_ignored(self, _name, path):
|
||||||
|
communicator = WebsocketCommunicator(application, path)
|
||||||
|
connected, _ = await communicator.connect()
|
||||||
|
self.assertTrue(connected)
|
||||||
|
|
||||||
|
with mock.patch(
|
||||||
|
"chat_backend.consumers.title_generator.generate_async",
|
||||||
|
new_callable=mock.AsyncMock,
|
||||||
|
) as title_chat, mock.patch(
|
||||||
|
"chat_backend.consumers_graph.title_generator.generate_async",
|
||||||
|
new_callable=mock.AsyncMock,
|
||||||
|
) as title_graph:
|
||||||
|
await communicator.send_json_to({"type": "ping", "email": "a@b.com"})
|
||||||
|
# No reply expected; give the event loop a tick.
|
||||||
|
self.assertTrue(await communicator.receive_nothing(timeout=0.2))
|
||||||
|
title_chat.assert_not_called()
|
||||||
|
title_graph.assert_not_called()
|
||||||
|
|
||||||
|
count = await sync_to_async(Conversation.objects.count)()
|
||||||
|
self.assertEqual(count, 0)
|
||||||
|
await communicator.disconnect()
|
||||||
|
|
||||||
|
@parameterized.expand(
|
||||||
|
[("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")]
|
||||||
|
)
|
||||||
|
async def test_empty_message_is_rejected(self, _name, path):
|
||||||
|
communicator = WebsocketCommunicator(application, path)
|
||||||
|
connected, _ = await communicator.connect()
|
||||||
|
self.assertTrue(connected)
|
||||||
|
|
||||||
|
with mock.patch(
|
||||||
|
"chat_backend.consumers.title_generator.generate_async",
|
||||||
|
new_callable=mock.AsyncMock,
|
||||||
|
) as title_chat, mock.patch(
|
||||||
|
"chat_backend.consumers_graph.title_generator.generate_async",
|
||||||
|
new_callable=mock.AsyncMock,
|
||||||
|
) as title_graph:
|
||||||
|
await communicator.send_json_to(
|
||||||
|
{"message": " ", "email": "a@b.com", "conversation_id": None}
|
||||||
|
)
|
||||||
|
response = await communicator.receive_from(timeout=1)
|
||||||
|
payload = json.loads(response)
|
||||||
|
self.assertEqual(payload["type"], "error")
|
||||||
|
self.assertIn("empty", payload["content"].lower())
|
||||||
|
title_chat.assert_not_called()
|
||||||
|
title_graph.assert_not_called()
|
||||||
|
|
||||||
|
count = await sync_to_async(Conversation.objects.count)()
|
||||||
|
self.assertEqual(count, 0)
|
||||||
|
await communicator.disconnect()
|
||||||
|
|||||||
@@ -58,6 +58,22 @@ class PromptSerializerTestCase(TestCase):
|
|||||||
self.assertFalse(serializer.is_valid())
|
self.assertFalse(serializer.is_valid())
|
||||||
self.assertIn("message", serializer.errors)
|
self.assertIn("message", serializer.errors)
|
||||||
|
|
||||||
|
def test_message_rejects_blank_and_whitespace(self):
|
||||||
|
for payload in ("", " ", "\n\t"):
|
||||||
|
serializer = PromptSerializer(
|
||||||
|
data={"message": payload, "user_created": True}
|
||||||
|
)
|
||||||
|
self.assertFalse(serializer.is_valid(), payload)
|
||||||
|
self.assertIn("message", serializer.errors)
|
||||||
|
|
||||||
|
def test_message_is_stripped(self):
|
||||||
|
serializer = PromptSerializer(
|
||||||
|
data={"message": " hello ", "user_created": True}
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(serializer.is_valid(), serializer.errors)
|
||||||
|
self.assertEqual(serializer.validated_data["message"], "hello")
|
||||||
|
|
||||||
def test_user_created_is_required(self):
|
def test_user_created_is_required(self):
|
||||||
serializer = PromptSerializer(data={"message": "hi"})
|
serializer = PromptSerializer(data={"message": "hi"})
|
||||||
|
|
||||||
|
|||||||
@@ -11,7 +11,13 @@ from chat_backend.ollama_config import (
|
|||||||
ollama_llm_kwargs,
|
ollama_llm_kwargs,
|
||||||
ollama_model,
|
ollama_model,
|
||||||
)
|
)
|
||||||
from chat_backend.utils import extract_token_usage, last_day_of_month
|
from chat_backend.utils import (
|
||||||
|
extract_token_usage,
|
||||||
|
has_usable_user_prompt,
|
||||||
|
is_heartbeat_payload,
|
||||||
|
last_day_of_month,
|
||||||
|
normalize_user_message,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ExtractTokenUsageTestCase(SimpleTestCase):
|
class ExtractTokenUsageTestCase(SimpleTestCase):
|
||||||
@@ -134,3 +140,23 @@ class OllamaConfigFallbackTestCase(SimpleTestCase):
|
|||||||
with override_settings(OLLAMA_MODEL="llama3.2"):
|
with override_settings(OLLAMA_MODEL="llama3.2"):
|
||||||
del settings.OLLAMA_EMBED_MODEL
|
del settings.OLLAMA_EMBED_MODEL
|
||||||
self.assertEqual(ollama_embed_model(), "llama3.2")
|
self.assertEqual(ollama_embed_model(), "llama3.2")
|
||||||
|
|
||||||
|
|
||||||
|
class UserPromptGuardTestCase(SimpleTestCase):
|
||||||
|
def test_heartbeat_payload_detected(self):
|
||||||
|
self.assertTrue(is_heartbeat_payload({"type": "ping", "email": "a@b.com"}))
|
||||||
|
self.assertFalse(is_heartbeat_payload({"message": "hi"}))
|
||||||
|
self.assertFalse(is_heartbeat_payload(None))
|
||||||
|
|
||||||
|
def test_normalize_user_message(self):
|
||||||
|
self.assertIsNone(normalize_user_message(None))
|
||||||
|
self.assertIsNone(normalize_user_message(""))
|
||||||
|
self.assertIsNone(normalize_user_message(" \n\t"))
|
||||||
|
self.assertEqual(normalize_user_message(" hello "), "hello")
|
||||||
|
|
||||||
|
def test_has_usable_user_prompt(self):
|
||||||
|
self.assertFalse(has_usable_user_prompt(None))
|
||||||
|
self.assertFalse(has_usable_user_prompt(" "))
|
||||||
|
self.assertFalse(has_usable_user_prompt("", file="base64"))
|
||||||
|
self.assertTrue(has_usable_user_prompt("hi"))
|
||||||
|
self.assertTrue(has_usable_user_prompt(" hi "))
|
||||||
|
|||||||
@@ -1,6 +1,26 @@
|
|||||||
import datetime
|
import datetime
|
||||||
|
|
||||||
|
|
||||||
|
def is_heartbeat_payload(data) -> bool:
|
||||||
|
"""True for app-level WS keepalive frames (see FE buildHeartbeatPayload)."""
|
||||||
|
return isinstance(data, dict) and data.get("type") == "ping"
|
||||||
|
|
||||||
|
|
||||||
|
def normalize_user_message(message):
|
||||||
|
"""Return stripped message text, or None if missing/blank."""
|
||||||
|
if message is None:
|
||||||
|
return None
|
||||||
|
if not isinstance(message, str):
|
||||||
|
message = str(message)
|
||||||
|
stripped = message.strip()
|
||||||
|
return stripped or None
|
||||||
|
|
||||||
|
|
||||||
|
def has_usable_user_prompt(message, file=None) -> bool:
|
||||||
|
"""Reject empty/whitespace chat text. ``file`` kept for call-site clarity."""
|
||||||
|
return normalize_user_message(message) is not None
|
||||||
|
|
||||||
|
|
||||||
def last_day_of_month(any_day):
|
def last_day_of_month(any_day):
|
||||||
# The day 28 exists in every month. 4 days later, it's always next month
|
# The day 28 exists in every month. 4 days later, it's always next month
|
||||||
next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
|
next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
|
||||||
|
|||||||
@@ -469,6 +469,12 @@ class ConversationDetailView(APIView):
|
|||||||
|
|
||||||
# make sure that our model exists and it is running
|
# make sure that our model exists and it is running
|
||||||
prompt = request.data.get("prompt")
|
prompt = request.data.get("prompt")
|
||||||
|
if not isinstance(prompt, str) or not prompt.strip():
|
||||||
|
return Response(
|
||||||
|
{"detail": "Message text cannot be empty."},
|
||||||
|
status=status.HTTP_400_BAD_REQUEST,
|
||||||
|
)
|
||||||
|
prompt = prompt.strip()
|
||||||
|
|
||||||
conversation_id = request.data.get("conversation_id")
|
conversation_id = request.data.get("conversation_id")
|
||||||
is_user = bool(request.data.get("is_user"))
|
is_user = bool(request.data.get("is_user"))
|
||||||
|
|||||||
Reference in New Issue
Block a user