From f2707bb8302e828d1efb78dcb876c7acac1a2d5e Mon Sep 17 00:00:00 2001 From: Quentin BEY Date: Mon, 28 Jul 2025 14:00:57 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=97=83=EF=B8=8F(chat)=20enforce=20`messag?= =?UTF-8?q?es`=20stored=20JSON=20format?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This enforces the models JSON format to `UIMessage`. It would be nicer to decouple the Vercel format from the stored one, but it's also more convenient for now, and will be quite easy to update later if needed. --- CHANGELOG.md | 1 + src/backend/chat/clients/pydantic_ai.py | 5 - .../0003_alter_chatconversation_messages.py | 32 ++ src/backend/chat/models.py | 9 +- src/backend/chat/serializers.py | 21 ++ .../tests/views/chat/test_conversation.py | 309 ++++++++++-------- src/backend/chat/views.py | 12 +- src/backend/pyproject.toml | 1 + 8 files changed, 240 insertions(+), 150 deletions(-) create mode 100644 src/backend/chat/migrations/0003_alter_chatconversation_messages.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d1763b..4ce649a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to ### Changed - ♻️(chat) rewrite backend using Pydantic AI SDK #4 +- 🗃️(chat) enforce messages stored JSON format #6 ### Added diff --git a/src/backend/chat/clients/pydantic_ai.py b/src/backend/chat/clients/pydantic_ai.py index 3fe4f39..b1e07ef 100644 --- a/src/backend/chat/clients/pydantic_ai.py +++ b/src/backend/chat/clients/pydantic_ai.py @@ -278,11 +278,6 @@ class AIAgentService: model_message_to_ui_message(msg) for msg in chain(history, [_merged_final_output_request, _merged_final_output_message]) ] - for message in self.conversation.messages: - logger.debug("conversation.messages: %s %s", type(message), message) - self.conversation.messages = [ - msg.model_dump(mode="json") for msg in self.conversation.messages if msg - ] self.conversation.agent_usage = usage logger.debug( diff --git a/src/backend/chat/migrations/0003_alter_chatconversation_messages.py b/src/backend/chat/migrations/0003_alter_chatconversation_messages.py new file mode 100644 index 0000000..a15e456 --- /dev/null +++ b/src/backend/chat/migrations/0003_alter_chatconversation_messages.py @@ -0,0 +1,32 @@ +# Generated by Django 5.2.3 on 2025-07-28 09:59 + +import django.core.serializers.json +from django.db import migrations + +import django_pydantic_field.compat.django +import django_pydantic_field.fields + +import chat.ai_sdk_types + + +class Migration(migrations.Migration): + dependencies = [ + ("chat", "0002_initial"), + ] + + operations = [ + migrations.AlterField( + model_name="chatconversation", + name="messages", + field=django_pydantic_field.fields.PydanticSchemaField( + blank=True, + config=None, + default=list, + encoder=django.core.serializers.json.DjangoJSONEncoder, + help_text="Stored messages for the chat conversation, sent to frontend", + schema=django_pydantic_field.compat.django.GenericContainer( + list, (chat.ai_sdk_types.UIMessage,) + ), + ), + ), + ] diff --git a/src/backend/chat/models.py b/src/backend/chat/models.py index e8e9a27..94e2741 100644 --- a/src/backend/chat/models.py +++ b/src/backend/chat/models.py @@ -1,10 +1,16 @@ """Models for chat conversations.""" +from typing import Sequence + from django.contrib.auth import get_user_model from django.db import models +from django_pydantic_field import SchemaField + from core.models import BaseModel +from chat.ai_sdk_types import UIMessage + User = get_user_model() @@ -48,7 +54,8 @@ class ChatConversation(BaseModel): blank=True, help_text="OpenAI messages for the chat conversation, not used", ) - messages = models.JSONField( + messages: Sequence[UIMessage] = SchemaField( + schema=list[UIMessage], default=list, blank=True, help_text="Stored messages for the chat conversation, sent to frontend", diff --git a/src/backend/chat/serializers.py b/src/backend/chat/serializers.py index a12b29c..73ab91c 100644 --- a/src/backend/chat/serializers.py +++ b/src/backend/chat/serializers.py @@ -1,16 +1,37 @@ """Serializers for chat application.""" +from django_pydantic_field.rest_framework import SchemaField # pylint: disable=no-name-in-module from rest_framework import serializers from chat import models +from chat.ai_sdk_types import UIMessage class ChatConversationSerializer(serializers.ModelSerializer): """Serializer for chat conversations.""" owner = serializers.HiddenField(default=serializers.CurrentUserDefault()) + messages = SchemaField(schema=list[UIMessage], read_only=True) class Meta: # pylint: disable=missing-class-docstring model = models.ChatConversation fields = ["id", "title", "created_at", "updated_at", "messages", "owner"] read_only_fields = ["id", "created_at", "updated_at", "messages"] + + +class ChatConversationInputSerializer(serializers.Serializer): + """ + Used to serialize input from Vercel AI SDK when using conversation endpoint. + + See ChatViewSet().post_conversation(...) method for more details. + """ + + messages = SchemaField(schema=list[UIMessage]) + + def update(self, instance, validated_data): + """Update method is not applicable in this context.""" + raise NotImplementedError("`update()` should not be used in this context.") + + def create(self, validated_data): + """Create method is not applicable in this context.""" + raise NotImplementedError("`create()` should not be used in this context.") diff --git a/src/backend/chat/tests/views/chat/test_conversation.py b/src/backend/chat/tests/views/chat/test_conversation.py index f064b3f..7f931ce 100644 --- a/src/backend/chat/tests/views/chat/test_conversation.py +++ b/src/backend/chat/tests/views/chat/test_conversation.py @@ -13,6 +13,13 @@ from rest_framework import status from core.factories import UserFactory +from chat.ai_sdk_types import ( + Attachment, + TextUIPart, + ToolInvocationCall, + ToolInvocationUIPart, + UIMessage, +) from chat.factories import ChatConversationFactory # enable database transactions for tests: @@ -435,29 +442,34 @@ def test_post_conversation_data_protocol(api_client, mock_openai_stream): ] assert len(chat_conversation.messages) == 2 - assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[0] == { - "annotations": None, - "content": "Hello", - "experimental_attachments": None, - "id": "", # ID is not set in the response - "parts": [{"text": "Hello", "type": "text"}], - "reasoning": None, - "role": "user", - "toolInvocations": None, - } - assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[1] == { - "annotations": None, - "content": "Hello there", - "experimental_attachments": None, - "id": "", # ID is not set in the response - "parts": [{"text": "Hello there", "type": "text"}], - "reasoning": None, - "role": "assistant", - "toolInvocations": None, - } + assert chat_conversation.messages[0].createdAt is not None + chat_conversation.messages[0].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[0] == UIMessage( + id="", + createdAt=None, + content="Hello", + reasoning=None, + experimental_attachments=None, + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Hello")], + ) + + assert chat_conversation.messages[1].createdAt is not None + chat_conversation.messages[1].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[1] == UIMessage( + id="", + createdAt=None, + content="Hello there", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Hello there")], + ) assert chat_conversation.openai_messages == [ { @@ -541,29 +553,33 @@ def test_post_conversation_text_protocol(api_client, mock_openai_stream): ] assert len(chat_conversation.messages) == 2 - assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[0] == { - "annotations": None, - "content": "Hello", - "experimental_attachments": None, - "id": "", # ID is not set in the response - "parts": [{"text": "Hello", "type": "text"}], - "reasoning": None, - "role": "user", - "toolInvocations": None, - } + assert chat_conversation.messages[0].createdAt is not None + chat_conversation.messages[0].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[0] == UIMessage( + id="", + createdAt=None, + content="Hello", + reasoning=None, + experimental_attachments=None, + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Hello")], + ) - assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[1] == { - "annotations": None, - "content": "Hello there", - "experimental_attachments": None, - "id": "", # ID is not set in the response - "parts": [{"text": "Hello there", "type": "text"}], - "reasoning": None, - "role": "assistant", - "toolInvocations": None, - } + assert chat_conversation.messages[1].createdAt is not None + chat_conversation.messages[1].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[1] == UIMessage( + id="", + createdAt=None, + content="Hello there", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Hello there")], + ) assert chat_conversation.openai_messages == [ { @@ -707,39 +723,43 @@ def test_post_conversation_with_image(api_client, mock_openai_stream_image): ] assert len(chat_conversation.messages) == 2 - assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[0] == { - "annotations": None, - "content": "Hello, what do you see on this picture?", - "experimental_attachments": [ - { - "contentType": "image/png", - "name": None, - "url": ( - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+wS" - "zIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEAAA" - "AASUVORK5CYII=" + assert chat_conversation.messages[0].createdAt is not None + chat_conversation.messages[0].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[0] == UIMessage( + id="", + createdAt=None, + content="Hello, what do you see on this picture?", + reasoning=None, + experimental_attachments=[ + Attachment( + name=None, + contentType="image/png", + url=( + "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAAIAQMAAAD+w" + "SzIAAAABlBMVEX///+/v7+jQ3Y5AAAADklEQVQI12P4AIX8EAgALgAD/aNpbtEA" + "AAAASUVORK5CYII=" ), - } + ) ], - "id": "", - "parts": [{"text": "Hello, what do you see on this picture?", "type": "text"}], - "reasoning": None, - "role": "user", - "toolInvocations": None, - } + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Hello, what do you see on this picture?")], + ) - assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[1] == { - "annotations": None, - "content": "I see a cat in the picture.", - "experimental_attachments": None, - "id": "", - "parts": [{"text": "I see a cat in the picture.", "type": "text"}], - "reasoning": None, - "role": "assistant", - "toolInvocations": None, - } + assert chat_conversation.messages[1].createdAt is not None + chat_conversation.messages[1].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[1] == UIMessage( + id="", + createdAt=None, + content="I see a cat in the picture.", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="I see a cat in the picture.")], + ) assert chat_conversation.openai_messages == [ { @@ -860,41 +880,45 @@ def test_post_conversation_tool_call(api_client, mock_openai_stream_tool, settin ] assert len(chat_conversation.messages) == 2 - assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[0] == { - "annotations": None, - "content": "Weather in Paris?", - "experimental_attachments": None, - "id": "", - "parts": [{"text": "Weather in Paris?", "type": "text"}], - "reasoning": None, - "role": "user", - "toolInvocations": None, - } + assert chat_conversation.messages[0].createdAt is not None + chat_conversation.messages[0].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[0] == UIMessage( + id="", + createdAt=None, + content="Weather in Paris?", + reasoning=None, + experimental_attachments=None, + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Weather in Paris?")], + ) - assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[1] == { - "annotations": None, - "content": "The current weather in Paris is nice", - "experimental_attachments": None, - "id": "", - "parts": [ - { - "toolInvocation": { - "args": {"location": "Paris", "unit": "celsius"}, - "state": "call", - "step": None, - "toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", - "toolName": "get_current_weather", - }, - "type": "tool-invocation", - }, - {"text": "The current weather in Paris is nice", "type": "text"}, + assert chat_conversation.messages[1].createdAt is not None + chat_conversation.messages[1].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[1] == UIMessage( + id="", + createdAt=None, + content="The current weather in Paris is nice", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[ + ToolInvocationUIPart( + type="tool-invocation", + toolInvocation=ToolInvocationCall( + toolCallId="xLDcIljdsDrz0idal7tATWSMm2jhMj47", + toolName="get_current_weather", + args={"unit": "celsius", "location": "Paris"}, + state="call", + step=None, + ), + ), + TextUIPart(type="text", text="The current weather in Paris is nice"), ], - "reasoning": None, - "role": "assistant", - "toolInvocations": None, - } + ) assert chat_conversation.openai_messages == [ { @@ -1040,41 +1064,46 @@ def test_post_conversation_tool_call_fails(api_client, mock_openai_stream_tool, ] assert len(chat_conversation.messages) == 2 - assert chat_conversation.messages[0].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[0] == { - "annotations": None, - "content": "Weather in Paris?", - "experimental_attachments": None, - "id": "", - "parts": [{"text": "Weather in Paris?", "type": "text"}], - "reasoning": None, - "role": "user", - "toolInvocations": None, - } - assert chat_conversation.messages[1].pop("createdAt") # Remove timestamp for comparison - assert chat_conversation.messages[1] == { - "annotations": None, - "content": "I cannot give you an answer to that.", - "experimental_attachments": None, - "id": "", - "parts": [ - { - "toolInvocation": { - "args": {"location": "Paris", "unit": "celsius"}, - "state": "call", - "step": None, - "toolCallId": "xLDcIljdsDrz0idal7tATWSMm2jhMj47", - "toolName": "get_current_weather", - }, - "type": "tool-invocation", - }, - {"text": "I cannot give you an answer to that.", "type": "text"}, + assert chat_conversation.messages[0].createdAt is not None + chat_conversation.messages[0].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[0] == UIMessage( + id="", + createdAt=None, + content="Weather in Paris?", + reasoning=None, + experimental_attachments=None, + role="user", + annotations=None, + toolInvocations=None, + parts=[TextUIPart(type="text", text="Weather in Paris?")], + ) + + assert chat_conversation.messages[1].createdAt is not None + chat_conversation.messages[1].createdAt = None # Remove timestamp for comparison + assert chat_conversation.messages[1] == UIMessage( + id="", + createdAt=None, + content="I cannot give you an answer to that.", + reasoning=None, + experimental_attachments=None, + role="assistant", + annotations=None, + toolInvocations=None, + parts=[ + ToolInvocationUIPart( + type="tool-invocation", + toolInvocation=ToolInvocationCall( + toolCallId="xLDcIljdsDrz0idal7tATWSMm2jhMj47", + toolName="get_current_weather", + args={"unit": "celsius", "location": "Paris"}, + state="call", + step=None, + ), + ), + TextUIPart(type="text", text="I cannot give you an answer to that."), ], - "reasoning": None, - "role": "assistant", - "toolInvocations": None, - } + ) assert chat_conversation.openai_messages == [ { diff --git a/src/backend/chat/views.py b/src/backend/chat/views.py index 786eb81..68094fd 100644 --- a/src/backend/chat/views.py +++ b/src/backend/chat/views.py @@ -8,11 +8,10 @@ from django.http import StreamingHttpResponse from rest_framework import decorators, filters, mixins, permissions, status, viewsets from rest_framework.response import Response -from core.api.viewsets import Pagination +from core.api.viewsets import Pagination, SerializerPerActionMixin from core.filters import remove_accents from chat import models, serializers -from chat.ai_sdk_types import UIMessage from chat.clients.pydantic_ai import AIAgentService logger = logging.getLogger(__name__) @@ -28,7 +27,8 @@ class ChatConversationFilter(filters.BaseFilterBackend): return queryset -class ChatViewSet( +class ChatViewSet( # pylint: disable=too-many-ancestors + SerializerPerActionMixin, mixins.CreateModelMixin, mixins.RetrieveModelMixin, mixins.ListModelMixin, @@ -51,6 +51,7 @@ class ChatViewSet( permissions.IsAuthenticated, ] serializer_class = serializers.ChatConversationSerializer + post_conversation_serializer_class = serializers.ChatConversationInputSerializer filter_backends = [filters.OrderingFilter, ChatConversationFilter] ordering = ["-created_at"] ordering_fields = ["created_at", "updated_at"] @@ -102,7 +103,10 @@ class ChatViewSet( conversation.ui_messages = request.data.get("messages", []) conversation.save() - messages = [UIMessage(**msg) for msg in request.data.get("messages", [])] + serializer = self.get_serializer(data=request.data) + if not serializer.is_valid(): + return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST) + messages = serializer.validated_data["messages"] logger.info("Received messages: %s", messages) logger.info("Using protocol: %s", protocol) diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 0a8210f..e029f9f 100644 --- a/src/backend/pyproject.toml +++ b/src/backend/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "django-filter==25.1", "django-lasuite[all]==0.0.9", "django-parler==2.3", + "django-pydantic-field==0.3.13", "django-redis==5.4.0", "django-storages[s3]==1.14.6", "django-timezone-field>=5.1",