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 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.chat_tenant_scope import (
    ChatTenantScopeError,
    asgi_user_or_none,
    create_conversation_for_user,
    get_workspace_for_scope,
    resolve_chat_company_scope,
    resolve_chat_user as resolve_chat_user_sync,
)
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 .services.grounded_chat import citations_frame, prepare_grounded_chat
from chat_backend.ollama_config import ollama_model_for_role, resolve_chat_role
from .utils import (
    TokenUsageCollector,
    aiter_text_chunks,
    extract_token_usage,
    has_usable_user_prompt,
    is_heartbeat_payload,
    normalize_user_message,
)
from monetization.services.quotas import FeatureNotAllowed, QuotaExceeded, check_generation_allowed

logger = logging.getLogger(__name__)

CHANNEL_NAME: str = "llm_messages"
PROMPT_CLASSIFIER = PromptClassifier()

# --- Database Helpers (Reused) ---

@database_sync_to_async
def create_conversation(prompt, email, title, user=None):
    if user is None:
        user = CustomUser.objects.get(email=email)
    return create_conversation_for_user(user, title)


@database_sync_to_async
def resolve_chat_user(
    email=None, conversation_id=None, token=None, authenticated_user=None
):
    return resolve_chat_user_sync(
        email=email,
        token=token,
        authenticated_user=authenticated_user,
        conversation_id=conversation_id,
    )


@database_sync_to_async
def enforce_generation_gates(user, feature="text_generation"):
    return check_generation_allowed(user, feature=feature)


@database_sync_to_async
def enforce_feature_gate(user, feature):
    from monetization.services.quotas import assert_feature_allowed

    assert_feature_allowed(user, feature)

@database_sync_to_async
def get_workspace(conversation_id, user=None):
    if user is None:
        raise ChatTenantScopeError(
            "Authenticated chat user is required.",
            code="user_not_found",
        )
    scope = resolve_chat_company_scope(user, conversation_id)
    return get_workspace_for_scope(scope)


@database_sync_to_async
def resolve_tenant_scope(user, conversation_id=None):
    return resolve_chat_company_scope(user, conversation_id)

@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, citations=None):
    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()
        if citations is not None:
            Prompt.objects.filter(pk=prompt_instance.pk).update(citations=citations)
    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
    chat_user: Any
    citations: List[Dict[str, Any]]
    resolved_model: 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 + plan gate: 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."}}
        chat_user = state.get("chat_user")
        if chat_user is not None:
            try:
                await enforce_feature_gate(chat_user, "image_generation")
            except FeatureNotAllowed as exc:
                return {
                    "response_generator": {
                        "type": "error",
                        "code": exc.code,
                        "content": exc.message,
                    }
                }
        return {"response_generator": {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}}

    # Feature Flag: Internet Access / always-on grounding handled below for chat.
    if prompt_type == PromptType.RAG:
        chat_user = state.get("chat_user")
        if chat_user is not None:
            try:
                await enforce_feature_gate(chat_user, "rag")
            except FeatureNotAllowed as exc:
                return {
                    "response_generator": {
                        "type": "error",
                        "code": exc.code,
                        "content": exc.message,
                    }
                }
        service = AsyncRAGService()
        workspace = await get_workspace(conversation_id, user=chat_user)
        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 / SEARCH / UNKNOWN — always-on grounding (#62).
        # FAST selects a smaller model; it no longer skips search.
        grounded = await prepare_grounded_chat(
            message=state["message"],
            messages=messages,
            model_name=state.get("model_name"),
            conversation_id=conversation_id,
        )
        if grounded.error:
            return {
                "response_generator": grounded.error,
                "citations": [],
                "resolved_model": grounded.model_name or "",
            }
        return {
            "response_generator": grounded.generator,
            "citations": grounded.citations,
            "resolved_model": grounded.model_name,
        }


# --- 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):
        # Connection already closing — do not call self.close() again
        # (triggers ASGI 'websocket.close' after close completed).
        pass

    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)
            # Keepalive frames must not create conversations or hit the LLM.
            if is_heartbeat_payload(data):
                return

            model = data.get("modelName", "Turbo")
            message = normalize_user_message(data.get("message", None))
            conversation_id = data.get("conversation_id", None)
            email = data.get("email", None)
            token = data.get("token") or data.get("access")
            file = data.get("file", None)
            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

            chat_user = await resolve_chat_user(
                email=email,
                conversation_id=conversation_id,
                token=token,
                authenticated_user=asgi_user_or_none(self.scope.get("user")),
            )
            if chat_user is None:
                await self.send_json_message(
                    json.dumps(
                        {
                            "type": "error",
                            "code": "user_not_found",
                            "content": "Unable to resolve user for this chat session.",
                        }
                    )
                )
                return

            try:
                await enforce_generation_gates(chat_user, feature="text_generation")
            except (QuotaExceeded, FeatureNotAllowed) as exc:
                await self.send_json_message(
                    json.dumps(
                        {
                            "type": "error",
                            "code": exc.code,
                            "content": exc.message,
                            "details": getattr(exc, "details", {}),
                        }
                    )
                )
                return

            if not conversation_id:
                title = await title_generator.generate_async(message)
                conversation_id = await create_conversation(
                    message, email, title, user=chat_user
                )

            try:
                tenant_scope = await resolve_tenant_scope(chat_user, conversation_id)
            except ChatTenantScopeError as exc:
                logger.warning(
                    "websocket tenant validation failed conversation_id=%s user_id=%s code=%s",
                    conversation_id,
                    chat_user.id,
                    exc.code,
                )
                await self.send_json_message(
                    json.dumps(
                        {
                            "type": "error",
                            "code": exc.code,
                            "content": exc.message,
                        }
                    )
                )
                return

            logger.info(
                "chat_scope_validated conversation_id=%s user_id=%s company_id=%s workspace_id=%s",
                tenant_scope.conversation_id,
                tenant_scope.user_id,
                tenant_scope.company_id,
                tenant_scope.workspace_id,
            )

            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)

                resolved_model = ollama_model_for_role(resolve_chat_role(model))
                prompt_metric = await create_prompt_metric(
                    prompt_instance.id,
                    prompt_instance.message,
                    True if file else False,
                    file_type,
                    resolved_model,
                    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,
                    "chat_user": chat_user,
                    "citations": [],
                    "resolved_model": resolved_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 = ""
                tokens_in = tokens_out = None

                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
                    tokens_in, tokens_out = extract_token_usage(
                        response_generator_or_dict
                    )
                else:
                    # Stream raw LLM chunks so final Ollama generation_info
                    # (prompt_eval_count / eval_count) is not stripped.
                    usage = TokenUsageCollector()
                    async for chunk in aiter_text_chunks(
                        response_generator_or_dict, usage
                    ):
                        full_response += chunk
                        await self.send_json_message(chunk)
                    tokens_in, tokens_out = usage.pair

                await self.send("END_OF_THE_STREAM_ENDER_GAME_42")

                citations = final_state.get("citations") or []
                if citations:
                    await self.send_json_message(json.dumps(citations_frame(citations)))

                final_model = final_state.get("resolved_model") or resolved_model
                if final_model and final_model != prompt_metric.model_name:
                    prompt_metric.model_name = final_model
                    await database_sync_to_async(prompt_metric.save)(
                        update_fields=["model_name"]
                    )

                await save_generated_message(
                    conversation_id, full_response, citations=citations
                )
                await finish_prompt_metric(
                    prompt_metric,
                    len(full_response),
                    tokens_in=tokens_in,
                    tokens_out=tokens_out,
                )