westfarn
fc06202a1f
## Summary - Image attachments (PNG, JPEG, WebP, GIF) skip the document parser and go to a vision-capable Ollama model (`OLLAMA_MODEL_VISION`, default `gemma4:latest`). - The image is stored on the prompt. `PromptSerializer` returns `file_type` and base64 `image` so history can redraw it. - Unsupported image types and files over 8 MB return a websocket error and are not saved. - A later text turn does not reuse a previous image as a spreadsheet. Pairs with `chat_web_app` `feat/128-vision-image-prompt` ([#128](ai_ml_operations/chat_web_app#128)). ## Test plan - [ ] `manage.py test chat_backend.tests.test_vision chat_backend.tests.test_consumers.GraphNodeTestCase` - [ ] Send a prompt plus PNG over `ws/chat_again/` and confirm the vision model is called. - [ ] `GET conversation_details` includes `image` and `file_type` for that turn. - [ ] CSV upload still follows the document path. Reviewed-on: #82
405 lines
12 KiB
Python
405 lines
12 KiB
Python
from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
|
|
from rest_framework import serializers
|
|
from django.db.models import Count, Q, Sum
|
|
|
|
from .models import (
|
|
CustomUser,
|
|
Announcement,
|
|
Company,
|
|
Conversation,
|
|
Prompt,
|
|
PromptFeedback,
|
|
PromptMetric,
|
|
Feedback,
|
|
FEEDBACK_CATEGORIES,
|
|
DocumentWorkspace,
|
|
Document,
|
|
DriveConnection,
|
|
)
|
|
|
|
|
|
class MyTokenObtainPairSerializer(TokenObtainPairSerializer):
|
|
@classmethod
|
|
def get_token(cls, user):
|
|
token = super(MyTokenObtainPairSerializer, cls).get_token(user)
|
|
|
|
# add custom claim
|
|
token["company"] = "something here"
|
|
|
|
return token
|
|
|
|
def validate(self, attrs):
|
|
data = super().validate(attrs)
|
|
if getattr(self.user, "deleted", False):
|
|
from rest_framework_simplejwt.exceptions import AuthenticationFailed
|
|
|
|
raise AuthenticationFailed(
|
|
"No active account found with the given credentials",
|
|
code="no_active_account",
|
|
)
|
|
return data
|
|
|
|
|
|
class CompanySerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = Company
|
|
fields = "__all__"
|
|
|
|
|
|
class AnnouncmentSerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = Announcement
|
|
fields = "__all__"
|
|
|
|
|
|
class FeedbackSerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = Feedback
|
|
fields = "__all__"
|
|
|
|
|
|
class CustomUserSerializer(serializers.ModelSerializer):
|
|
email = serializers.EmailField(required=True)
|
|
username = serializers.CharField()
|
|
password = serializers.CharField(min_length=8, write_only=True)
|
|
company = CompanySerializer()
|
|
has_usable_password = serializers.BooleanField()
|
|
subscription = serializers.SerializerMethodField()
|
|
|
|
class Meta:
|
|
model = CustomUser
|
|
fields = "__all__"
|
|
extra_kwargs = {"password": {"write_only": True}}
|
|
|
|
def get_subscription(self, obj):
|
|
from monetization.services.plans import needs_checkout, plan_to_dict
|
|
from monetization.models import UserSubscription
|
|
|
|
try:
|
|
sub = obj.subscription
|
|
except UserSubscription.DoesNotExist:
|
|
return {
|
|
"plan": None,
|
|
"status": UserSubscription.Status.NONE,
|
|
"source": UserSubscription.Source.NONE,
|
|
"needs_checkout": True,
|
|
"cancel_at_period_end": False,
|
|
"current_period_end": None,
|
|
}
|
|
return {
|
|
"plan": plan_to_dict(sub.plan) if sub.plan_id else None,
|
|
"status": sub.status,
|
|
"source": sub.source,
|
|
"needs_checkout": needs_checkout(obj),
|
|
"cancel_at_period_end": bool(sub.cancel_at_period_end),
|
|
"current_period_end": (
|
|
sub.current_period_end.isoformat() if sub.current_period_end else None
|
|
),
|
|
}
|
|
|
|
|
|
class SelfServeRegistrationSerializer(serializers.Serializer):
|
|
"""Minimal payload for public self-serve sign-up (gated by settings)."""
|
|
|
|
email = serializers.EmailField(required=True)
|
|
password = serializers.CharField(min_length=8, write_only=True)
|
|
first_name = serializers.CharField(
|
|
required=False, allow_blank=True, max_length=150, default=""
|
|
)
|
|
last_name = serializers.CharField(
|
|
required=False, allow_blank=True, max_length=150, default=""
|
|
)
|
|
company_name = serializers.CharField(
|
|
required=False, allow_blank=True, max_length=256, default=""
|
|
)
|
|
|
|
def validate_email(self, value: str) -> str:
|
|
email = value.strip().lower()
|
|
if CustomUser.objects.filter(email__iexact=email).exists():
|
|
raise serializers.ValidationError("A user with this email already exists.")
|
|
if CustomUser.objects.filter(username__iexact=email).exists():
|
|
raise serializers.ValidationError("A user with this email already exists.")
|
|
return email
|
|
|
|
def create(self, validated_data):
|
|
from monetization.services.plans import try_redeem_backer_email
|
|
|
|
email = validated_data["email"]
|
|
password = validated_data["password"]
|
|
first_name = (validated_data.get("first_name") or "").strip()
|
|
last_name = (validated_data.get("last_name") or "").strip()
|
|
company_name = (validated_data.get("company_name") or "").strip()
|
|
if not company_name:
|
|
company_name = f"{email}'s workspace"
|
|
|
|
company = Company.objects.create(
|
|
name=company_name,
|
|
state="NA",
|
|
zipcode="00000",
|
|
address="N/A",
|
|
)
|
|
user = CustomUser.objects.create_user(
|
|
username=email,
|
|
email=email,
|
|
password=password,
|
|
first_name=first_name,
|
|
last_name=last_name,
|
|
company=company,
|
|
is_company_manager=True,
|
|
)
|
|
try_redeem_backer_email(user)
|
|
return user
|
|
|
|
|
|
def _conversation_token_totals(conversation_id: int):
|
|
"""
|
|
Sum PromptMetric tokens for a conversation.
|
|
|
|
Returns nulls when the provider never reported usage (never fabricate 0).
|
|
"""
|
|
agg = PromptMetric.objects.filter(conversation_id=conversation_id).aggregate(
|
|
tin=Sum("tokens_in"),
|
|
tout=Sum("tokens_out"),
|
|
with_in=Count("id", filter=Q(tokens_in__isnull=False)),
|
|
with_out=Count("id", filter=Q(tokens_out__isnull=False)),
|
|
)
|
|
return (
|
|
agg["tin"] if agg["with_in"] else None,
|
|
agg["tout"] if agg["with_out"] else None,
|
|
)
|
|
|
|
|
|
def _prompt_token_pair(prompt_id: int):
|
|
metric = (
|
|
PromptMetric.objects.filter(prompt_id=prompt_id)
|
|
.order_by("-created")
|
|
.only("tokens_in", "tokens_out")
|
|
.first()
|
|
)
|
|
if metric is None:
|
|
return None, None
|
|
return metric.tokens_in, metric.tokens_out
|
|
|
|
|
|
class ConversationSerializer(serializers.ModelSerializer):
|
|
tokens_in = serializers.SerializerMethodField()
|
|
tokens_out = serializers.SerializerMethodField()
|
|
|
|
class Meta:
|
|
model = Conversation
|
|
fields = (
|
|
"title",
|
|
"created",
|
|
"last_modified",
|
|
"id",
|
|
"tokens_in",
|
|
"tokens_out",
|
|
)
|
|
|
|
def _token_pair(self, obj):
|
|
cache = self.context.setdefault("_conversation_token_cache", {})
|
|
if obj.id not in cache:
|
|
cache[obj.id] = _conversation_token_totals(obj.id)
|
|
return cache[obj.id]
|
|
|
|
def get_tokens_in(self, obj):
|
|
tin, _ = self._token_pair(obj)
|
|
return tin
|
|
|
|
def get_tokens_out(self, obj):
|
|
_, tout = self._token_pair(obj)
|
|
return tout
|
|
|
|
|
|
class PromptFeedbackSerializer(serializers.ModelSerializer):
|
|
prompt_id = serializers.IntegerField(source="prompt.id", read_only=True)
|
|
|
|
class Meta:
|
|
model = PromptFeedback
|
|
fields = (
|
|
"id",
|
|
"prompt_id",
|
|
"rating",
|
|
"reason",
|
|
"comment",
|
|
"created",
|
|
"last_modified",
|
|
)
|
|
read_only_fields = ("id", "prompt_id", "created", "last_modified")
|
|
|
|
|
|
class PromptFeedbackUpsertSerializer(serializers.Serializer):
|
|
prompt_id = serializers.IntegerField()
|
|
rating = serializers.ChoiceField(choices=PromptFeedback.Rating.choices)
|
|
reason = serializers.ChoiceField(
|
|
choices=PromptFeedback.Reason.choices,
|
|
required=False,
|
|
allow_null=True,
|
|
allow_blank=True,
|
|
)
|
|
comment = serializers.CharField(
|
|
required=False, allow_null=True, allow_blank=True, max_length=1024
|
|
)
|
|
|
|
def validate_reason(self, value):
|
|
if value == "":
|
|
return None
|
|
return value
|
|
|
|
def validate_comment(self, value):
|
|
if value is None:
|
|
return None
|
|
stripped = str(value).strip()
|
|
return stripped or None
|
|
|
|
|
|
class PromptSerializer(serializers.ModelSerializer):
|
|
tokens_in = serializers.SerializerMethodField()
|
|
tokens_out = serializers.SerializerMethodField()
|
|
feedback = serializers.SerializerMethodField()
|
|
image = serializers.SerializerMethodField()
|
|
|
|
class Meta:
|
|
model = Prompt
|
|
fields = (
|
|
"message",
|
|
"user_created",
|
|
"created",
|
|
"id",
|
|
"tokens_in",
|
|
"tokens_out",
|
|
"citations",
|
|
"feedback",
|
|
"file_type",
|
|
"image",
|
|
)
|
|
read_only_fields = ("citations", "feedback", "file_type", "image")
|
|
|
|
def _token_pair(self, obj):
|
|
cache = self.context.setdefault("_prompt_token_cache", {})
|
|
if obj.id not in cache:
|
|
cache[obj.id] = _prompt_token_pair(obj.id)
|
|
return cache[obj.id]
|
|
|
|
def get_tokens_in(self, obj):
|
|
tin, _ = self._token_pair(obj)
|
|
return tin
|
|
|
|
def get_tokens_out(self, obj):
|
|
_, tout = self._token_pair(obj)
|
|
return tout
|
|
|
|
def get_feedback(self, obj):
|
|
"""Current caller's rating for this prompt, if any."""
|
|
request = self.context.get("request")
|
|
if request is None or not getattr(request, "user", None):
|
|
return None
|
|
user = request.user
|
|
if not user.is_authenticated:
|
|
return None
|
|
|
|
by_prompt = self.context.get("_prompt_feedback_by_id")
|
|
if by_prompt is None:
|
|
prompt_ids = self.context.get("_prompt_ids_for_feedback")
|
|
qs = PromptFeedback.objects.filter(user=user).only(
|
|
"prompt_id", "rating", "reason", "comment"
|
|
)
|
|
if prompt_ids is not None:
|
|
qs = qs.filter(prompt_id__in=prompt_ids)
|
|
by_prompt = {
|
|
row.prompt_id: {
|
|
"rating": row.rating,
|
|
"reason": row.reason,
|
|
"comment": row.comment,
|
|
}
|
|
for row in qs
|
|
}
|
|
self.context["_prompt_feedback_by_id"] = by_prompt
|
|
return by_prompt.get(obj.id)
|
|
|
|
def get_image(self, obj):
|
|
from chat_backend.services.vision import image_base64
|
|
|
|
return image_base64(obj)
|
|
|
|
def validate_message(self, value: str) -> str:
|
|
if value is None or not str(value).strip():
|
|
raise serializers.ValidationError("Message text cannot be empty.")
|
|
return str(value).strip()
|
|
|
|
|
|
class BasicUserSerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = CustomUser
|
|
fields = (
|
|
"email",
|
|
"first_name",
|
|
"last_name",
|
|
"is_active",
|
|
"has_usable_password",
|
|
"is_company_manager",
|
|
"has_signed_tos",
|
|
)
|
|
|
|
|
|
# document serializers
|
|
class DocumentWorkspaceSerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = DocumentWorkspace
|
|
fields = ["id", "name", "created"]
|
|
read_only_fields = ["id", "created"]
|
|
|
|
|
|
class DocumentSerializer(serializers.ModelSerializer):
|
|
class Meta:
|
|
model = Document
|
|
fields = [
|
|
"id",
|
|
"workspace",
|
|
"file",
|
|
"uploaded_at",
|
|
"processed",
|
|
"created",
|
|
"active",
|
|
]
|
|
read_only_fields = ["id", "uploaded_at", "processed", "created"]
|
|
|
|
|
|
# drive connection serializers (#47-#52)
|
|
class DriveConnectionSerializer(serializers.ModelSerializer):
|
|
"""Never exposes access_token/refresh_token to the client."""
|
|
|
|
class Meta:
|
|
model = DriveConnection
|
|
fields = [
|
|
"id",
|
|
"provider",
|
|
"kind",
|
|
"external_account_email",
|
|
"selected_resource_ids",
|
|
"selected_resource_labels",
|
|
"last_sync_at",
|
|
"last_sync_status",
|
|
"last_sync_error",
|
|
"sync_total",
|
|
"sync_processed",
|
|
"sync_added",
|
|
"sync_updated",
|
|
"sync_failed",
|
|
"is_active",
|
|
"created",
|
|
]
|
|
read_only_fields = fields
|
|
|
|
|
|
class DriveConnectionResourcesSerializer(serializers.Serializer):
|
|
resource_ids = serializers.ListField(
|
|
child=serializers.CharField(max_length=512), allow_empty=True
|
|
)
|
|
resource_labels = serializers.ListField(
|
|
child=serializers.CharField(max_length=512, allow_blank=True),
|
|
allow_empty=True,
|
|
required=False,
|
|
default=list,
|
|
)
|