Adding more tests for prepdocs

Author: pamelafox

app/backend/prepdocs.py

@@ -158,7 +158,7 @@ def setup_file_processors(
 ):
     html_parser: Parser
     pdf_parser: Parser
-    doc_int_parser: DocumentAnalysisParser
+    doc_int_parser: DocumentAnalysisParser = None
 
     # check if Azure Document Intelligence credentials are provided
     if document_intelligence_service is not None:
@@ -178,23 +178,33 @@ def setup_file_processors(
     else:
         html_parser = doc_int_parser
     sentence_text_splitter = SentenceTextSplitter(has_image_embeddings=search_images)
-    return {
+
+    # These file formats can always be parsed, thanks to local packages
+    file_processors = {
         ".pdf": FileProcessor(pdf_parser, sentence_text_splitter),
         ".html": FileProcessor(html_parser, sentence_text_splitter),
         ".json": FileProcessor(JsonParser(), SimpleTextSplitter()),
-        ".docx": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".pptx": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".xlsx": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".png": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".jpg": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".jpeg": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".tiff": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".bmp": FileProcessor(doc_int_parser, sentence_text_splitter),
-        ".heic": FileProcessor(doc_int_parser, sentence_text_splitter),
         ".md": FileProcessor(TextParser(), sentence_text_splitter),
         ".txt": FileProcessor(TextParser(), sentence_text_splitter),
     }
 
+    # These file formats require Document Intelligence
+    if doc_int_parser is not None:
+        file_processors.update(
+            {
+                ".docx": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".pptx": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".xlsx": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".png": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".jpg": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".jpeg": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".tiff": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".bmp": FileProcessor(doc_int_parser, sentence_text_splitter),
+                ".heic": FileProcessor(doc_int_parser, sentence_text_splitter),
+            }
+        )
+    return file_processors
+
 
 def setup_image_embeddings_service(
     azure_credential: AsyncTokenCredential, vision_endpoint: Union[str, None], search_images: bool

tests/test_app_config.py

@@ -2,6 +2,7 @@
 from unittest import mock
 
 import pytest
+import quart
 
 import app
 
@@ -52,6 +53,65 @@ async def test_app_azure_custom_identity(monkeypatch, minimal_env):
         assert quart_app.config[app.CONFIG_OPENAI_CLIENT].base_url == "http://azureapi.com/api/v1/openai/"
 
 
+@pytest.mark.asyncio
+async def test_app_user_upload_processors(monkeypatch, minimal_env):
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+
+    quart_app = app.create_app()
+    async with quart_app.test_app():
+        ingester = quart_app.config[app.CONFIG_INGESTER]
+        assert ingester is not None
+        assert len(ingester.file_processors.keys()) == 5
+
+
+@pytest.mark.asyncio
+async def test_app_user_upload_processors_docint(monkeypatch, minimal_env):
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("AZURE_DOCUMENTINTELLIGENCE_SERVICE", "test-docint-service")
+
+    quart_app = app.create_app()
+    async with quart_app.test_app():
+        ingester = quart_app.config[app.CONFIG_INGESTER]
+        assert ingester is not None
+        assert len(ingester.file_processors.keys()) == 14
+
+
+@pytest.mark.asyncio
+async def test_app_user_upload_processors_docint_localpdf(monkeypatch, minimal_env):
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("AZURE_DOCUMENTINTELLIGENCE_SERVICE", "test-docint-service")
+    monkeypatch.setenv("USE_LOCAL_PDF_PARSER", "true")
+
+    quart_app = app.create_app()
+    async with quart_app.test_app():
+        ingester = quart_app.config[app.CONFIG_INGESTER]
+        assert ingester is not None
+        assert len(ingester.file_processors.keys()) == 14
+        assert ingester.file_processors[".pdf"] is not ingester.file_processors[".pptx"]
+
+
+@pytest.mark.asyncio
+async def test_app_user_upload_processors_docint_localhtml(monkeypatch, minimal_env):
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("AZURE_DOCUMENTINTELLIGENCE_SERVICE", "test-docint-service")
+    monkeypatch.setenv("USE_LOCAL_HTML_PARSER", "true")
+
+    quart_app = app.create_app()
+    async with quart_app.test_app():
+        ingester = quart_app.config[app.CONFIG_INGESTER]
+        assert ingester is not None
+        assert len(ingester.file_processors.keys()) == 14
+        assert ingester.file_processors[".html"] is not ingester.file_processors[".pptx"]
+
+
 @pytest.mark.asyncio
 async def test_app_config_default(monkeypatch, minimal_env):
     quart_app = app.create_app()
@@ -105,6 +165,7 @@ async def test_app_config_semanticranker_free(monkeypatch, minimal_env):
         assert result["showGPT4VOptions"] is False
         assert result["showSemanticRankerOption"] is True
         assert result["showVectorOption"] is True
+        assert result["showUserUpload"] is False
 
 
 @pytest.mark.asyncio
@@ -119,6 +180,75 @@ async def test_app_config_semanticranker_disabled(monkeypatch, minimal_env):
         assert result["showGPT4VOptions"] is False
         assert result["showSemanticRankerOption"] is False
         assert result["showVectorOption"] is True
+        assert result["showUserUpload"] is False
+
+
+@pytest.mark.asyncio
+async def test_app_config_user_upload(monkeypatch, minimal_env):
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    quart_app = app.create_app()
+    async with quart_app.test_app() as test_app:
+        client = test_app.test_client()
+        response = await client.get("/config")
+        assert response.status_code == 200
+        result = await response.get_json()
+        assert result["showGPT4VOptions"] is False
+        assert result["showSemanticRankerOption"] is True
+        assert result["showVectorOption"] is True
+        assert result["showUserUpload"] is True
+
+
+@pytest.mark.asyncio
+async def test_app_config_user_upload_novectors(monkeypatch, minimal_env):
+    """Check that this combo works correctly with prepdocs.py embedding service."""
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("USE_VECTORS", "false")
+    quart_app = app.create_app()
+    async with quart_app.test_app() as test_app:
+        client = test_app.test_client()
+        response = await client.get("/config")
+        assert response.status_code == 200
+        result = await response.get_json()
+        assert result["showGPT4VOptions"] is False
+        assert result["showSemanticRankerOption"] is True
+        assert result["showVectorOption"] is False
+        assert result["showUserUpload"] is True
+
+
+@pytest.mark.asyncio
+async def test_app_config_user_upload_bad_openai_config(monkeypatch, minimal_env):
+    """Check that this combo works correctly with prepdocs.py embedding service."""
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("OPENAI_HOST", "openai")
+    quart_app = app.create_app()
+    with pytest.raises(
+        quart.testing.app.LifespanError, match="OpenAI key is required when using the non-Azure OpenAI API"
+    ):
+        async with quart_app.test_app() as test_app:
+            test_app.test_client()
+
+
+@pytest.mark.asyncio
+async def test_app_config_user_upload_openaicom(monkeypatch, minimal_env):
+    """Check that this combo works correctly with prepdocs.py embedding service."""
+    monkeypatch.setenv("AZURE_USERSTORAGE_ACCOUNT", "test-user-storage-account")
+    monkeypatch.setenv("AZURE_USERSTORAGE_CONTAINER", "test-user-storage-container")
+    monkeypatch.setenv("USE_USER_UPLOAD", "true")
+    monkeypatch.setenv("OPENAI_HOST", "openai")
+    monkeypatch.setenv("OPENAI_API_KEY", "pretendkey")
+    quart_app = app.create_app()
+    async with quart_app.test_app() as test_app:
+        client = test_app.test_client()
+        response = await client.get("/config")
+        assert response.status_code == 200
+        result = await response.get_json()
+        assert result["showUserUpload"] is True
 
 
 @pytest.mark.asyncio