westfarn
093e5462a5
## Summary - Closes Phase 4 of [#62](#62): `evals/suite.json` (≥40 graded questions), `run_evals` management command, and manually-triggered `.gitea/workflows/run-evals.yml`. - Emits versioned WS `status` frames during grounded chat (evaluating / searching / reading_sources / refining / writing) for [chat_web_app#96](ai_ml_operations/chat_web_app#96). - Implements [#63](#63): Redis/Celery optional infra, `AgentRun`/`AgentStep`, tool registry (SSRF-safe `fetch_url`, tenant-scoped docs), LangGraph orchestrator, progress frames, REST `GET/POST /api/agent_runs/…`, gated by `ALLOW_AGENTIC_TASKS` (default off). ## Test plan - [x] `SKIP_RAG_INIT=1 uv run python manage.py test` for evals, ws frames, agent tools, consumers, grounding - [ ] Manual: with `ALLOW_AGENTIC_TASKS=false`, chat identical to today - [ ] Manual: status frames visible in FE with #96 branch - [ ] Manual (GPU): `python manage.py run_evals --runs 3` - [ ] Manual: `ALLOW_AGENTIC_TASKS=true` multi-step research prompt creates AgentRun + framesReviewed-on: #71
691 lines
28 KiB
Python
691 lines
28 KiB
Python
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 prepare_grounded_chat
|
|
from .services.status_context import (
|
|
emit_status,
|
|
reset_status_emitter,
|
|
set_status_emitter,
|
|
)
|
|
from .services.ws_frames import citations_frame, status_frame
|
|
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_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 monetization.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,
|
|
}
|
|
await emit_status("retrieving_docs")
|
|
service = AsyncRAGService()
|
|
workspace = await get_workspace(
|
|
conversation_id, user=chat_user
|
|
)
|
|
await emit_status("refining")
|
|
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."}
|
|
await emit_status("analysing")
|
|
return service.generate_response(prompt_instance.message, decoded_file, file_type)
|
|
|
|
else:
|
|
# GENERAL_CHAT / SEARCH / UNKNOWN — agentic (#63) or grounded (#62).
|
|
from chat_backend.services.agent import (
|
|
run_agentic_turn,
|
|
should_use_agent,
|
|
)
|
|
|
|
if should_use_agent(input_dict["message"]):
|
|
try:
|
|
await enforce_feature_gate(
|
|
chat_user, "agentic_tasks"
|
|
)
|
|
except FeatureNotAllowed as exc:
|
|
return {
|
|
"type": "error",
|
|
"code": exc.code,
|
|
"content": exc.message,
|
|
}
|
|
|
|
async def _ws_send(raw: str):
|
|
await self.send_json_message(raw)
|
|
|
|
_run, answer = await run_agentic_turn(
|
|
user=chat_user,
|
|
scope=tenant_scope,
|
|
conversation_id=conversation_id,
|
|
goal=input_dict["message"],
|
|
prompt=prompt_instance,
|
|
ws_send=_ws_send,
|
|
)
|
|
input_dict["_resolved_model"] = (
|
|
_run.model_orchestrator or ""
|
|
)
|
|
input_dict["_citations"] = []
|
|
|
|
async def _agent_answer_gen():
|
|
yield answer
|
|
|
|
return _agent_answer_gen()
|
|
|
|
# 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,
|
|
}
|
|
|
|
# Send stream markers early so status frames reach the client
|
|
# during moderation / grounding (#96).
|
|
await self.send("CONVERSATION_ID")
|
|
await self.send(str(conversation_id))
|
|
await self.send("START_OF_THE_STREAM_ENDER_GAME_42")
|
|
|
|
async def _send_status(stage, detail=None):
|
|
await self.send_json_message(
|
|
json.dumps(status_frame(stage, detail=detail))
|
|
)
|
|
|
|
status_token = set_status_emitter(_send_status)
|
|
try:
|
|
await emit_status("queued")
|
|
await emit_status("moderating")
|
|
step1 = await check_moderation(pipeline_input)
|
|
step2 = await classify_prompt_step(step1)
|
|
|
|
response_generator_or_dict = await generate_response_step(step2)
|
|
|
|
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:
|
|
await emit_status("writing")
|
|
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))
|
|
)
|
|
|
|
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,
|
|
)
|
|
finally:
|
|
reset_status_emitter(status_token)
|
|
|
|
if bytes_data:
|
|
logger.info("we have byte data")
|