westfarn
ee3d47c8c3
## 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
453 lines
20 KiB
Python
453 lines
20 KiB
Python
import json
|
|
import base64
|
|
import logging
|
|
import pandas as pd
|
|
from datetime import datetime
|
|
from django.utils import timezone
|
|
from django.conf import settings
|
|
from django.core.files.base import ContentFile
|
|
from channels.generic.websocket import AsyncWebsocketConsumer
|
|
from channels.db import database_sync_to_async
|
|
from channels.layers import get_channel_layer
|
|
from asgiref.sync import sync_to_async, async_to_sync
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
|
from langchain_community.vectorstores import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_community.tools import DuckDuckGoSearchRun
|
|
from chat_backend.ollama_config import ollama_embeddings_kwargs
|
|
from django.conf import settings as django_settings
|
|
from langchain_core.runnables import RunnableLambda, RunnableBranch, RunnablePassthrough
|
|
from langchain_core.tracers.context import collect_runs
|
|
|
|
from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, Document, CustomUser
|
|
from .serializers import PromptSerializer
|
|
from .services.llm_service import AsyncLLMService
|
|
from .services.rag_services import AsyncRAGService
|
|
from .services.title_generator import title_generator
|
|
from .services.moderation_classifier import moderation_classifier, ModerationLabel
|
|
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
|
|
from .services.data_analysis_service import AsyncDataAnalysisService
|
|
from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CHANNEL_NAME: str = "llm_messages"
|
|
MODEL_NAME: str = "llama3.2"
|
|
PROMPT_CLASSIFIER = PromptClassifier()
|
|
|
|
@database_sync_to_async
|
|
def create_conversation(prompt, email, title):
|
|
# return the conversation id
|
|
conversation = Conversation.objects.create(title=title)
|
|
conversation.save()
|
|
|
|
user = CustomUser.objects.get(email=email)
|
|
conversation.user_id = user.id
|
|
conversation.save()
|
|
return conversation.id
|
|
|
|
|
|
@database_sync_to_async
|
|
def get_workspace(conversation_id):
|
|
conversation = Conversation.objects.get(id=conversation_id)
|
|
return DocumentWorkspace.objects.get(company=conversation.user.company)
|
|
|
|
|
|
@database_sync_to_async
|
|
def get_messages(conversation_id, prompt, file_string: str = None, file_type: str = ""):
|
|
messages = []
|
|
|
|
conversation = Conversation.objects.get(id=conversation_id)
|
|
logger.debug(file_string)
|
|
|
|
# add the prompt to the conversation
|
|
serializer = PromptSerializer(
|
|
data={
|
|
"message": prompt,
|
|
"user_created": True,
|
|
"created": timezone.now(),
|
|
}
|
|
)
|
|
if serializer.is_valid(raise_exception=True):
|
|
prompt_instance = serializer.save()
|
|
prompt_instance.conversation_id = conversation.id
|
|
prompt_instance.save()
|
|
if file_string:
|
|
file_name = f"prompt_{prompt_instance.id}_data.{file_type}"
|
|
f = ContentFile(file_string, name=file_name)
|
|
prompt_instance.file.save(file_name, f)
|
|
prompt_instance.file_type = file_type
|
|
prompt_instance.save()
|
|
|
|
for prompt_obj in Prompt.objects.filter(conversation__id=conversation_id):
|
|
messages.append(
|
|
{
|
|
"content": prompt_obj.message,
|
|
"role": "user" if prompt_obj.user_created else "assistant",
|
|
"has_file": prompt_obj.file_exists(),
|
|
"file": prompt_obj.file if prompt_obj.file_exists() else None,
|
|
"file_type": prompt_obj.file_type if prompt_obj.file_exists() else None,
|
|
}
|
|
)
|
|
|
|
# now transform the messages
|
|
transformed_messages = []
|
|
for message in messages:
|
|
|
|
if message["has_file"] and message["file_type"] != None:
|
|
if "csv" in message["file_type"]:
|
|
file_type = "csv"
|
|
altered_message = f"{message['content']}\n The file type is csv and the file contents are: {message['file'].read()}"
|
|
elif "xlsx" in message["file_type"]:
|
|
file_type = "xlsx"
|
|
df = pd.read_excel(message["file"].read())
|
|
altered_message = f"{message['content']}\n The file type is xlsx and the file contents are: {df}"
|
|
elif "txt" in message["file_type"]:
|
|
file_type = "txt"
|
|
altered_message = f"{message['content']}\n The file type is csv and the file contents are: {message['file'].read()}"
|
|
else:
|
|
altered_message = message["content"]
|
|
else:
|
|
altered_message = message["content"]
|
|
|
|
transformed_message = (
|
|
AIMessage(content=altered_message)
|
|
if message["role"] == "assistant"
|
|
else HumanMessage(content=altered_message)
|
|
)
|
|
transformed_messages.append(transformed_message)
|
|
|
|
return transformed_messages, prompt_instance
|
|
|
|
|
|
@database_sync_to_async
|
|
def save_generated_message(conversation_id, message):
|
|
conversation = Conversation.objects.get(id=conversation_id)
|
|
|
|
# add the prompt to the conversation
|
|
serializer = PromptSerializer(
|
|
data={
|
|
"message": message,
|
|
"user_created": False,
|
|
"created": timezone.now(),
|
|
}
|
|
)
|
|
if serializer.is_valid():
|
|
prompt_instance = serializer.save()
|
|
prompt_instance.conversation_id = conversation.id
|
|
prompt_instance = serializer.save()
|
|
else:
|
|
print(serializer.errors)
|
|
|
|
|
|
@database_sync_to_async
|
|
def create_prompt_metric(
|
|
prompt_id, prompt, has_file, file_type, model_name, conversation_id, tokens_in=None
|
|
):
|
|
prompt_metric = PromptMetric.objects.create(
|
|
prompt_id=prompt_id,
|
|
start_time=timezone.now(),
|
|
prompt_length=len(prompt),
|
|
tokens_in=tokens_in,
|
|
has_file=has_file,
|
|
file_type=file_type,
|
|
model_name=model_name,
|
|
conversation_id=conversation_id,
|
|
)
|
|
prompt_metric.save()
|
|
return prompt_metric
|
|
|
|
|
|
@database_sync_to_async
|
|
def update_prompt_metric(prompt_metric, status):
|
|
prompt_metric.event = status
|
|
prompt_metric.save()
|
|
|
|
|
|
@database_sync_to_async
|
|
def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_out=None):
|
|
logger.info(f"finish_prompt_metric: {response_length}")
|
|
prompt_metric.end_time = timezone.now()
|
|
prompt_metric.reponse_length = response_length
|
|
prompt_metric.event = "FINISHED"
|
|
update_fields = ["end_time", "reponse_length", "event"]
|
|
if tokens_in is not None:
|
|
prompt_metric.tokens_in = tokens_in
|
|
update_fields.append("tokens_in")
|
|
if tokens_out is not None:
|
|
prompt_metric.tokens_out = tokens_out
|
|
update_fields.append("tokens_out")
|
|
prompt_metric.save(update_fields=update_fields)
|
|
logger.info("finish_prompt_metric saved")
|
|
|
|
|
|
@database_sync_to_async
|
|
def get_retriever(conversation_id):
|
|
logger.info(f"getting workspace from conversation: {conversation_id}")
|
|
conversation = Conversation.objects.get(id=conversation_id)
|
|
logger.info(f"Got conversation: {conversation}")
|
|
workspace = DocumentWorkspace.objects.get(company=conversation.user.company)
|
|
logger.info(f"Got workspace: {conversation}")
|
|
persist_directory = getattr(
|
|
django_settings, "CHROMA_PERSIST_DIRECTORY", "./chroma_db/"
|
|
)
|
|
vectorstore = Chroma(
|
|
persist_directory=persist_directory,
|
|
embedding=OllamaEmbeddings(**ollama_embeddings_kwargs()),
|
|
)
|
|
return vectorstore.as_retriever()
|
|
|
|
async def get_conversation_file_async(conversation_id):
|
|
try:
|
|
# Get the very first prompt in the conversation that has a file
|
|
prompt_with_file = await Prompt.objects.filter(
|
|
conversation_id=conversation_id
|
|
).exclude(file='').order_by('created').afirst()
|
|
|
|
if prompt_with_file and prompt_with_file.file:
|
|
# Opening a DatabaseStorage file hits the DB, so read inside the thread.
|
|
file_data = await sync_to_async(lambda: prompt_with_file.file.read())()
|
|
file_type = prompt_with_file.file_type
|
|
return file_data, file_type
|
|
except Exception as e:
|
|
logger.error(f"Error retrieving file from conversation history: {e}")
|
|
return None, None
|
|
|
|
class ChatConsumerAgain(AsyncWebsocketConsumer):
|
|
async def connect(self):
|
|
await self.accept()
|
|
|
|
async def disconnect(self, close_code):
|
|
await self.close()
|
|
|
|
async def send_json_message(self, data_str):
|
|
"""
|
|
Ensures that the message sent over the websocket is a valid JSON object.
|
|
If data_str is a plain string, it wraps it in {"type": "text", "content": ...}.
|
|
"""
|
|
try:
|
|
# Test if it's already a valid JSON object string
|
|
json.loads(data_str)
|
|
# If it is, send it as is
|
|
await self.send(data_str)
|
|
except (json.JSONDecodeError, TypeError):
|
|
# If it's a plain string or not JSON-decodable, wrap it
|
|
await self.send(data_str)
|
|
|
|
async def receive(self, text_data=None, bytes_data=None):
|
|
logger.debug(f"Text Data: {text_data}")
|
|
logger.debug(f"Bytes Data: {bytes_data}")
|
|
if text_data:
|
|
data = json.loads(text_data)
|
|
# 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)
|
|
email = data.get("email", None)
|
|
file = data.get("file", None)
|
|
file_type = data.get("fileType", "")
|
|
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:
|
|
# we need to create a new conversation
|
|
# we will generate a name for it too
|
|
title = await title_generator.generate_async(message)
|
|
conversation_id = await create_conversation(message, email, title)
|
|
|
|
if conversation_id:
|
|
decoded_file = None
|
|
|
|
if file:
|
|
decoded_file = base64.b64decode(file)
|
|
logger.debug(decoded_file)
|
|
# The `altered_message` should only be created if a file exists
|
|
# and you want to pass its content directly to the classifier.
|
|
# Here, we'll let the classifier decide based on the user's prompt
|
|
# and then handle the file content separately.
|
|
altered_message = message
|
|
if "csv" in file_type:
|
|
file_type = "csv"
|
|
#altered_message = f"{message}\n The file type is csv and the file contents are: {decoded_file}"
|
|
elif "xmlformats-officedocument" in file_type:
|
|
file_type = "xlsx"
|
|
#df = pd.read_excel(decoded_file)
|
|
#altered_message = f"{message}\n The file type is xlsx and the file contents are: {df}"
|
|
elif "word" in file_type:
|
|
file_type = "docx"
|
|
elif "pdf" in file_type:
|
|
file_type = "pdf"
|
|
elif "text" in file_type:
|
|
file_type = "txt"
|
|
#altered_message = f"{message}\n The file type is txt and the file contents are: {decoded_file}"
|
|
else:
|
|
file_type = "Not Sure"
|
|
|
|
logger.info(f'received: "{message}" for conversation {conversation_id}')
|
|
|
|
# --- LangSmith Pipeline Construction ---
|
|
|
|
async def check_moderation(input_dict):
|
|
msg = input_dict["message"]
|
|
label = await moderation_classifier.classify_async(msg)
|
|
return {**input_dict, "moderation_label": label}
|
|
|
|
async def classify_prompt_step(input_dict):
|
|
if input_dict["moderation_label"] == ModerationLabel.NSFW:
|
|
return {**input_dict, "prompt_type": None} # Skip classification
|
|
|
|
msg = input_dict["message"]
|
|
decoded_file = input_dict.get("decoded_file")
|
|
|
|
prompt_type = await PROMPT_CLASSIFIER.classify_async(msg)
|
|
|
|
# Override logic
|
|
if decoded_file and (prompt_type == PromptType.DATA_ANALYSIS or 'analyze' in msg.lower() or 'data' in msg.lower()):
|
|
prompt_type = PromptType.DATA_ANALYSIS
|
|
elif decoded_file:
|
|
prompt_type = PromptType.GENERAL_CHAT
|
|
|
|
return {**input_dict, "prompt_type": prompt_type}
|
|
|
|
async def generate_response_step(input_dict):
|
|
if input_dict["moderation_label"] == ModerationLabel.NSFW:
|
|
response = "Prompt has been marked as NSFW. If this is in error, submit a feedback with the prompt text."
|
|
return {"type": "error", "content": response}
|
|
|
|
prompt_type = input_dict["prompt_type"]
|
|
messages = input_dict["messages"]
|
|
prompt_instance = input_dict["prompt_instance"]
|
|
conversation_id = input_dict["conversation_id"]
|
|
decoded_file = input_dict.get("decoded_file")
|
|
file_type = input_dict.get("file_type")
|
|
|
|
# Feature Flag: Image Generation
|
|
if prompt_type == PromptType.IMAGE_GENERATION:
|
|
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
|
|
return {"type": "text", "content": "Image Generation is disabled."}
|
|
# If enabled, proceed (assuming implementation exists, but user said "have it set to false for now")
|
|
return {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}
|
|
|
|
if prompt_type == PromptType.SEARCH:
|
|
# Check modelName first - if FAST, we skip search regardless of settings
|
|
if input_dict.get("model_name") == "FAST":
|
|
pass # Skip search
|
|
elif getattr(settings, "ALLOW_INTERNET_ACCESS", False):
|
|
try:
|
|
search = DuckDuckGoSearchRun()
|
|
search_results = search.run(input_dict["message"])
|
|
messages.append(HumanMessage(content=f"Search Results: {search_results}"))
|
|
except Exception as e:
|
|
logger.error(f"Search failed: {e}")
|
|
# If search fails, we proceed without it, essentially falling back to general chat
|
|
pass
|
|
else:
|
|
# If search is disabled, we could notify the user, but for now we'll just proceed
|
|
# potentially adding a system message or just letting the LLM handle it with its training data
|
|
pass
|
|
|
|
if prompt_type == PromptType.RAG:
|
|
service = AsyncRAGService()
|
|
workspace = await get_workspace(conversation_id)
|
|
return service.generate_response(messages, prompt_instance.message, workspace)
|
|
|
|
elif prompt_type == PromptType.DATA_ANALYSIS:
|
|
service = AsyncDataAnalysisService()
|
|
print(file_type)
|
|
if not decoded_file:
|
|
return {"type": "text", "content": "Please upload a file to perform data analysis."}
|
|
return service.generate_response(prompt_instance.message, decoded_file, file_type)
|
|
|
|
else: # GENERAL_CHAT or others
|
|
service = AsyncLLMService()
|
|
return service.generate_response(messages, prompt_instance.message, conversation_id)
|
|
|
|
# --- Execution ---
|
|
|
|
# Pre-fetch messages and file
|
|
messages, prompt_instance = await get_messages(
|
|
conversation_id, message, decoded_file, file_type
|
|
)
|
|
if not decoded_file:
|
|
decoded_file, file_type = await get_conversation_file_async(conversation_id)
|
|
|
|
if file:
|
|
# udpate with the altered_message (logic from original)
|
|
# Note: altered_message was defined in original but not fully used in the messages list construction in the same way
|
|
# In original: messages = messages[:-1] + [HumanMessage(content=altered_message)]
|
|
# I need to replicate that if I want exact behavior.
|
|
# But altered_message was only set if file was present.
|
|
pass # Logic is already in get_messages for the most part, but the original code had a specific override at the end.
|
|
# Let's trust get_messages for now or add the override if needed.
|
|
# Original:
|
|
# if file:
|
|
# messages = messages[:-1] + [HumanMessage(content=altered_message)]
|
|
# I'll add it to the input_dict if needed.
|
|
|
|
prompt_metric = await create_prompt_metric(
|
|
prompt_instance.id,
|
|
prompt_instance.message,
|
|
True if file else False,
|
|
file_type,
|
|
MODEL_NAME,
|
|
conversation_id,
|
|
)
|
|
|
|
pipeline_input = {
|
|
"message": message,
|
|
"conversation_id": conversation_id,
|
|
"decoded_file": decoded_file,
|
|
"file_type": file_type,
|
|
"messages": messages,
|
|
"prompt_instance": prompt_instance,
|
|
"model_name": model
|
|
}
|
|
|
|
# Run the pipeline steps manually to handle the async generator return type of generate_response_step
|
|
# A pure RunnableSequence might struggle with the async generator return.
|
|
# So I'll chain them in python but conceptually it's one pipeline.
|
|
|
|
step1 = await check_moderation(pipeline_input)
|
|
step2 = await classify_prompt_step(step1)
|
|
|
|
# Send start markers
|
|
await self.send("CONVERSATION_ID")
|
|
await self.send(str(conversation_id))
|
|
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
|
|
|
|
response_generator_or_dict = await generate_response_step(step2)
|
|
|
|
full_response = ""
|
|
|
|
if isinstance(response_generator_or_dict, dict):
|
|
# It's an error or simple message
|
|
content = response_generator_or_dict.get("content", "")
|
|
await self.send_json_message(json.dumps(response_generator_or_dict))
|
|
full_response = content
|
|
else:
|
|
# It's an async generator
|
|
async for chunk in response_generator_or_dict:
|
|
full_response += chunk
|
|
await self.send_json_message(chunk)
|
|
|
|
await self.send("END_OF_THE_STREAM_ENDER_GAME_42")
|
|
|
|
await save_generated_message(conversation_id, full_response)
|
|
await finish_prompt_metric(prompt_metric, len(full_response))
|
|
|
|
if bytes_data:
|
|
logger.info("we have byte data")
|