chore/use get_user_model everywhere

django_app/redbox_app/redbox_core/consumers.py

@@ -8,6 +8,7 @@
 from channels.db import database_sync_to_async
 from channels.generic.websocket import AsyncWebsocketConsumer
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.forms.models import model_to_dict
 from django.utils import timezone
 from langchain_core.documents import Document
@@ -26,9 +27,9 @@
     ChatRoleEnum,
     Citation,
     File,
-    User,
 )
 
+User = get_user_model()
 OptFileSeq = Sequence[File] | None
 logger = logging.getLogger(__name__)
 logger.info("WEBSOCKET_SCHEME is: %s", settings.WEBSOCKET_SCHEME)

django_app/redbox_app/redbox_core/serializers.py

@@ -1,6 +1,9 @@
+from django.contrib.auth import get_user_model
 from rest_framework import serializers
 
-from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, File, User
+from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, File
+
+User = get_user_model()
 
 
 class FileSerializer(serializers.ModelSerializer):

django_app/redbox_app/redbox_core/views/document_views.py

@@ -3,6 +3,7 @@
 from collections.abc import MutableSequence, Sequence
 from pathlib import Path
 
+from django.contrib.auth import get_user_model
 from django.contrib.auth.decorators import login_required
 from django.core.exceptions import FieldError, ValidationError
 from django.core.files.uploadedfile import UploadedFile
@@ -15,9 +16,10 @@
 from django_q.tasks import async_task
 from requests.exceptions import RequestException
 
-from redbox_app.redbox_core.models import File, StatusEnum, User
+from redbox_app.redbox_core.models import File, StatusEnum
 from redbox_app.worker import ingest
 
+User = get_user_model()
 logger = logging.getLogger(__name__)
 CHUNK_SIZE = 1024
 # move this somewhere

django_app/tests/conftest.py

@@ -8,6 +8,7 @@
 import pytest
 from botocore.exceptions import ClientError
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.core.files.uploadedfile import SimpleUploadedFile, UploadedFile
 from django.core.management import call_command
 from django.utils import timezone
@@ -22,9 +23,10 @@
     Citation,
     File,
     StatusEnum,
-    User,
 )
 
+User = get_user_model()
+
 logger = logging.getLogger(__name__)
 
 

django_app/tests/management/test_commands.py

@@ -6,13 +6,17 @@
 import pytest
 from botocore.exceptions import UnknownClientMethodError
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.core.management import CommandError, call_command
 from django.utils import timezone
 from freezegun import freeze_time
 from magic_link.models import MagicLink
 from requests_mock import Mocker
 
-from redbox_app.redbox_core.models import Chat, ChatMessage, ChatRoleEnum, File, StatusEnum, User
+from redbox_app.redbox_core.models import Chat, ChatMessage, ChatRoleEnum, File, StatusEnum
+
+User = get_user_model()
+
 
 # === check_file_status command tests ===
 

django_app/tests/test_admin.py

@@ -6,13 +6,16 @@
 
 import pytest
 from bs4 import BeautifulSoup
+from django.contrib.auth import get_user_model
 from django.test import Client
 from django.urls import reverse
 from yarl import URL
 
-from redbox_app.redbox_core.models import ChatMessage, User
+from redbox_app.redbox_core.models import ChatMessage
 from redbox_app.redbox_core.serializers import ChatMessageSerializer, ChatSerializer, UserSerializer
 
+User = get_user_model()
+
 logger = logging.getLogger(__name__)
 
 

django_app/tests/test_consumers.py

@@ -9,6 +9,7 @@
 import pytest
 from channels.db import database_sync_to_async
 from channels.testing import WebsocketCommunicator
+from django.contrib.auth import get_user_model
 from django.db.models import Model
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseChatModel
@@ -20,9 +21,11 @@
 from redbox.models.chat import MetadataDetail
 from redbox_app.redbox_core import error_messages
 from redbox_app.redbox_core.consumers import ChatConsumer
-from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File, User
+from redbox_app.redbox_core.models import Chat, ChatMessage, ChatMessageTokenUse, ChatRoleEnum, File
 from redbox_app.redbox_core.prompts import CHAT_MAP_QUESTION_PROMPT
 
+User = get_user_model()
+
 logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"))
 logger = logging.getLogger(__name__)
 

django_app/tests/views/test_chat_views.py

@@ -5,15 +5,17 @@
 
 import pytest
 from bs4 import BeautifulSoup
+from django.contrib.auth import get_user_model
 from django.test import Client
 from django.urls import reverse
 
 from redbox_app.redbox_core.models import (
     Chat,
     ChatMessage,
-    User,
 )
 
+User = get_user_model()
+
 logger = logging.getLogger(__name__)
 
 

django_app/tests/views/test_citation_views.py

@@ -5,6 +5,7 @@
 
 import pytest
 from bs4 import BeautifulSoup
+from django.contrib.auth import get_user_model
 from django.test import Client
 from django.urls import reverse
 
@@ -14,9 +15,10 @@
     ChatRoleEnum,
     Citation,
     File,
-    User,
 )
 
+User = get_user_model()
+
 logger = logging.getLogger(__name__)
 
 

django_app/tests/views/test_document_views.py

@@ -6,10 +6,13 @@
 import pytest
 from botocore.exceptions import ClientError
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.test import Client
 from django.urls import reverse
 
-from redbox_app.redbox_core.models import File, StatusEnum, User
+from redbox_app.redbox_core.models import File, StatusEnum
+
+User = get_user_model()
 
 logger = logging.getLogger(__name__)
 

django_app/tests/views/test_ratings_views.py

@@ -3,10 +3,13 @@
 from http import HTTPStatus
 
 import pytest
+from django.contrib.auth import get_user_model
 from django.test import Client
 from django.urls import reverse
 
-from redbox_app.redbox_core.models import ChatMessage, User
+from redbox_app.redbox_core.models import ChatMessage
+
+User = get_user_model()
 
 logger = logging.getLogger(__name__)