chat_backend/llm_be/chat_backend/consumers.py

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 chat_backend.ollama_config import (
    ollama_embeddings_kwargs,
    ollama_model_for_role,
    resolve_chat_role,
)
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, build_chat_service
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 .utils import (
    TokenUsageCollector,
    aiter_text_chunks,
    extract_token_usage,
    has_usable_user_prompt,
    is_heartbeat_payload,
    normalize_user_message,
)
from finance.services.quotas import (
    FeatureNotAllowed,
    QuotaExceeded,
    check_generation_allowed,
)

logger = logging.getLogger(__name__)

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


@database_sync_to_async
def create_conversation(prompt, email, title, user=None):
    """Create a conversation for ``user`` (preferred) or legacy ``email``."""
    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
):
    # conversation_id intentionally unused for identity — ownership is checked
    # via resolve_chat_company_scope after the principal is known.
    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 finance.services.quotas import assert_feature_allowed

    assert_feature_allowed(user, feature)


@database_sync_to_async
def get_workspace(conversation_id, user=None):
    """Resolve workspace only after conversation ownership is validated."""
    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)
    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, citations=None):
    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
        if citations:
            prompt_instance.citations = citations
        prompt_instance = serializer.save()
        # Ensure citations survive even if serializer omits write.
        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,
    )
    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, user=None):
    """Legacy helper — always applies a workspace metadata filter (fail closed)."""
    if user is None:
        raise ChatTenantScopeError(
            "Authenticated chat user is required.",
            code="user_not_found",
        )
    logger.info(f"getting workspace from conversation: {conversation_id}")
    scope = resolve_chat_company_scope(user, conversation_id)
    workspace = get_workspace_for_scope(scope)
    logger.info(f"Got workspace: {workspace.id} company={scope.company_id}")
    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(
        search_type="similarity",
        search_kwargs={"k": 4, "filter": {"workspace_id": workspace.id}},
    )

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):
        # 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):
        """
        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)
            token = data.get("token") or data.get("access")
            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

            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:
                # 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, 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:
                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 + plan gate: Image Generation
                    if prompt_type == PromptType.IMAGE_GENERATION:
                        if not getattr(settings, "ALLOW_IMAGE_GENERATION", False):
                             return {"type": "text", "content": "Image Generation is disabled."}
                        try:
                            await enforce_feature_gate(
                                chat_user, "image_generation"
                            )
                        except FeatureNotAllowed as exc:
                            return {
                                "type": "error",
                                "code": exc.code,
                                "content": exc.message,
                            }
                        return {"type": "text", "content": "Image Generation is not supported at this time, but it will be soon."}

                    if prompt_type == PromptType.RAG:
                        try:
                            await enforce_feature_gate(chat_user, "rag")
                        except FeatureNotAllowed as exc:
                            return {
                                "type": "error",
                                "code": exc.code,
                                "content": exc.message,
                            }
                        service = AsyncRAGService()
                        workspace = await get_workspace(
                            conversation_id, user=chat_user
                        )
                        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 / SEARCH / UNKNOWN — always-on grounding (#62).
                        # FAST selects a smaller model; it no longer skips search.
                        grounded = await prepare_grounded_chat(
                            message=input_dict["message"],
                            messages=messages,
                            model_name=input_dict.get("model_name"),
                            conversation_id=conversation_id,
                        )
                        if grounded.error:
                            return grounded.error
                        # Stash citations/model on the input for the caller.
                        input_dict["_citations"] = grounded.citations
                        input_dict["_resolved_model"] = grounded.model_name
                        return grounded.generator

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

                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,
                )

                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,
                    "_citations": [],
                    "_resolved_model": resolved_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 = ""
                tokens_in = tokens_out = None

                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
                    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 = step2.get("_citations") or []
                if citations:
                    await self.send_json_message(json.dumps(citations_frame(citations)))

                # Prefer model actually used by the grounded path when present.
                final_model = step2.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,
                )

        if bytes_data:
            logger.info("we have byte data")