westfarn
a6c45b0882
Implements #16/#17/#36: Founders/Standard/Pro/Business/Backer catalog, Backer email whitelist, prompt-window + token-period gates, and tokens_in/out on conversation/prompt + subscription usage APIs.
480 lines
18 KiB
Python
480 lines
18 KiB
Python
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
|
|
from .utils import (
|
|
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"
|
|
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 resolve_chat_user(email=None, conversation_id=None):
|
|
if email:
|
|
user = CustomUser.objects.filter(email__iexact=email).first()
|
|
if user:
|
|
return user
|
|
if conversation_id:
|
|
conversation = (
|
|
Conversation.objects.select_related("user")
|
|
.filter(id=conversation_id)
|
|
.first()
|
|
)
|
|
if conversation and conversation.user_id:
|
|
return conversation.user
|
|
return None
|
|
|
|
|
|
@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):
|
|
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
|
|
chat_user: Any
|
|
|
|
|
|
# --- 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
|
|
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)
|
|
# 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)
|
|
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
|
|
)
|
|
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)
|
|
|
|
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,
|
|
"chat_user": chat_user,
|
|
}
|
|
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)
|
|
tokens_in, tokens_out = extract_token_usage(
|
|
response_generator_or_dict
|
|
if isinstance(response_generator_or_dict, dict)
|
|
else None
|
|
)
|
|
await finish_prompt_metric(
|
|
prompt_metric,
|
|
len(full_response),
|
|
tokens_in=tokens_in,
|
|
tokens_out=tokens_out,
|
|
)
|