electron-rare/find
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

✨(backend) albert AI client & pdf conversion

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

Signed-off-by: Fabre Florian <ffabre@hybird.org>
This commit is contained in:
Fabre Florian
2025-11-25 16:10:11 +01:00
parent 7e5fc7c138
commit 56f0fb8df3
11 changed files with 499 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/backend/core/schemas.py
+1 -1
View File
@@ -7,8 +7,8 @@ from django.utils.text import slugify
from pydantic import (
UUID4,
AwareDatetime,
AnyUrl,
AwareDatetime,
BaseModel,
BeforeValidator,
ConfigDict,
src/backend/core/services/albert.py
+107
View File
@@ -0,0 +1,107 @@
"""Albert AI related tools"""
import logging
from io import BytesIO
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 # V2 : Rename as ALBERT_API_KEY
self.timeout = (
settings.EMBEDDING_REQUEST_TIMEOUT
) # V2 : Rename as ALBERT_REQUEST_TIMEOUT
self.doc_parse_url = settings.ALBERT_PARSE_ENDPOINT
self.embedding_url = (
settings.EMBEDDING_API_PATH
) # V2 : Rename as ALBERT_EMBEDDING_ENDPOINT
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
except requests.RequestException as e:
raise AlbertAIError(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",
):
"""
Convert the content (only pdf) to markdown, json or html using the Albert API
"""
if isinstance(content, str):
content = BytesIO(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
+6 -2
View File
@@ -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
+40 -18
View File
@@ -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
@@ -12,10 +12,10 @@ import requests
from core import enums
from core.models import IndexDocument, Service
from .converters import pdf_to_markdown
from . import converters
from .albert import AlbertAI, AlbertAIError
from .opensearch import (
check_hybrid_search_enabled,
embed_text,
ensure_index_exists,
format_document,
opensearch_client,
@@ -222,7 +222,8 @@ class IndexerTaskService:
and embedding
"""
converters = {"application/pdf": pdf_to_markdown}
# V2 : Use settings to define the list of converters
converters = {"application/pdf": converters.pdf_to_markdown}
def __init__(
self, service: Service, batch_size=100, client=None, force_refresh=None
@@ -251,9 +252,18 @@ 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
converter = self.get_converter(document.mimetype)
output = converter(stream) if converter is not None else stream.read()
if converter is not None:
# A converter only accepts a BytesIO
stream = (
BytesIO(content.encode()) if isinstance(content, str) else content
)
output = converter(stream)
else:
# no need to convert
output = content if isinstance(content, str) else content.read()
return output.decode() if isinstance(output, bytes) else output
except IndexContentError as e:
e.id = document.id
@@ -304,7 +314,7 @@ class IndexerTaskService:
for hit in self.search(query, sort=sort, batch_size=batch_size):
yield hit_to_doc(hit)
def preprocess_all(self, batch_size=None):
def process_all(self, batch_size=None):
"""
Gets all the documents in waiting status to load and convert their content
Returns an error dict.
@@ -335,7 +345,8 @@ class IndexerTaskService:
) as actions:
for doc in docs:
try:
content = self.process_content(doc, StringIO(doc.content))
# V2 : Use asyncio loop to parallelize conversion
content = self.process_content(doc, doc.content)
actions.update(
doc.id,
@@ -355,7 +366,7 @@ class IndexerTaskService:
return errors
def load_all(self, batch_size=None):
def load_n_process_all(self, batch_size=None):
"""
Gets all the documents in waiting status to load and convert their content
Returns an error dict.
@@ -420,6 +431,7 @@ class IndexerTaskService:
index_name = self.service.index_name
batch_size = batch_size or self.batch_size
model_name = model_name or settings.EMBEDDING_API_MODEL_NAME
albert = AlbertAI()
doc_batches = self.search_documents(
query={
"bool": {
@@ -452,13 +464,19 @@ class IndexerTaskService:
) as actions:
for doc in docs:
try:
embedding = embed_text(format_document(doc.title, doc.content))
if embedding is None:
raise IndexerError(
"Unable to build embedding for the document", _id=doc.id
# V2 : Use asyncio loop to parallelize embedding
embedding = albert.embedding(
text=format_document(doc.title, doc.content),
model=model_name,
)
except AlbertAIError as e:
errors.append(
IndexerError(
f"Unable to build embedding for the document : {e.message}",
_id=doc.id,
)
)
else:
actions.update(
doc.id,
data={
@@ -470,8 +488,6 @@ class IndexerTaskService:
if_seq_no=doc.hit["_seq_no"],
if_primary_term=doc.hit["_primary_term"],
)
except IndexerError as e:
errors.append(e)
errors.extend(actions.errors())
@@ -490,14 +506,20 @@ class IndexerTaskService:
) as actions:
for doc in documents:
if doc.content_uri and not doc.content:
# Without content and a dowload uri : set WAIT status
doc.content_status = enums.ContentStatusEnum.WAIT
elif not is_allowed_mimetype(doc.mimetype, INDEXABLE_MIMETYPES):
# A content but not directly indexable (e.g xml or html content) : process them
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:
# If process has failed, set LOADED status for a retry.
# V2 : Add retry mechanism ?
doc.content_status = enums.ContentStatusEnum.LOADED
errors.append(IndexBulkError(str(err), _id=doc.id))
else:
# The content exists and is indexable : set READY
doc.content_status = enums.ContentStatusEnum.READY
actions.index_document(doc)
src/backend/core/services/opensearch.py
+1
View File
@@ -230,6 +230,7 @@ def format_document(title, content):
def embed_text(text):
"""
Get embedding vector for the given text from any OpenAI-compatible embedding API
V2 : Deprecated, use AlbertAI.embedding() instead.
"""
response = requests.post(
settings.EMBEDDING_API_PATH,
src/backend/core/tasks/indexer.py
+18 -24
View File
@@ -24,35 +24,27 @@ def get_service(service_id):
@app.task
def loading_task(service_id):
def load_n_process_task(service_id):
"""Celery Task : Re-index documents with deferred loading."""
service = get_service(service_id)
if service is not None:
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
logger.info("Start deferred loading on index %s", service.index_name)
indexer.load_all()
# Trigger the embedding task if enabled
if check_hybrid_search_enabled():
embedding_task.apply_async((service_id,))
indexer.load_n_process_all()
@app.task
def preprocess_task(service_id):
def process_task(service_id):
"""Celery Task : Re-index documents with deferred loading."""
service = get_service(service_id)
if service is not None:
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
logger.info("Start deferred preprocessing on index %s", service.index_name)
indexer.preprocess_all()
# Trigger the embedding task if enabled
if check_hybrid_search_enabled():
embedding_task.apply_async((service_id,))
indexer.process_all()
@app.task
@@ -61,7 +53,7 @@ def embedding_task(service_id):
service = get_service(service_id)
if service is not None:
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
logger.info("Start embedding on index %s", service.index_name)
indexer.embed_all()
@@ -72,18 +64,20 @@ def dispatch_indexing_tasks(service, documents: List[models.IndexDocument]):
Trigger task related to the different status of the documents
"""
countdown = settings.INDEXER_TASK_COUNTDOWN
waiting = any(doc.is_waiting for doc in documents)
not_ready = any(doc.is_loaded for doc in documents)
not_embed = check_hybrid_search_enabled() and any(
should_load = any(doc.is_waiting for doc in documents)
should_preprocess = any(doc.is_loaded for doc in documents)
should_embed = check_hybrid_search_enabled() and any(
doc.is_ready and not doc.embedding for doc in documents
)
# Trigger tasks for deferred loading
if waiting:
loading_task.apply_async((service.pk,), countdown=countdown)
# Trigger task for deferred loading if the file is too big
if should_load:
load_n_process_task.apply_async((service.pk,), countdown=countdown)
if not_ready:
preprocess_task.apply_async((service.pk,), countdown=countdown)
# Trigger task for deferred preprocessing of the content (picture analysis for instance)
if should_preprocess:
process_task.apply_async((service.pk,), countdown=countdown)
if not_embed:
# Trigger task for semantic indexation
if should_embed:
embedding_task.apply_async((service.pk,), countdown=countdown)
src/backend/core/tests/test_albert_services.py
+216
View File
@@ -0,0 +1,216 @@
"""Tests albert AI service"""
from io import BytesIO
from json import dumps as json_dumps
from unittest import mock
import pytest
import responses
from core.services import albert
from core.tests.mock import albert_embedding_response
from core.tests.utils import (
enable_hybrid_search,
)
pytestmark = pytest.mark.django_db
@responses.activate
def test_albert_service_embedding(settings):
"""Should return the embedding vector from Albert API"""
enable_hybrid_search(settings)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
json=albert_embedding_response.response,
status=200,
)
assert (
albert.AlbertAI().embedding("any text")
== (albert_embedding_response.response["data"][0]["embedding"])
)
responses.assert_call_count(settings.EMBEDDING_API_PATH, 1)
assert (
responses.calls[0].request.body
== json_dumps(
{
"input": "any text",
"model": settings.EMBEDDING_API_MODEL_NAME,
"dimensions": settings.EMBEDDING_DIMENSION,
"encoding_format": "float",
}
).encode()
)
@responses.activate
def test_albert_service_embedding__arguments(settings):
"""Should return the embedding vector from Albert API with custom arguments"""
enable_hybrid_search(settings)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
json=albert_embedding_response.response,
status=200,
)
assert (
albert.AlbertAI().embedding("any text", dimensions=123, model="mymodel")
== (albert_embedding_response.response["data"][0]["embedding"])
)
responses.assert_call_count(settings.EMBEDDING_API_PATH, 1)
assert (
responses.calls[0].request.body
== json_dumps(
{
"input": "any text",
"model": "mymodel",
"dimensions": 123,
"encoding_format": "float",
}
).encode()
)
def test_albert_service_embedding__not_configured(settings):
"""Should raise if not configured"""
settings.EMBEDDING_API_PATH = ""
with pytest.raises(albert.AlbertAIError):
albert.AlbertAI().embedding("any text")
@responses.activate
def test_albert_service_embedding__unexpected_content(settings):
"""Should raise if the API response is invalid"""
enable_hybrid_search(settings)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
body="invalid !!",
status=200,
)
with pytest.raises(albert.AlbertAIError) as err:
albert.AlbertAI().embedding("any text")
assert err.value.message == "Unexpected content : invalid !!"
@responses.activate
def test_albert_service_embedding__invalid_response(settings):
"""Should raise if the API returned error"""
enable_hybrid_search(settings)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
json={"detail": "Invalid request"},
status=400,
)
with pytest.raises(albert.AlbertAIError) as err:
albert.AlbertAI().embedding("any text")
assert err.value.message == "Invalid request"
@pytest.mark.parametrize(
"content",
[
"PDF content",
b"PDF content",
BytesIO(b"PDF content"),
],
)
@responses.activate
def test_albert_service_convert(settings, content):
"""Should return a converted PDF from Albert API"""
settings.ALBERT_PARSE_ENDPOINT = "https://test.albert.api/v1/parse"
responses.add(
responses.POST,
settings.ALBERT_PARSE_ENDPOINT,
json={
"data": [
{"content": "Markdown line 1"},
{"content": "Markdown line 2"},
]
},
status=200,
)
assert albert.AlbertAI().convert(content) == ("Markdown line 1\nMarkdown line 2")
responses.assert_call_count(settings.ALBERT_PARSE_ENDPOINT, 1)
def test_albert_service_convert__arguments(settings):
"""Should return a converted PDF from Albert API"""
settings.ALBERT_PARSE_ENDPOINT = "https://test.albert.api/v1/parse"
with mock.patch("core.services.albert.AlbertAI._request_api") as mock_api:
content = BytesIO(b"PDF content")
assert albert.AlbertAI().convert(content, pages=5, output="json") == ""
mock_api.assert_called_once()
assert mock_api.call_args == mock.call(
"https://test.albert.api/v1/parse",
files={
"file": ("input", content, "application/pdf"),
},
data={
"output_format": "json",
"page_range": "0-5",
},
)
def test_albert_service_convert__not_configured(settings):
"""Should raise if not configured"""
settings.ALBERT_PARSE_ENDPOINT = ""
with pytest.raises(albert.AlbertAIError):
albert.AlbertAI().convert(BytesIO(b"PDF content"))
@responses.activate
def test_albert_service_convert__unexpected_content(settings):
"""Should raise if the API response is invalid"""
settings.ALBERT_PARSE_ENDPOINT = "https://test.albert.api/v1/parse"
responses.add(
responses.POST,
settings.ALBERT_PARSE_ENDPOINT,
body="invalid !!",
status=200,
)
with pytest.raises(albert.AlbertAIError) as err:
albert.AlbertAI().convert(BytesIO(b"PDF content"))
assert err.value.message == "Unexpected content : invalid !!"
@responses.activate
def test_albert_service_convert__invalid_response(settings):
"""Should raise if the API returned error"""
settings.ALBERT_PARSE_ENDPOINT = "https://test.albert.api/v1/parse"
responses.add(
responses.POST,
settings.ALBERT_PARSE_ENDPOINT,
json={"detail": "Invalid request"},
status=400,
)
with pytest.raises(albert.AlbertAIError) as err:
albert.AlbertAI().convert(BytesIO(b"PDF content"))
assert err.value.message == "Invalid request"
src/backend/core/tests/test_api_documents_index_single.py
+4
View File
@@ -56,6 +56,10 @@ def test_api_documents_index_single_hybrid_enabled_success(settings):
If hybrid search is enabled, the indexing should have embedding of
dimension settings.EMBEDDING_DIMENSION.
"""
# Force opensearch index refresh because the task are runned synchronously
# within the tests.
settings.INDEXER_FORCE_REFRESH = True
service = factories.ServiceFactory()
enable_hybrid_search(settings)
responses.add(
src/backend/core/tests/test_indexer_services.py
+43 -40
View File
@@ -1,4 +1,4 @@
"""Tests indexing documents"""
"""Test index task service for documents"""
from dataclasses import asdict as dataasdict
from io import StringIO
@@ -162,9 +162,9 @@ def test_services_openbulk__raise_on_status():
def test_services_process_content():
"""Should convert the document content with the according converter"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
indexer.converters = {"application/pdf": mock_pdf}
assert (
@@ -201,9 +201,9 @@ def test_services_process_content():
def test_services_process_content__error():
"""Document content process should raise an error"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
indexer.converters = {"application/pdf": mock_pdf}
mock_pdf.side_effect = KeyError()
@@ -217,7 +217,7 @@ def test_services_process_content__error():
def test_services_search_documents():
"""Search document as an iterable of IndexDocument"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, batch_size=3)
indexer = IndexerTaskService(service, batch_size=3, force_refresh=True)
active_docs = factories.IndexDocumentFactory.create_batch(5)
inactive_docs = factories.IndexDocumentFactory.create_batch(3, is_active=False)
docs = active_docs + inactive_docs
@@ -240,7 +240,7 @@ def test_services_search_documents():
def test_services_search_documents__as_batch():
"""Search document as an iterable of list of IndexDocument (batches)"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, batch_size=3)
indexer = IndexerTaskService(service, batch_size=3, force_refresh=True)
docs = [
factories.IndexDocumentFactory(
id=f"00000000-0000-0000-0000-0000000000{index:02d}" # fake uuid to know order
@@ -285,12 +285,12 @@ def test_services_search_documents__as_batch():
def test_services_index(content_uri, content, mimetype, processed, status):
"""Add documents to index and set the content_status depending of their properties"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
doc = factories.IndexDocumentFactory(
mimetype=mimetype, content_uri=content_uri, content=content
)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.return_value = "processed content"
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.index([doc])
@@ -311,7 +311,7 @@ def test_services_index__process_errors():
the index() error list
"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf", content="this is a pdf"
)
@@ -322,7 +322,7 @@ def test_services_index__process_errors():
mimetype="application/unknown", content="this is a ???"
)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.side_effect = Exception()
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.index([pdf_doc, text_doc, unknown_doc])
@@ -344,9 +344,9 @@ def test_services_index__bulk_errors():
Opensearch bulk errors should appear in the index() error list
"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
doc_a, doc_b = factories.IndexDocumentFactory.create_batch(
2, mimetype="application/pdf", content="this is a pdf"
2, mimetype="text/plain", content="this is a text"
)
with mock.patch.object(indexer.client, "bulk") as mock_client_bulk:
@@ -424,10 +424,10 @@ def test_services_index__bulk_errors():
),
)
@responses.activate
def test_services_load_all(status, mimetype, is_active, expected):
def test_services_load_n_process_all(status, mimetype, is_active, expected):
"""Documents content should be downloaded and processed if needed"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
mimetype=mimetype,
@@ -452,10 +452,10 @@ def test_services_load_all(status, mimetype, is_active, expected):
assert [d.content_status for d in indexed_docs] == [status] * 3
assert [d.is_active for d in indexed_docs] == [is_active] * 3
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.return_value = "processed content"
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.load_all()
errors = indexer.load_n_process_all()
assert len(errors) == 0
@@ -465,10 +465,10 @@ def test_services_load_all(status, mimetype, is_active, expected):
@responses.activate
def test_services_load_all__errors():
def test_services_load_n_process_all__errors():
"""Should return download and processing errors"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf",
content_uri="http://localhost/mydoc",
@@ -503,10 +503,10 @@ def test_services_load_all__errors():
status=400,
)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.side_effect = Exception()
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.load_all()
errors = indexer.load_n_process_all()
assert len(errors) == 3
assert sorted(((e.id, e.message) for e in errors)) == sorted(
@@ -534,10 +534,10 @@ def test_services_load_all__errors():
@responses.activate
def test_services_load_all__bulk_errors():
def test_services_load_n_process_all__bulk_errors():
"""Should return bulk indexation errors"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
doc_a, doc_b = factories.IndexDocumentFactory.create_batch(
2,
mimetype="text/plain",
@@ -572,7 +572,7 @@ def test_services_load_all__bulk_errors():
with mock.patch.object(indexer.client, "bulk") as mock_client_bulk:
mock_client_bulk.return_value = bulk_response
errors = indexer.load_all()
errors = indexer.load_n_process_all()
assert [(e.id, e.message) for e in errors] == [
(doc_a.id, "Unknown error"),
@@ -639,10 +639,10 @@ def test_services_load_all__bulk_errors():
),
)
@responses.activate
def test_services_preprocess_all(status, mimetype, is_active, expected):
def test_services_process_all(status, mimetype, is_active, expected):
"""Documents content should be processed if needed"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
mimetype=mimetype,
@@ -667,10 +667,10 @@ def test_services_preprocess_all(status, mimetype, is_active, expected):
assert [d.content_status for d in indexed_docs] == [status] * 3
assert [d.is_active for d in indexed_docs] == [is_active] * 3
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.return_value = "processed content"
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.preprocess_all()
errors = indexer.process_all()
assert len(errors) == 0
@@ -679,10 +679,10 @@ def test_services_preprocess_all(status, mimetype, is_active, expected):
assert [d.content for d in indexed_docs] == [expected["content"]] * 3
def test_services_preprocess_all__errors():
def test_services_process_all__errors():
"""Documents content should be processed if needed"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf",
content_status=enums.ContentStatusEnum.LOADED,
@@ -696,10 +696,10 @@ def test_services_preprocess_all__errors():
actions.index_document(pdf_doc)
actions.index_document(unknown_doc)
with mock.patch("core.services.indexer_services.pdf_to_markdown") as mock_pdf:
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.side_effect = Exception()
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.preprocess_all()
errors = indexer.process_all()
assert len(errors) == 2
assert sorted(((e.id, e.message) for e in errors)) == sorted(
@@ -784,7 +784,7 @@ def test_services_embed_all(settings, status, embedding_model, is_active, expect
assert check_hybrid_search_enabled()
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
content_status=status,
@@ -826,7 +826,7 @@ def test_services_embed_all__errors(settings):
enable_hybrid_search(settings)
service = factories.ServiceFactory()
indexer = IndexerTaskService(service)
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(3)
with openbulk(service.index_name, refresh=True) as actions:
@@ -836,14 +836,17 @@ def test_services_embed_all__errors(settings):
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
json={"message": "Invalid request"},
json={"detail": "Invalid request"},
status=400,
)
errors = indexer.embed_all(model_name="mymodel")
assert sorted([(e.id, e.message) for e in errors]) == sorted(
[(d.id, "Unable to build embedding for the document") for d in docs]
[
(d.id, "Unable to build embedding for the document : Invalid request")
for d in docs
]
)
indexed_docs = indexer.search_documents({"match_all": {}})
@@ -853,11 +856,11 @@ def test_services_embed_all__errors(settings):
)
@mock.patch("core.services.indexer_services.IndexerTaskService.load_all")
@mock.patch("core.services.indexer_services.IndexerTaskService.preprocess_all")
@mock.patch("core.services.indexer_services.IndexerTaskService.load_n_process_all")
@mock.patch("core.services.indexer_services.IndexerTaskService.process_all")
@mock.patch("core.services.indexer_services.IndexerTaskService.embed_all")
def test_dispatch_indexing_tasks(
mock_load_all, mock_process_all, mock_embed_all, settings
mock_load_n_process_all, mock_process_all, mock_embed_all, settings
):
"""An different task for should be started depending of the state of the documents"""
settings.INDEXER_TASK_COUNTDOWN = 0
@@ -886,6 +889,6 @@ def test_dispatch_indexing_tasks(
dispatch_indexing_tasks(service, [waiting_doc, loaded_doc, ready_doc])
mock_load_all.assert_called()
mock_load_n_process_all.assert_called()
mock_process_all.assert_called()
mock_embed_all.assert_called()
src/backend/demo/management/commands/albert.py
+58
View File
@@ -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
+5
View File
@@ -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.",
@@ -574,6 +578,7 @@ class Test(Base):
CELERY_TASK_ALWAYS_EAGER = values.BooleanValue(True)
INDEXER_FORCE_REFRESH = True
HYBRID_SEARCH_ENABLED = False
def __init__(self):
# pylint: disable=invalid-name