diff --git a/Dockerfile b/Dockerfile index b898def..d664ec7 100644 --- a/Dockerfile +++ b/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/* diff --git a/src/backend/core/services/albert.py b/src/backend/core/services/albert.py new file mode 100644 index 0000000..d92e474 --- /dev/null +++ b/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 diff --git a/src/backend/core/services/converters.py b/src/backend/core/services/converters.py index 39452d7..2c5fc66 100644 --- a/src/backend/core/services/converters.py +++ b/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") diff --git a/src/backend/core/services/indexer_services.py b/src/backend/core/services/indexer_services.py index 57f7ab7..636a528 100644 --- a/src/backend/core/services/indexer_services.py +++ b/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)) diff --git a/src/backend/demo/management/commands/albert.py b/src/backend/demo/management/commands/albert.py new file mode 100644 index 0000000..c874c03 --- /dev/null +++ b/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 diff --git a/src/backend/find/settings.py b/src/backend/find/settings.py index f47c47d..093d456 100755 --- a/src/backend/find/settings.py +++ b/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.", diff --git a/src/backend/pyproject.toml b/src/backend/pyproject.toml index 6a35e1c..f2e5c82 100644 --- a/src/backend/pyproject.toml +++ b/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]