updates
This commit is contained in:
westfarn committed
1 parent
91bdb2fd2d
commit
eed1abedc8
8 files changed
+1156
-532
No files matched your search
@@ -0,0 +1,357 @@
|
||||
import json
|
||||
import base64
|
||||
import logging
|
||||
import pandas as pd
|
||||
from datetime import datetime
|
||||
from typing import TypedDict, Annotated, List, Union, Dict, Any
|
||||
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 langchain_core.messages import HumanMessage, AIMessage, BaseMessage
|
||||
from langchain_community.tools import DuckDuckGoSearchRun
|
||||
from langgraph.graph import StateGraph, END
|
||||
|
||||
from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, 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
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CHANNEL_NAME: str = "llm_messages"
|
||||
MODEL_NAME: str = "llama3.2"
|
||||
PROMPT_CLASSIFIER = PromptClassifier()
|
||||
|
||||
# --- Database Helpers (Reused) ---
|
||||
|
||||
@database_sync_to_async
|
||||
def create_conversation(prompt, email, title):
|
||||
conversation = Conversation.objects.create(title=title)
|
||||
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)
|
||||
|
||||
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,
|
||||
}
|
||||
)
|
||||
|
||||
transformed_messages = []
|
||||
for message in messages:
|
||||
if message["has_file"] and message["file_type"] != None:
|
||||
# Simplified handling compared to original, as we rely on services to handle files now
|
||||
# But we keep the structure for context
|
||||
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)
|
||||
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.save()
|
||||
else:
|
||||
print(serializer.errors)
|
||||
|
||||
@database_sync_to_async
|
||||
def create_prompt_metric(prompt_id, prompt, has_file, file_type, model_name, conversation_id):
|
||||
prompt_metric = PromptMetric.objects.create(
|
||||
prompt_id=prompt_id,
|
||||
start_time=timezone.now(),
|
||||
prompt_length=len(prompt),
|
||||
has_file=has_file,
|
||||
file_type=file_type,
|
||||
model_name=model_name,
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
return prompt_metric
|
||||
|
||||
@database_sync_to_async
|
||||
def finish_prompt_metric(prompt_metric, response_length):
|
||||
prompt_metric.end_time = timezone.now()
|
||||
prompt_metric.reponse_length = response_length
|
||||
prompt_metric.event = "FINISHED"
|
||||
prompt_metric.save(update_fields=["end_time", "reponse_length", "event"])
|
||||
|
||||
async def get_conversation_file_async(conversation_id):
|
||||
try:
|
||||
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:
|
||||
file_data = await sync_to_async(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
|
||||
|
||||
|
||||
# --- LangGraph State ---
|
||||
|
||||
class ChatState(TypedDict):
|
||||
message: str
|
||||
conversation_id: int
|
||||
decoded_file: Union[bytes, None]
|
||||
file_type: Union[str, None]
|
||||
messages: List[BaseMessage]
|
||||
prompt_instance: Any # Django model instance
|
||||
moderation_label: Union[ModerationLabel, None]
|
||||
prompt_type: Union[PromptType, None]
|
||||
response_generator: Any # AsyncGenerator or dict
|
||||
error: Union[str, None]
|
||||
|
||||
|
||||
# --- LangGraph Nodes ---
|
||||
|
||||
async def moderation_node(state: ChatState) -> ChatState:
|
||||
msg = state["message"]
|
||||
label = await moderation_classifier.classify_async(msg)
|
||||
return {"moderation_label": label}
|
||||
|
||||
async def classification_node(state: ChatState) -> ChatState:
|
||||
if state.get("moderation_label") == ModerationLabel.NSFW:
|
||||
return {"prompt_type": None}
|
||||
|
||||
msg = state["message"]
|
||||
decoded_file = state.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 {"prompt_type": prompt_type}
|
||||
|
||||
async def generation_node(state: ChatState) -> ChatState:
|
||||
if state.get("moderation_label") == ModerationLabel.NSFW:
|
||||
response = "Prompt has been marked as NSFW. If this is in error, submit a feedback with the prompt text."
|
||||
return {"response_generator": {"type": "error", "content": response}}
|
||||
|
||||
prompt_type = state["prompt_type"]
|
||||
messages = state["messages"]
|
||||
prompt_instance = state["prompt_instance"]
|
||||
conversation_id = state["conversation_id"]
|
||||
decoded_file = state.get("decoded_file")
|
||||
file_type = state.get("file_type")
|
||||
|
||||
# Feature Flag: Image Generation
|
||||
if prompt_type == PromptType.IMAGE_GENERATION:
|
||||
if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
|
||||
return {"response_generator": {"type": "text", "content": "Image Generation is disabled."}}
|
||||
return {"response_generator": {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}}
|
||||
|
||||
# Feature Flag: Internet Access
|
||||
if prompt_type == PromptType.SEARCH:
|
||||
if getattr(settings, "ALLOW_INTERNET_ACCESS", False):
|
||||
try:
|
||||
search = DuckDuckGoSearchRun()
|
||||
search_results = search.run(state["message"])
|
||||
messages.append(HumanMessage(content=f"Search Results: {search_results}"))
|
||||
except Exception as e:
|
||||
logger.error(f"Search failed: {e}")
|
||||
pass
|
||||
else:
|
||||
pass
|
||||
|
||||
if prompt_type == PromptType.RAG:
|
||||
service = AsyncRAGService()
|
||||
workspace = await get_workspace(conversation_id)
|
||||
generator = service.generate_response(messages, prompt_instance.message, workspace)
|
||||
return {"response_generator": generator}
|
||||
|
||||
elif prompt_type == PromptType.DATA_ANALYSIS:
|
||||
service = AsyncDataAnalysisService()
|
||||
if not decoded_file:
|
||||
return {"response_generator": {"type": "text", "content": "Please upload a file to perform data analysis."}}
|
||||
generator = service.generate_response(prompt_instance.message, decoded_file, file_type)
|
||||
return {"response_generator": generator}
|
||||
|
||||
else: # GENERAL_CHAT or others
|
||||
service = AsyncLLMService()
|
||||
generator = service.generate_response(messages, prompt_instance.message, conversation_id)
|
||||
return {"response_generator": generator}
|
||||
|
||||
|
||||
# --- LangGraph Definition ---
|
||||
|
||||
workflow = StateGraph(ChatState)
|
||||
|
||||
workflow.add_node("moderation", moderation_node)
|
||||
workflow.add_node("classification", classification_node)
|
||||
workflow.add_node("generation", generation_node)
|
||||
|
||||
workflow.set_entry_point("moderation")
|
||||
|
||||
workflow.add_edge("moderation", "classification")
|
||||
workflow.add_edge("classification", "generation")
|
||||
workflow.add_edge("generation", END)
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
|
||||
# --- Consumer ---
|
||||
|
||||
class ChatConsumerGraph(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):
|
||||
try:
|
||||
json.loads(data_str)
|
||||
await self.send(data_str)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
await self.send(data_str)
|
||||
|
||||
async def receive(self, text_data=None, bytes_data=None):
|
||||
logger.debug(f"Text Data: {text_data}")
|
||||
print("Text Data: ", text_data)
|
||||
if text_data:
|
||||
data = json.loads(text_data)
|
||||
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", "")
|
||||
|
||||
if not conversation_id:
|
||||
title = await title_generator.generate_async(message)
|
||||
conversation_id = await create_conversation(message, email, title)
|
||||
|
||||
if conversation_id:
|
||||
print("Conversation ID: ", conversation_id)
|
||||
decoded_file = None
|
||||
if file:
|
||||
decoded_file = base64.b64decode(file)
|
||||
if "csv" in file_type: file_type = "csv"
|
||||
elif "xmlformats-officedocument" in file_type: file_type = "xlsx"
|
||||
elif "word" in file_type: file_type = "docx"
|
||||
elif "pdf" in file_type: file_type = "pdf"
|
||||
elif "text" in file_type: file_type = "txt"
|
||||
else: file_type = "Not Sure"
|
||||
|
||||
# Pre-fetch messages and file
|
||||
messages, prompt_instance = await get_messages(
|
||||
conversation_id, message, decoded_file, file_type
|
||||
)
|
||||
print("Messages: ", messages)
|
||||
if not decoded_file:
|
||||
decoded_file, file_type = await get_conversation_file_async(conversation_id)
|
||||
|
||||
prompt_metric = await create_prompt_metric(
|
||||
prompt_instance.id,
|
||||
prompt_instance.message,
|
||||
True if file else False,
|
||||
file_type,
|
||||
MODEL_NAME,
|
||||
conversation_id,
|
||||
)
|
||||
|
||||
# Initialize State
|
||||
initial_state = {
|
||||
"message": message,
|
||||
"conversation_id": conversation_id,
|
||||
"decoded_file": decoded_file,
|
||||
"file_type": file_type,
|
||||
"messages": messages,
|
||||
"prompt_instance": prompt_instance,
|
||||
"moderation_label": None,
|
||||
"prompt_type": None,
|
||||
"response_generator": None,
|
||||
"error": None
|
||||
}
|
||||
print("Initial State: ", initial_state)
|
||||
|
||||
# Run Graph
|
||||
final_state = await app.ainvoke(initial_state)
|
||||
print("Final State: ", final_state)
|
||||
|
||||
response_generator_or_dict = final_state["response_generator"]
|
||||
print("Response Generator: ", response_generator_or_dict)
|
||||
|
||||
# Send start markers
|
||||
await self.send("CONVERSATION_ID")
|
||||
await self.send(str(conversation_id))
|
||||
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
|
||||
|
||||
full_response = ""
|
||||
|
||||
if isinstance(response_generator_or_dict, dict):
|
||||
content = response_generator_or_dict.get("content", "")
|
||||
await self.send_json_message(json.dumps(response_generator_or_dict))
|
||||
full_response = content
|
||||
else:
|
||||
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))
|
||||
Reference in new issue
Block a user