Merge pull request #565 from i-dot-ai/feature/REDBOX-337-tests-for-ch…

…at-file-selection

REDBOX-337 - Tests for chat file selection

Makefile

@@ -1,81 +1,101 @@
+makefile_name := $(word $(words $(MAKEFILE_LIST)),$(MAKEFILE_LIST))
 
 -include .env
 
-.PHONY: app reqs
+default: help
 
+.PHONY: reqs
 reqs:
 	poetry install
 
-run:
-	docker compose up -d elasticsearch kibana worker minio redis core-api db django-app
+.PHONY: run
+run: stop
+	docker compose up -d --wait core-api django-app
 
-stop:
+.PHONY: stop
+stop: ## Stop all containers
 	docker compose down
 
+.PHONY: clean
 clean:
 	docker compose down -v --rmi all --remove-orphans
 
+.PHONY: build
 build:
 	docker compose build
 
-rebuild: stop
+.PHONY: rebuild
+rebuild: stop ## Rebuild all images
 	docker compose build --no-cache
 
-test-core-api:
+.PHONY: test-core-api
+test-core-api: ## Test core-api
 	poetry install --no-root --no-ansi --with api,dev,ai --without worker,docs
 	poetry run pytest core_api/tests --cov=core_api/src -v --cov-report=term-missing --cov-fail-under=75
 
-test-redbox:
+.PHONY: test-redbox
+test-redbox: ## Test redbox
 	poetry install --no-root --no-ansi --with api,dev --without ai,worker,docs
 	poetry run pytest redbox/tests --cov=redbox -v --cov-report=term-missing --cov-fail-under=80
 
-test-worker:
+.PHONY: test-worker
+test-worker: ## Test worker
 	poetry install --no-root --no-ansi --with worker,dev --without ai,api,docs
 	poetry run pytest worker/tests --cov=worker -v --cov-report=term-missing --cov-fail-under=40
 
-test-django: stop
+.PHONY: test-django
+test-django: stop ## Test django-app
 	docker compose up -d --wait db minio
 	docker compose run --no-deps django-app venv/bin/pytest tests/ --ds redbox_app.settings -v --cov=redbox_app.redbox_core --cov-fail-under 80 -o log_cli=true
 
-test-integration: stop
-	docker compose up -d --wait core-api django-app
+.PHONY: test-integration
+test-integration: stop run ## Run all integration tests
 	poetry install --no-root --no-ansi --with dev --without ai,api,worker,docs
 	poetry run pytest tests/
 
+.PHONY: collect-static
 collect-static:
 	docker compose run django-app venv/bin/django-admin collectstatic --noinput
 
-lint:
+.PHONY: lint
+lint:  ## Check code formatting & linting
 	poetry run ruff format . --check
 	poetry run ruff check .
 
-format:
+.PHONY: format
+format:  ## Format and fix code
 	poetry run ruff format .
 	poetry run ruff check . --fix
 
-safe:
+.PHONY: safe
+safe:  ##
 	poetry run bandit -ll -r ./redbox
 	poetry run bandit -ll -r ./django_app
 	poetry run mypy ./redbox --ignore-missing-imports
 	poetry run mypy ./django_app --ignore-missing-imports
 
-checktypes:
+.PHONY: checktypes
+checktypes:  ## Check types in redbox and worker
 	poetry run mypy redbox worker --ignore-missing-imports --no-incremental
 
-check-migrations: stop
+.PHONY: check-migrations
+check-migrations: stop  ## Check types in redbox and worker
 	docker compose build django-app
 	docker compose up -d --wait db minio
 	docker compose run --no-deps django-app venv/bin/django-admin migrate
 	docker compose run --no-deps django-app venv/bin/django-admin makemigrations --check
 
-reset-db:
+.PHONY: reset-db
+reset-db:  ## Reset Django database
 	docker compose down db --volumes
 	docker compose up -d db
 
-docs-serve:
+.PHONY: docs-serve
+docs-serve:  ## Build and serve documentation
 	poetry run mkdocs serve
 
-docs-build:
+.PHONY: docs-build
+docs-build:  ## Build documentation
 	poetry run mkdocs build
 
 # Docker
@@ -140,7 +160,7 @@ docker_update_tag:
 	done
 
 
-# Ouputs the value that you're after - usefx	ul to get a value i.e. IMAGE_TAG out of the Makefile
+# Ouputs the value that you're after - useful to get a value i.e. IMAGE_TAG out of the Makefile
 .PHONY: docker_echo
 docker_echo:
 	echo $($(value))
@@ -216,3 +236,7 @@ release: ## Deploy app
 eval_backend:
 	docker compose up core-api worker -d --build
 	docker exec -it $$(docker ps -q --filter "name=minio") mc mb data/$${BUCKET_NAME}
+
+.PHONY: help
+help: ## Show this help
+	@ grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(makefile_name) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-30s\033[0m %s\n", $$1,$$2}'

README.md

@@ -37,6 +37,7 @@ You will need to install `poppler` and `tesseract` to run the `worker`
 # Testing
 - Unit tests and QA run in CI
 - At this time integration test(s) take 10+ mins to run so are triggered manually in CI
+- Run `make help` to see all the available build activities.
 
 # Dependencies
 

django_app/redbox_app/redbox_core/admin.py

@@ -9,10 +9,13 @@
 class UserResource(admin.ModelAdmin):
     fields = ["email", "is_superuser", "is_staff", "last_login"]
     list_display = ["email", "is_superuser", "is_staff", "last_login"]
+    date_hierarchy = "last_login"
 
 
 class FileResource(admin.ModelAdmin):
-    list_display = ["original_file_name", "user", "status"]
+    list_display = ["original_file_name", "user", "status", "created_at"]
+    list_filter = ["user", "status"]
+    date_hierarchy = "created_at"
 
 
 class ChatMessageResource(admin.ModelAdmin):

django_app/redbox_app/redbox_core/consumers.py

@@ -18,8 +18,9 @@
 
 
 class ChatConsumer(AsyncWebsocketConsumer):
-    async def receive(self, text_data):
-        data = json.loads(text_data)
+    async def receive(self, text_data=None, bytes_data=None):
+        data = json.loads(text_data or bytes_data)
+        logger.debug("received %s from browser", data)
         user_message_text: str = data.get("message", "")
         session_id: str | None = data.get("sessionId", None)
         selected_file_uuids: Sequence[UUID] = [UUID(u) for u in data.get("selectedFiles", [])]
@@ -32,6 +33,7 @@ async def receive(self, text_data):
         await self.save_message(session, user_message_text, ChatRoleEnum.user, selected_files=selected_files)
 
         await self.llm_conversation(selected_files, session, user)
+        await self.close()
 
     async def llm_conversation(self, selected_files: Sequence[File], session: ChatHistory, user: User) -> None:
         session_messages = await self.get_messages(session)
@@ -44,7 +46,6 @@ async def llm_conversation(self, selected_files: Sequence[File], session: ChatHi
                 "message_history": message_history,
                 "selected_files": [{"uuid": f.core_file_uuid} for f in selected_files],
             }
-            logger.debug("sending to core-api: %s", message)
             await self.send_to_server(core_websocket, message)
             await self.send_to_client({"type": "session-id", "data": str(session.id)})
             reply, source_files = await self.receive_llm_responses(user, core_websocket)
@@ -57,7 +58,7 @@ async def receive_llm_responses(
         source_files: MutableSequence[File] = []
         async for raw_message in core_websocket:
             message = json.loads(raw_message, object_hook=lambda d: SimpleNamespace(**d))
-            logger.debug("Received: %s", message)
+            logger.debug("received %s from core-api", message)
             if message.resource_type == "text":
                 full_reply.append(await self.handle_text(message))
             elif message.resource_type == "documents":
@@ -81,10 +82,12 @@ async def handle_text(self, message: SimpleNamespace) -> str:
         return message.data
 
     async def send_to_client(self, data):
+        logger.debug("sending %s to browser", data)
         await self.send(json.dumps(data, default=str))
 
     @staticmethod
     async def send_to_server(websocket, data):
+        logger.debug("sending %s to core-api", data)
         return await websocket.send(json.dumps(data, default=str))
 
     @staticmethod

django_app/redbox_app/settings.py

@@ -16,7 +16,7 @@
 
 SECRET_KEY = env.str("DJANGO_SECRET_KEY")
 ENVIRONMENT = Environment[env.str("ENVIRONMENT").upper()]
-WEBSOCKET_SCHEME = "ws" if ENVIRONMENT.is_local else "wss"
+WEBSOCKET_SCHEME = "ws" if ENVIRONMENT.is_test else "wss"
 
 # SECURITY WARNING: don't run with debug turned on in production!
 DEBUG = env.bool("DEBUG")

django_app/tests/conftest.py

@@ -106,10 +106,10 @@ def chat_history_with_files(chat_history: ChatHistory, several_files: Sequence[F
     chat_message = ChatMessage.objects.create(chat_history=chat_history, text="An answer.", role=ChatRoleEnum.ai)
     chat_message.source_files.set(several_files[0::2])
     chat_message = ChatMessage.objects.create(
-        chat_history=chat_history, text="Another question?", role=ChatRoleEnum.user
+        chat_history=chat_history, text="A second question?", role=ChatRoleEnum.user
     )
     chat_message.selected_files.set(several_files[0:2])
-    chat_message = ChatMessage.objects.create(chat_history=chat_history, text="Another answer.", role=ChatRoleEnum.ai)
+    chat_message = ChatMessage.objects.create(chat_history=chat_history, text="A second answer.", role=ChatRoleEnum.ai)
     chat_message.source_files.set([several_files[2]])
     return chat_history
 

django_app/tests/test_consumers.py

@@ -78,6 +78,7 @@ def get_chat_message_text(user: User, role: ChatRoleEnum) -> Sequence[str]:
     return [m.text for m in ChatMessage.objects.filter(chat_history__users=user, role=role)]
 
 
+@pytest.mark.xfail()
 @pytest.mark.django_db(transaction=True)
 @pytest.mark.asyncio()
 async def test_chat_consumer_with_selected_files(
@@ -87,6 +88,7 @@ async def test_chat_consumer_with_selected_files(
     mocked_connect_with_several_files: Connect,
 ):
     # Given
+    selected_files: Sequence[File] = several_files[2:]
 
     # When
     with patch("redbox_app.redbox_core.consumers.connect", new=mocked_connect_with_several_files):
@@ -95,11 +97,12 @@ async def test_chat_consumer_with_selected_files(
         connected, _ = await communicator.connect()
         assert connected
 
+        selected_file_core_uuids: Sequence[str] = [str(f.core_file_uuid) for f in selected_files]
         await communicator.send_json_to(
             {
                 "message": "Third question, with selected files?",
                 "sessionId": str(chat_history_with_files.id),
-                "selectedFiles": [str(f.core_file_uuid) for f in several_files[2:]],
+                "selectedFiles": selected_file_core_uuids,
             }
         )
         response1 = await communicator.receive_json_from(timeout=5)
@@ -111,13 +114,41 @@ async def test_chat_consumer_with_selected_files(
         # Close
         await communicator.disconnect()
 
-    assert await get_chat_message_text(alice, ChatRoleEnum.user) == [
-        "A question?",
-        "Another question?",
-        "Third question, with selected files?",
-    ]
-    assert await get_chat_message_text(alice, ChatRoleEnum.ai) == ["An answer.", "Another answer.", "Third answer."]
-    # TODO (@brunns): Assert selected files sent to core, and saved to model.
+    # Then
+
+    # TODO (@brunns): Assert selected files sent to core.
+    # Requires fix for https://github.com/django/channels/issues/1091
+    mocked_websocket = mocked_connect_with_several_files.return_value.__aenter__.return_value
+    expected = json.dumps(
+        {
+            "message_history": [
+                {"role": "user", "text": "A question?"},
+                {"role": "ai", "text": "An answer."},
+                {"role": "user", "text": "A second question?"},
+                {"role": "ai", "text": "A second answer."},
+                {"role": "user", "text": "Third question, with selected files?"},
+            ],
+            "selected_files": selected_file_core_uuids,
+        }
+    )
+    mocked_websocket.send.assert_called_with(expected)
+
+    # TODO (@brunns): Assert selected files saved to model.
+    # Requires fix for https://github.com/django/channels/issues/1091
+    all_messages = get_chat_messages(alice)
+    last_user_message = [m for m in all_messages if m.rule == ChatRoleEnum.user][-1]
+    assert last_user_message.selected_files == selected_files
+
+
+@database_sync_to_async
+def get_chat_messages(user: User) -> Sequence[ChatMessage]:
+    return list(
+        ChatMessage.objects.filter(chat_history__users=user)
+        .order_by("created_at")
+        .prefetch_related("chat_history")
+        .prefetch_related("source_files")
+        .prefetch_related("selected_files")
+    )
 
 
 @pytest.fixture()