Merge pull request #211 from vintasoftware/fix-bulk-create-messages

Fix save_django_messages on DBs where can_return_rows_from_bulk_insert=False

Author: fjsj

django_ai_assistant/helpers/django_messages.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING
 
-from django.db import transaction
+from django.db import connections, transaction
 
 from langchain_core.messages import (
     BaseMessage,
@@ -35,9 +35,23 @@ def save_django_messages(messages: list[BaseMessage], thread: "Thread") -> list[
 
     messages_to_create = [m for m in messages if m.id not in existing_message_ids]
 
-    created_messages = DjangoMessage.objects.bulk_create(
-        [DjangoMessage(thread=thread, message={}) for _ in messages_to_create]
-    )
+    # Insert in bulk only if primary keys are then assigned by the DB.
+    # Please check https://docs.djangoproject.com/en/4.0/ref/models/querysets/#django.db.models.query.QuerySet.bulk_create
+    # for more context on why this is required
+    can_bulk_insert = connections[
+        DjangoMessage.objects.db
+    ].features.can_return_rows_from_bulk_insert
+    if can_bulk_insert:
+        created_messages = DjangoMessage.objects.bulk_create(
+            [DjangoMessage(thread=thread, message={}) for _ in messages_to_create],
+        )
+    else:
+        for message in (
+            created_messages := [
+                DjangoMessage(thread=thread, message={}) for _ in messages_to_create
+            ]
+        ):
+            message.save()
 
     # Update langchain message IDs with Django message IDs
     for idx, created_message in enumerate(created_messages):

tests/test_helpers/test_django_messages.py

@@ -0,0 +1,52 @@
+from unittest.mock import patch
+
+from django.contrib.auth.models import User
+from django.db import connection, connections
+from django.db.backends.sqlite3.features import DatabaseFeatures
+
+import pytest
+from langchain_core.messages import HumanMessage
+from model_bakery import baker
+
+from django_ai_assistant.helpers.django_messages import save_django_messages
+from django_ai_assistant.models import Message, Thread
+
+
+@pytest.mark.django_db()
+def test_django_messages_with_can_return_rows_from_bulk_insert_true():
+    class MockFeatures(DatabaseFeatures):
+        can_return_rows_from_bulk_insert = True
+
+    mock_features = MockFeatures(connections[Message.objects.db])
+
+    thread = baker.make(Thread, created_by=baker.make(User))
+    with patch.object(
+        Message.objects,
+        "bulk_create",
+        wraps=Message.objects.bulk_create,
+    ) as mock_bulk_create:
+        with patch.object(connection, "features", mock_features):
+            save_django_messages([HumanMessage(content="Hello")], thread=thread)
+    mock_bulk_create.assert_called_once()
+    assert Message.objects.count() == 1
+    assert Message.objects.first().message["data"]["content"] == "Hello"
+
+
+@pytest.mark.django_db()
+def test_django_messages_with_can_return_rows_from_bulk_insert_false():
+    class MockFeatures(DatabaseFeatures):
+        can_return_rows_from_bulk_insert = False
+
+    mock_features = MockFeatures(connections[Message.objects.db])
+
+    thread = baker.make(Thread, created_by=baker.make(User))
+    with patch.object(
+        Message.objects,
+        "bulk_create",
+        wraps=Message.objects.bulk_create,
+    ) as mock_bulk_create:
+        with patch.object(connection, "features", mock_features):
+            save_django_messages([HumanMessage(content="Hello")], thread=thread)
+    mock_bulk_create.assert_not_called()
+    assert Message.objects.count() == 1
+    assert Message.objects.first().message["data"]["content"] == "Hello"