Skip to content

Commit

Permalink
use get_user_model everywhere
Browse files Browse the repository at this point in the history
  • Loading branch information
George Burton committed Sep 12, 2024
1 parent 869352e commit a1c6b28
Show file tree
Hide file tree
Showing 11 changed files with 39 additions and 11 deletions.
3 changes: 2 additions & 1 deletion django_app/redbox_app/redbox_core/consumers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down
5 changes: 4 additions & 1 deletion django_app/redbox_app/redbox_core/serializers.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down
4 changes: 3 additions & 1 deletion django_app/redbox_app/redbox_core/views/document_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
4 changes: 3 additions & 1 deletion django_app/tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -22,9 +23,10 @@
Citation,
File,
StatusEnum,
User,
)

User = get_user_model()

logger = logging.getLogger(__name__)


Expand Down
6 changes: 5 additions & 1 deletion django_app/tests/management/test_commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ===

Expand Down
5 changes: 4 additions & 1 deletion django_app/tests/test_admin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)


Expand Down
5 changes: 4 additions & 1 deletion django_app/tests/test_consumers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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__)

Expand Down
4 changes: 3 additions & 1 deletion django_app/tests/views/test_chat_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)


Expand Down
4 changes: 3 additions & 1 deletion django_app/tests/views/test_citation_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -14,9 +15,10 @@
ChatRoleEnum,
Citation,
File,
User,
)

User = get_user_model()

logger = logging.getLogger(__name__)


Expand Down
5 changes: 4 additions & 1 deletion django_app/tests/views/test_document_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down
5 changes: 4 additions & 1 deletion django_app/tests/views/test_ratings_views.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down

0 comments on commit a1c6b28

Please sign in to comment.