feat: add support for Azure Document Intelligence in ETL pipeline

surfsense_backend/app/config/__init__.py

@@ -397,6 +397,10 @@ class Config:
         # LlamaCloud API Key
         LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
 
+    elif ETL_SERVICE == "AZURE_DI":
+        AZURE_DI_ENDPOINT = os.getenv("AZURE_DI_ENDPOINT")
+        AZURE_DI_KEY = os.getenv("AZURE_DI_KEY")
+
     # Residential Proxy Configuration (anonymous-proxies.net)
     # Used for web crawling and YouTube transcript fetching to avoid IP bans.
     RESIDENTIAL_PROXY_USERNAME = os.getenv("RESIDENTIAL_PROXY_USERNAME")

surfsense_backend/app/etl_pipeline/etl_pipeline_service.py

@@ -56,7 +56,7 @@ class EtlPipelineService:
         if not etl_service:
             raise EtlServiceUnavailableError(
                 "No ETL_SERVICE configured. "
-                "Set ETL_SERVICE to UNSTRUCTURED, LLAMACLOUD, or DOCLING in your .env"
+                "Set ETL_SERVICE to UNSTRUCTURED, LLAMACLOUD, DOCLING, or AZURE_DI in your .env"
             )
 
         ext = PurePosixPath(request.filename).suffix.lower()
@@ -80,6 +80,12 @@ class EtlPipelineService:
             content = await parse_with_llamacloud(
                 request.file_path, request.estimated_pages
             )
+        elif etl_service == "AZURE_DI":
+            from app.etl_pipeline.parsers.azure_doc_intelligence import (
+                parse_with_azure_doc_intelligence,
+            )
+
+            content = await parse_with_azure_doc_intelligence(request.file_path)
         else:
             raise EtlServiceUnavailableError(f"Unknown ETL_SERVICE: {etl_service}")
 

surfsense_backend/app/etl_pipeline/parsers/azure_doc_intelligence.py

@@ -0,0 +1,93 @@
+import asyncio
+import logging
+import os
+import random
+
+from app.config import config as app_config
+
+MAX_RETRIES = 5
+BASE_DELAY = 10
+MAX_DELAY = 120
+
+
+async def parse_with_azure_doc_intelligence(file_path: str) -> str:
+    from azure.ai.documentintelligence.aio import DocumentIntelligenceClient
+    from azure.ai.documentintelligence.models import DocumentContentFormat
+    from azure.core.credentials import AzureKeyCredential
+    from azure.core.exceptions import (
+        ClientAuthenticationError,
+        HttpResponseError,
+        ServiceRequestError,
+        ServiceResponseError,
+    )
+
+    file_size_mb = os.path.getsize(file_path) / (1024 * 1024)
+    retryable_exceptions = (ServiceRequestError, ServiceResponseError)
+
+    last_exception = None
+    attempt_errors: list[str] = []
+
+    for attempt in range(1, MAX_RETRIES + 1):
+        try:
+            client = DocumentIntelligenceClient(
+                endpoint=app_config.AZURE_DI_ENDPOINT,
+                credential=AzureKeyCredential(app_config.AZURE_DI_KEY),
+            )
+            async with client:
+                with open(file_path, "rb") as f:
+                    poller = await client.begin_analyze_document(
+                        "prebuilt-read",
+                        body=f,
+                        output_content_format=DocumentContentFormat.MARKDOWN,
+                    )
+                result = await poller.result()
+
+            if attempt > 1:
+                logging.info(
+                    f"Azure Document Intelligence succeeded on attempt {attempt} "
+                    f"after {len(attempt_errors)} failures"
+                )
+
+            if not result.content:
+                return ""
+
+            return result.content
+
+        except ClientAuthenticationError:
+            raise
+        except HttpResponseError as e:
+            if e.status_code and 400 <= e.status_code < 500:
+                raise
+            last_exception = e
+            error_type = type(e).__name__
+            error_msg = str(e)[:200]
+            attempt_errors.append(f"Attempt {attempt}: {error_type} - {error_msg}")
+        except retryable_exceptions as e:
+            last_exception = e
+            error_type = type(e).__name__
+            error_msg = str(e)[:200]
+            attempt_errors.append(f"Attempt {attempt}: {error_type} - {error_msg}")
+
+        if attempt < MAX_RETRIES:
+            base_delay = min(BASE_DELAY * (2 ** (attempt - 1)), MAX_DELAY)
+            jitter = base_delay * 0.25 * (2 * random.random() - 1)
+            delay = base_delay + jitter
+
+            logging.warning(
+                f"Azure Document Intelligence failed "
+                f"(attempt {attempt}/{MAX_RETRIES}): "
+                f"{attempt_errors[-1]}. File: {file_size_mb:.1f}MB. "
+                f"Retrying in {delay:.0f}s..."
+            )
+            await asyncio.sleep(delay)
+        else:
+            logging.error(
+                f"Azure Document Intelligence failed after {MAX_RETRIES} "
+                f"attempts. File size: {file_size_mb:.1f}MB. "
+                f"Errors: {'; '.join(attempt_errors)}"
+            )
+
+    raise last_exception or RuntimeError(
+        f"Azure Document Intelligence parsing failed after {MAX_RETRIES} retries. "
+        f"File size: {file_size_mb:.1f}MB"
+    )

surfsense_backend/app/utils/file_extensions.py

@@ -93,6 +93,22 @@ UNSTRUCTURED_DOCUMENT_EXTENSIONS: frozenset[str] = frozenset(
     }
 )
 
+AZURE_DI_DOCUMENT_EXTENSIONS: frozenset[str] = frozenset(
+    {
+        ".pdf",
+        ".docx",
+        ".xlsx",
+        ".pptx",
+        ".png",
+        ".jpg",
+        ".jpeg",
+        ".bmp",
+        ".tiff",
+        ".tif",
+        ".heif",
+    }
+)
+
 # ---------------------------------------------------------------------------
 # Union (used by classify_file for routing) + service lookup
 # ---------------------------------------------------------------------------
@@ -101,12 +117,14 @@ DOCUMENT_EXTENSIONS: frozenset[str] = (
     DOCLING_DOCUMENT_EXTENSIONS
     | LLAMAPARSE_DOCUMENT_EXTENSIONS
     | UNSTRUCTURED_DOCUMENT_EXTENSIONS
+    | AZURE_DI_DOCUMENT_EXTENSIONS
 )
 
 _SERVICE_MAP: dict[str, frozenset[str]] = {
     "DOCLING": DOCLING_DOCUMENT_EXTENSIONS,
     "LLAMACLOUD": LLAMAPARSE_DOCUMENT_EXTENSIONS,
     "UNSTRUCTURED": UNSTRUCTURED_DOCUMENT_EXTENSIONS,
+    "AZURE_DI": AZURE_DI_DOCUMENT_EXTENSIONS,
 }