WIP✨(backend) albert AI client & pdf conversion

Add AlbertAI client to wrap embedding & conversion API calls
Implement working pdf to markdown converter using Albert
Add unstructured dependencies

Signed-off-by: Fabre Florian <ffabre@hybird.org>

Dockerfile

@@ -63,8 +63,11 @@ RUN apt-get update && \
       libcairo2 \
       libffi-dev \
       libgdk-pixbuf2.0-0 \
+      libmagic-dev \
       libpango-1.0-0 \
       libpangocairo-1.0-0 \
+      poppler-utils \
+      pandoc \
       shared-mime-info && \
   rm -rf /var/lib/apt/lists/*
 

src/backend/core/services/albert.py

@@ -0,0 +1,101 @@
+"""Albert AI related tools"""
+
+import logging
+from io import BytesIO, StringIO
+
+from django.conf import settings
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+class AlbertAIError(Exception):
+    """Albert AI errors"""
+
+    def __init__(self, message):
+        super().__init__(message)
+        self.message = message
+
+
+class AlbertAI:
+    """
+    Client for Albert AI API
+    https://albert.api.etalab.gouv.fr/swagger#/
+    """
+
+    def __init__(self):
+        self.api_key = settings.EMBEDDING_API_KEY
+        self.timeout = settings.EMBEDDING_REQUEST_TIMEOUT
+        self.doc_parse_url = settings.ALBERT_PARSE_ENDPOINT
+        self.embedding_url = settings.EMBEDDING_API_PATH
+
+    def _request_api(self, url, **kwargs):
+        """Make authenticated api call"""
+        try:
+            response = requests.post(
+                url,
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                timeout=self.timeout,
+                **kwargs,
+            )
+            response.raise_for_status()
+            return response
+        except requests.HTTPError as e:
+            raise AlbertAIError(e.response.json().get("detail", str(e))) from e
+
+    def embedding(self, text, dimensions=None, model=None):
+        """
+        Get embedding vector for the given text from Albert OpenAI-compatible embedding API
+        """
+        dimensions = dimensions or settings.EMBEDDING_DIMENSION
+        model = model or settings.EMBEDDING_API_MODEL_NAME
+
+        response = self._request_api(
+            self.embedding_url,
+            json={
+                "input": text,
+                "model": model,
+                "dimensions": dimensions,
+                "encoding_format": "float",
+            },
+        )
+
+        try:
+            return response.json()["data"][0]["embedding"]
+        except Exception as e:
+            raise AlbertAIError(f"Unexpected content : {response.text}") from e
+
+    # pylint: disable=too-many-arguments, too-many-positional-arguments
+    def convert(
+        self,
+        content,
+        mimetype="application/pdf",
+        pages=None,
+        output="markdown",
+        encoding="utf-8",
+    ):  # noqa : PLR0913
+        """
+        Convert the content (only pdf) to markdown, json or html using the Albert API
+        """
+        if isinstance(content, str):
+            content = StringIO(content.encode(encoding))
+        elif isinstance(content, bytes):
+            content = BytesIO(content)
+
+        response = self._request_api(
+            self.doc_parse_url,
+            files={
+                "file": ("input", content, mimetype),
+            },
+            data={
+                "output_format": output,
+                "page_range": f"0-{pages}" if pages else None,
+            },
+        )
+
+        try:
+            data = response.json()["data"]
+            return "\n".join([page["content"] for page in data])
+        except Exception as e:
+            raise AlbertAIError(f"Unexpected content : {response.text}") from e

src/backend/core/services/converters.py

@@ -1,6 +1,10 @@
 """Document content conversion tools"""
 
+from io import BytesIO
 
-def pdf_to_markdown(content):
+from .albert import AlbertAI
+
+
+def pdf_to_markdown(content: BytesIO):
     """Convert PDF stream into markdown"""
-    return content.read()
+    return AlbertAI().convert(content=content, mimetype="application/pdf")

src/backend/core/services/indexer_services.py

@@ -3,7 +3,7 @@
 import logging
 from contextlib import contextmanager
 from dataclasses import asdict as dataasdict
-from io import StringIO
+from io import BytesIO
 
 from django.conf import settings
 
@@ -222,6 +222,7 @@ class IndexerTaskService:
     and embedding
     """
 
+    # V2 : Use settings to define the list of converters
     converters = {"application/pdf": pdf_to_markdown}
 
     def __init__(
@@ -251,7 +252,7 @@ class IndexerTaskService:
     def process_content(self, document, content):
         """Transforms the document file data into an indexable format"""
         try:
-            stream = StringIO(content) if isinstance(content, str) else content
+            stream = BytesIO(content.encode()) if isinstance(content, str) else content
             converter = self.get_converter(document.mimetype)
             output = converter(stream) if converter is not None else stream.read()
             return output.decode() if isinstance(output, bytes) else output
@@ -335,7 +336,7 @@ class IndexerTaskService:
             ) as actions:
                 for doc in docs:
                     try:
-                        content = self.process_content(doc, StringIO(doc.content))
+                        content = self.process_content(doc, doc.content)
 
                         actions.update(
                             doc.id,
@@ -493,7 +494,7 @@ class IndexerTaskService:
                     doc.content_status = enums.ContentStatusEnum.WAIT
                 elif not is_allowed_mimetype(doc.mimetype, INDEXABLE_MIMETYPES):
                     try:
-                        doc.content = self.process_content(doc, StringIO(doc.content))
+                        doc.content = self.process_content(doc, doc.content)
                         doc.content_status = enums.ContentStatusEnum.READY
                     except IndexContentError as err:
                         errors.append(IndexBulkError(str(err), _id=doc.id))

src/backend/demo/management/commands/albert.py

@@ -0,0 +1,58 @@
+"""Helper command to test Albert AI features"""
+
+from mimetypes import guess_type
+from pathlib import Path
+
+from django.core.management.base import BaseCommand, CommandError
+
+from core.services.albert import AlbertAI, AlbertAIError
+
+
+class Command(BaseCommand):
+    """A management command to test Albert AI features"""
+
+    help = __doc__
+
+    def add_arguments(self, parser):
+        """Command arguments."""
+        subparsers = parser.add_subparsers(help="sub-command help", dest="action")
+
+        parse = subparsers.add_parser("parse", help="parse help")
+        parse.add_argument(nargs="+", dest="files")
+        parse.add_argument(
+            "-f", "--format", dest="format", default="markdown", help="output format"
+        )
+        parse.add_argument(
+            "-p", "--pages", dest="page", default="", help="extracted pages"
+        )
+
+    def handle(self, *args, **options):
+        """Handling of the management command."""
+        action = options.get("action")
+
+        try:
+            handler = getattr(self, f"handle_{action}")
+        except AttributeError:
+            self.print_help("albert", action)
+            return
+
+        handler(options)
+
+    def handle_parse(self, options):
+        """Handling of the file convertion using Albert AI (only pdf)"""
+        paths = [Path(p) for p in options.get("files", [])]
+        albert = AlbertAI()
+
+        for path in paths:
+            with open(path, "rb") as fd:
+                try:
+                    self.stdout.write(
+                        albert.convert(
+                            content=fd,
+                            mimetype=guess_type(path)[0],
+                            output=options.get("format"),
+                            pages=options.get("pages"),
+                        )
+                    )
+                except AlbertAIError as e:
+                    raise CommandError(f"Unable to convert {path} : {e.message}") from e

src/backend/find/settings.py

@@ -258,6 +258,10 @@ class Base(Configuration):
         default=1, environ_name="INDEXER_TASK_COUNTDOWN", environ_prefix=None
     )
 
+    ALBERT_PARSE_ENDPOINT = values.Value(
+        default=None, environ_name="ALBERT_PARSE_ENDPOINT", environ_prefix=None
+    )
+
     SPECTACULAR_SETTINGS = {
         "TITLE": "Find API",
         "DESCRIPTION": "This is the find API schema.",

src/backend/pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
     "url-normalize==1.4.3",
     "opensearch-py==2.8.0",
     "whitenoise==6.8.2",
+    "unstructured[doc,docx,csv,odt,pptx,xlsx]==0.18.21",
 ]
 
 [project.urls]