chat_backend/llm_be/chat_backend/consumers_graph.py

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 asgiref.sync import sync_to_async
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, 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,
    )
    return prompt_metric

@database_sync_to_async
def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_out=None):
    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)

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:
            # 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


# --- 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]
    model_name: str


# --- 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:
         # Check modelName first - if FAST, we skip search regardless of settings
         if state.get("model_name") == "FAST":
             pass
         elif 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)
            model = data.get("modelName", "Turbo")
            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,
                    "model_name": model
                }
                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))