Compare commits

...

4 Commits

Author SHA1 Message Date
Fabre Florian a2d3b08fc4 (backend) add preprocess support for indexer
New processors mechanism in IndexerTaskService : after the loading & conversion
steps a list a functions can be chained to transform the document content (like
django middlewares)

Signed-off-by: Fabre Florian <ffabre@hybird.org>
2025-11-28 06:53:52 +01:00
Fabre Florian 55d97797fe (backend) fix service index prefixed name in create_demo
Use service.index_name instead of service.name in create_demo
command.

Signed-off-by: Fabre Florian <ffabre@hybird.org>
2025-11-27 06:49:33 +01:00
Fabre Florian 56f0fb8df3 (backend) albert AI client & pdf conversion
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>
2025-11-27 06:49:33 +01:00
Fabre Florian 7e5fc7c138 (backend) refactor indexation pipeline
Add support for deferred loading, preprocessing & embedding of documents.
Add mimetype, language, content_status & content_uri fields in document schema.

Signed-off-by: Fabre Florian <ffabre@hybird.org>
2025-11-27 06:49:33 +01:00
21 changed files with 2303 additions and 75 deletions
+2
View File
@@ -25,3 +25,5 @@ and this project adheres to
- 🐛(backend) Fix parallel test execution issues
- ✨(backend) Add OPENSEARCH_INDEX_PREFIX setting to prevent naming overlaping
issues if the opensearch database is shared between apps.
- ✨(backend) Add IndexerTaskService to handle heavy indexation task
asynchronously.
+3
View File
@@ -42,6 +42,9 @@ These are the environment variables you can set for the `find-backend` container
| FRONTEND_CSS_URL | To add a external css file to the app | |
| FRONTEND_HOMEPAGE_FEATURE_ENABLED | Frontend feature flag to display the homepage | false |
| FRONTEND_THEME | Frontend theme to use | |
| INDEXER_DOWNLOAD_TIMEOUT | Document download request timeout in seconds | 100 |
| INDEXER_FORCE_REFRESH | Enable Opensearch index refresh after bulk changes | false |
| INDEXER_TASK_COUNTDOWN | Indexation sub-tasks countdown duration in seconds | 1 |
| LANGUAGE_CODE | Default language | en-us |
| LOGGING_LEVEL_LOGGERS_APP | Application logging level. options are "DEBUG", "INFO", "WARN", "ERROR", "CRITICAL" | INFO |
| LOGGING_LEVEL_LOGGERS_ROOT | Default logging level. options are "DEBUG", "INFO", "WARN", "ERROR", "CRITICAL" | INFO |
+8
View File
@@ -13,6 +13,14 @@ class ReachEnum(str, Enum):
RESTRICTED = "restricted"
class ContentStatusEnum(str, Enum):
"""Content upload status for indexed documents"""
READY = "ready"
LOADED = "loaded"
WAIT = "wait"
# Fields
CREATED_AT = "created_at"
+39
View File
@@ -23,6 +23,7 @@ class DocumentSchemaFactory(factory.DictFactory):
title = factory.Sequence(lambda n: f"Test title {n!s}")
path = factory.Sequence(lambda n: f"000{n}")
content = factory.Sequence(lambda n: f"Test content {n!s}")
content_uri = None
created_at = factory.LazyFunction(
lambda: fake.date_time_this_decade(tzinfo=timezone.get_current_timezone())
)
@@ -56,3 +57,41 @@ class ServiceFactory(factory.django.DjangoModelFactory):
class Meta:
model = models.Service
class IndexDocumentFactory(factory.DictFactory):
"""
A factory for generating dictionaries that represent a document for
indexation for testing and development purposes.
"""
id = factory.LazyFunction(lambda: str(uuid4()))
title = factory.Sequence(lambda n: f"Test title {n!s}")
path = factory.Sequence(lambda n: f"000{n}")
content = factory.Sequence(lambda n: f"Test content {n!s}")
content_status = enums.ContentStatusEnum.READY
content_uri = ""
created_at = factory.LazyFunction(
lambda: fake.date_time_this_decade(tzinfo=timezone.get_current_timezone())
)
size = factory.LazyFunction(lambda: fake.random_int(min=0, max=1024**2))
users = factory.LazyFunction(lambda: [str(uuid4()) for _ in range(3)])
groups = factory.LazyFunction(lambda: [slugify(fake.word()) for _ in range(3)])
reach = factory.Iterator(list(enums.ReachEnum))
depth = 1
numchild = 0
is_active = True
class Meta:
"""Factory model"""
model = models.IndexDocument
@factory.lazy_attribute
def updated_at(self):
"""Ensure updated_at is after created_at and before now"""
return fake.date_time_between(
start_date=self.created_at,
end_date=timezone.now(),
tzinfo=timezone.get_current_timezone(),
)
+62
View File
@@ -2,6 +2,10 @@
import secrets
import string
from dataclasses import dataclass
from dataclasses import field as datafield
from typing import List, Optional
from uuid import UUID
from django.conf import settings
from django.contrib.auth.models import AbstractUser
@@ -9,8 +13,11 @@ from django.db import models
from django.db.models.functions import Length
from django.utils.functional import cached_property
from django.utils.text import slugify
from django.utils.timezone import datetime
from django.utils.translation import gettext_lazy as _
from . import enums
models.CharField.register_lookup(Length)
TOKEN_LENGTH = 50
@@ -73,3 +80,58 @@ class Service(models.Model):
def index_name(self):
"""Returns the opensearch index for the service"""
return get_opensearch_index_name(self.name)
# pylint: disable=too-many-instance-attributes
@dataclass
class IndexDocument:
"""Represents the _source data of opensearch entry"""
id: UUID
title: str = ""
depth: int = 0
path: str = ""
numchild: int = 0
content: str = ""
content_uri: str = ""
content_status: enums.ContentStatusEnum = enums.ContentStatusEnum.READY
created_at: Optional[datetime] = None
updated_at: Optional[datetime] = None
size: int = 0
users: List[str] = datafield(default_factory=list)
groups: List[str] = datafield(default_factory=list)
reach: enums.ReachEnum = enums.ReachEnum.RESTRICTED
is_active: bool = True
mimetype: str = "text/plain"
embedding: Optional[str] = None
embedding_model: Optional[str] = None
@property
def is_waiting(self):
"""Retuns true if in waiting status"""
return self.content_status == enums.ContentStatusEnum.WAIT
@property
def is_ready(self):
"""Retuns true if in ready status"""
return self.content_status == enums.ContentStatusEnum.READY
@property
def is_loaded(self):
"""Retuns true if in loaded status"""
return self.content_status == enums.ContentStatusEnum.LOADED
@staticmethod
def from_dict(data):
"""Create an instance from dict data"""
document = IndexDocument(**data)
document.reach = enums.ReachEnum(document.reach)
document.content_status = enums.ContentStatusEnum(document.content_status)
if isinstance(document.created_at, str):
document.created_at = datetime.fromisoformat(document.created_at)
if isinstance(document.updated_at, str):
document.updated_at = datetime.fromisoformat(document.updated_at)
return document
+3
View File
@@ -7,6 +7,7 @@ from django.utils.text import slugify
from pydantic import (
UUID4,
AnyUrl,
AwareDatetime,
BaseModel,
BeforeValidator,
@@ -29,6 +30,7 @@ class DocumentSchema(BaseModel):
path: Annotated[str, Field(max_length=300)]
numchild: Annotated[int, Field(ge=0)]
content: Annotated[str, Field(min_length=0)]
content_uri: Annotated[Optional[AnyUrl], Field(None)]
created_at: AwareDatetime
updated_at: AwareDatetime
size: Annotated[int, Field(ge=0, le=100 * 1024**3)] # File size limited to 100GB
@@ -38,6 +40,7 @@ class DocumentSchema(BaseModel):
)
reach: Optional[enums.ReachEnum] = Field(default=enums.ReachEnum.RESTRICTED)
is_active: bool
mimetype: Annotated[str, Field(default="text/plain")]
model_config = ConfigDict(
str_min_length=1, str_strip_whitespace=True, use_enum_values=True
+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
+10
View File
@@ -0,0 +1,10 @@
"""Document content conversion tools"""
from io import BytesIO
from .albert import AlbertAI
def pdf_to_markdown(content: BytesIO):
"""Convert PDF stream into markdown"""
return AlbertAI().convert(content=content, mimetype="application/pdf")
@@ -0,0 +1,567 @@
"""Indexation tasks tools"""
import logging
from contextlib import contextmanager
from dataclasses import asdict as dataasdict
from io import BytesIO
from django.conf import settings
import requests
from core import enums
from core.models import IndexDocument, Service
from . import converters
from .albert import AlbertAI, AlbertAIError
from .opensearch import (
check_hybrid_search_enabled,
ensure_index_exists,
format_document,
opensearch_client,
)
logger = logging.getLogger(__name__)
INDEXABLE_MIMETYPES = ["text/"]
class IndexerError(Exception):
"""Base exception for indexer service"""
def __init__(self, message, _id=None):
super().__init__(message)
self.id = _id
if isinstance(message, dict):
self.error_dict = message
elif isinstance(message, list):
self.error_list = message
else:
self.message = message
class IndexContentError(IndexerError):
"""Raised on content loading or conversion issues"""
class IndexBulkError(IndexerError):
"""Raised on IndexBulkTransaction commit issues"""
def __init__(self, message, action=None, _id=None, params=None):
super().__init__(message, _id=_id)
self.action = action
self.params = params
class IndexBulkTransaction:
"""
Helper for bulk changes in opensearch.
"""
def __init__(
self,
index_name,
client=None,
):
self.client = client or opensearch_client()
self.index_name = index_name
self._actions = []
self._errors = {}
def errors(self):
"""Returns errors from last commit"""
return self._errors
def create(self, uid, data, **kwargs):
"""Create an opensearch entry. Raises if exists"""
self._actions.append({"create": {"_id": uid, **kwargs}})
self._actions.append(data)
return self
def index_document(self, doc: IndexDocument, **kwargs):
"""Index a document model instance"""
data = dataasdict(doc)
_id = data.pop("id")
return self.index(_id, data=data, **kwargs)
def index(self, uid, data, **kwargs):
"""Index an openseach entry (create or replace)"""
self._actions.append({"index": {"_id": uid, **kwargs}})
self._actions.append(data)
return self
def update(self, uid, data, **kwargs):
"""Update an openseach entry"""
self._actions.append({"update": {"_id": uid, **kwargs}})
self._actions.append({"doc": data})
return self
def commit(self, raise_on_status=False, refresh=False):
"""Send all actions to opensearch database & create the index if needed."""
# No action, do nothing
if not self._actions:
return
# Build index if needed.
ensure_index_exists(self.index_name)
response = self.client.bulk(
index=self.index_name, body=self._actions, refresh=refresh
)
errors = self._errors = []
self._actions = []
for item in response["items"]:
for action, result in item.items():
status_code = result["status"]
if status_code >= 300:
message = result.get("error", {}).get("reason", "Unknown error")
errors.append(
IndexBulkError(
message, action=action, _id=result.get("_id"), params=result
)
)
if errors and raise_on_status:
raise IndexBulkError(errors)
@contextmanager
def openbulk(index_name, client=None, raise_on_status=False, refresh=False):
"""
Create a bulk transaction and commit it swiftly after use.
"""
transaction = IndexBulkTransaction(index_name, client=client)
yield transaction
transaction.commit(raise_on_status=raise_on_status, refresh=refresh)
def offsetpaginate(*index_names, query, sort=None, size=100, client=None):
"""
Paginator based on the next offset filter. Iterates on batches of 'limit' size.
"""
client = client or opensearch_client()
sort = sort or [{"_id": "asc"}]
index = ",".join(index_names)
body = {
"query": query or {},
"sort": sort,
"size": size,
}
page = client.search(
body=body,
index=index,
seq_no_primary_term=True,
ignore_unavailable=True,
)
hits = page["hits"]["hits"]
count = len(hits)
if count > 0:
yield hits
while count == size:
page = client.search(
index=index,
body={
**body,
"search_after": hits[-1]["sort"],
},
seq_no_primary_term=True,
ignore_unavailable=True,
)
hits = page["hits"]["hits"]
count = len(hits)
if count > 0:
yield hits
def hit_to_doc(hit):
"""Returns an IndexDocument instance from opensearch hit data"""
doc = IndexDocument.from_dict({**hit["_source"], "id": hit["_id"]})
doc.hit = hit
return doc
def match_mimetype_glob(mimetype, pattern):
"""
Returns true if the mimetype match with the pattern.
If a pattern ends with / all subtypes are valid, if not it have to
perfectly match. e.g :
- application/pdf only match "application/pdf" and not "application/pdf+bin"
- application/ match any "application/*"
"""
if len(mimetype) < 1:
return False
if pattern.endswith("/"):
return mimetype.startswith(pattern)
return mimetype == pattern
def is_allowed_mimetype(mimetype, patterns):
"""
Returns true if the mimetype is not empty and matches any of the allowed patterns.
"""
return len(mimetype) > 0 and any(
match_mimetype_glob(mimetype, pattern) for pattern in patterns
)
class IndexerTaskService:
"""
Service used by indexation tasks for content loading, format conversion
and embedding
"""
# V2 : Use settings to define the list of converters
converters = {"application/pdf": converters.pdf_to_markdown}
processors = [
# Put prepare_document_for_indexing() here !
]
def __init__(
self, service: Service, batch_size=100, client=None, force_refresh=None
):
self.service = service
self.batch_size = batch_size
self.client = client or opensearch_client()
self.force_refresh = force_refresh or settings.INDEXER_FORCE_REFRESH
self.download_timeout = settings.INDEXER_DOWNLOAD_TIMEOUT
def get_converter(self, mimetype):
"""
Retrieve a converter for the mimetype or None if the mimitype is allowed
for indexation.
"""
try:
if is_allowed_mimetype(mimetype, INDEXABLE_MIMETYPES):
return None
return self.converters[mimetype]
except KeyError as e:
raise IndexContentError(
f"No such converter for the unindexable mimetype {mimetype}"
) from e
def convert_content(self, content, document):
"""Transforms the document file data into an indexable format"""
try:
converter = self.get_converter(document.mimetype)
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
raise e
except Exception as e:
raise IndexContentError(
f"Unable to convert content with the mimetype {document.mimetype}",
_id=document.id,
) from e
def process_document(self, document):
"""Transforms the document file data into an indexable format"""
try:
processors = list(self.processors)
processors.reverse()
while len(processors) > 0:
document = processors.pop()(document)
return document
except IndexContentError as e:
e.id = document.id
raise e
except Exception as e:
raise IndexContentError(
f"Unable to process content : {e}",
_id=document.id,
) from e
def stream_content(self, document):
"""Open download stream containing the file data"""
try:
response = requests.get(
document.content_uri, stream=True, timeout=self.download_timeout
)
response.raise_for_status()
return response.raw
except requests.RequestException as e:
raise IndexerError(str(e), _id=document.id) from e
def search(self, query, *, sort=None, as_batch=False, batch_size=None):
"""Returns the paginated results from the opensearch query as an iterable"""
pages = offsetpaginate(
self.service.index_name,
query=query,
sort=sort,
size=batch_size or self.batch_size,
client=self.client,
)
if as_batch:
yield from pages
else:
for hits in pages:
yield from hits
def search_documents(self, query, *, sort=None, as_batch=False, batch_size=None):
"""
Returns the paginated results from the opensearch query as IndexDocument instances
"""
if as_batch:
for hits in self.search(
query, sort=sort, as_batch=True, batch_size=batch_size
):
yield [hit_to_doc(hit) for hit in hits]
else:
for hit in self.search(query, sort=sort, batch_size=batch_size):
yield hit_to_doc(hit)
def process_all(self, batch_size=None):
"""
Gets all the documents in waiting status to load and convert their content
Returns an error dict.
"""
errors = []
index_name = self.service.index_name
batch_size = batch_size or self.batch_size
doc_batches = self.search_documents(
query={
"bool": {
"must": [
{
"term": {
"content_status": enums.ContentStatusEnum.LOADED.value
}
},
{"term": {"is_active": True}},
]
}
},
batch_size=batch_size,
as_batch=True,
)
for docs in doc_batches:
with openbulk(
index_name, client=self.client, refresh=self.force_refresh
) as actions:
for doc in docs:
try:
# V2 : Use asyncio loop to parallelize conversion
doc = self.process_document(doc) # noqa : PLW2901
actions.update(
doc.id,
data={
"content_status": enums.ContentStatusEnum.READY.value,
"content": doc.content,
},
# if_seq_no and if_primary_term ensure we only update indexes
# if the document hasn't changed
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())
return errors
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.
"""
errors = []
index_name = self.service.index_name
batch_size = batch_size or self.batch_size
doc_batches = self.search_documents(
query={
"bool": {
"must": [
{
"term": {
"content_status": enums.ContentStatusEnum.WAIT.value
}
},
{"term": {"is_active": True}},
]
}
},
batch_size=batch_size,
as_batch=True,
)
for docs in doc_batches:
with openbulk(
index_name, client=self.client, refresh=self.force_refresh
) as actions:
for doc in docs:
try:
# V2 : Use asyncio loop to parallelize downloads
content = self.stream_content(doc)
doc.content = self.convert_content(content, doc)
doc.content_status = enums.ContentStatusEnum.LOADED
except IndexerError as e:
# V2 ; handle retry here and remove the entry if not working
errors.append(e)
continue
try:
doc = self.process_document(doc) # noqa : PLW2901
doc.content_status = enums.ContentStatusEnum.READY
except IndexerError as e:
errors.append(e)
actions.update(
doc.id,
data={
"content_status": doc.content_status.value,
"content": doc.content,
},
# if_seq_no and if_primary_term ensure we only update indexes
# if the document hasn't changed
if_seq_no=doc.hit["_seq_no"],
if_primary_term=doc.hit["_primary_term"],
)
errors.extend(actions.errors())
return errors
def embed_all(self, model_name=None, batch_size=None):
"""
Re-index all the document in "ready" status without embbeding or with a different model.
"""
errors = []
# Hybrid search is disabled, skip it
if not check_hybrid_search_enabled():
return ()
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": {
"must": [
{
"term": {
"content_status": enums.ContentStatusEnum.READY.value
}
},
{"term": {"is_active": True}},
],
"should": [
{"bool": {"must_not": {"exists": {"field": "embedding"}}}},
{
"bool": {
"must_not": {"term": {"embedding_model": model_name}}
}
},
],
"minimum_should_match": 1,
},
},
batch_size=batch_size,
as_batch=True,
)
for docs in doc_batches:
with openbulk(
index_name, client=self.client, refresh=self.force_refresh
) as actions:
for doc in docs:
try:
# 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={
"embedding": embedding,
"embedding_model": model_name,
},
# if_seq_no and if_primary_term ensure we only update indexes
# if the document hasn't changed
if_seq_no=doc.hit["_seq_no"],
if_primary_term=doc.hit["_primary_term"],
)
errors.extend(actions.errors())
return errors
def index(self, documents):
"""
Index all the documents and initialize the status depending of the availability
of their content.
Returns the errors from the commit and convertion
"""
errors = []
with openbulk(
self.service.index_name, client=self.client, refresh=self.force_refresh
) as actions:
for doc in documents:
# Without content and a dowload uri : set WAIT status
if doc.content_uri and not doc.content:
doc.content_status = enums.ContentStatusEnum.WAIT
actions.index_document(doc)
continue
# A content but not directly indexable (e.g xml or html content) : convert them
if not is_allowed_mimetype(doc.mimetype, INDEXABLE_MIMETYPES):
try:
doc.content = self.convert_content(doc.content, doc)
except IndexContentError as err:
errors.append(IndexBulkError(str(err), _id=doc.id))
continue
# Preprocess the content of the document
try:
doc = self.process_document(doc) # noqa: PLW2901
except IndexContentError as err:
# V2 : Add retry mechanism with LOADED status ?
errors.append(IndexBulkError(str(err), _id=doc.id))
continue
# The content exists and is indexable : set READY
doc.content_status = enums.ContentStatusEnum.READY
actions.index_document(doc)
errors.extend(actions.errors())
return errors
+4
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,
@@ -283,6 +284,8 @@ def ensure_index_exists(index_name):
},
"numchild": {"type": "integer"},
"content": {"type": "text"},
"content_uri": {"type": "text", "index": False},
"content_status": {"type": "keyword", "index": False},
"created_at": {"type": "date"},
"updated_at": {"type": "date"},
"size": {"type": "long"},
@@ -290,6 +293,7 @@ def ensure_index_exists(index_name):
"groups": {"type": "keyword"},
"reach": {"type": "keyword"},
"is_active": {"type": "boolean"},
"mimetype": {"type": "keyword", "index": False},
"embedding": {
# for simplicity, embedding is always present but is empty
# when hybrid search is disabled
View File
+83
View File
@@ -0,0 +1,83 @@
"""Celery task for deferred indexing"""
import logging
from typing import List
from django.conf import settings
from core import models
from core.services.indexer_services import IndexerTaskService
from core.services.opensearch import check_hybrid_search_enabled
from find.celery_app import app
logger = logging.getLogger(__file__)
def get_service(service_id):
"""Returns the required active service"""
try:
return models.Service.objects.get(pk=service_id, is_active=True)
except models.Service.DoesNotExist:
logging.warning("The service {service_id} does not exit or disabled")
return None
@app.task
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, force_refresh=True)
logger.info("Start deferred loading on index %s", service.index_name)
indexer.load_n_process_all()
@app.task
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, force_refresh=True)
logger.info("Start deferred preprocessing on index %s", service.index_name)
indexer.process_all()
@app.task
def embedding_task(service_id):
"""Celery Task : Re-index documents for embedding."""
service = get_service(service_id)
if service is not None:
indexer = IndexerTaskService(service, force_refresh=True)
logger.info("Start embedding on index %s", service.index_name)
indexer.embed_all()
def dispatch_indexing_tasks(service, documents: List[models.IndexDocument]):
"""
Trigger task related to the different status of the documents
"""
countdown = settings.INDEXER_TASK_COUNTDOWN
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 task for deferred loading if the file is too big
if should_load:
load_n_process_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)
# Trigger task for semantic indexation
if should_embed:
embedding_task.apply_async((service.pk,), countdown=countdown)
@@ -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"
@@ -193,6 +193,12 @@ def test_api_documents_index_bulk_ensure_index():
"bool_parsing",
"Input should be a valid boolean, unable to interpret input",
),
(
"content_uri",
"notanurl",
"url_parsing",
"Input should be a valid URL, relative URL without a base",
),
],
)
def test_api_documents_index_bulk_invalid_document(
@@ -391,13 +397,20 @@ def test_api_documents_index_opensearch_errors():
with mock.patch.object(opensearch.opensearch_client(), "bulk") as mock_bulk:
mock_bulk.return_value = {
"items": [
{"index": {"status": 201}},
{"index": {"_id": documents[0]["id"], "status": 201}},
{
"index": {
"_id": documents[1]["id"],
"status": 400,
}
},
{"index": {"status": 403, "error": {"reason": "This is forbidden"}}},
{
"index": {
"_id": documents[2]["id"],
"status": 403,
"error": {"reason": "This is forbidden"},
}
},
]
}
@@ -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(
@@ -80,9 +84,13 @@ def test_api_documents_index_single_hybrid_enabled_success(settings):
new_indexed_document = opensearch.opensearch_client().get(
index=service.index_name, id=str(document["id"])
)
assert new_indexed_document["_version"] == 1
# already second version because the embedding as been done synchrously
assert new_indexed_document["_version"] == 2
assert new_indexed_document["_source"]["title"] == document["title"].strip().lower()
assert new_indexed_document["_source"]["content"] == document["content"]
assert new_indexed_document["_source"]["content_status"] == "ready"
assert (
new_indexed_document["_source"]["embedding"]
== albert_embedding_response.response["data"][0]["embedding"]
@@ -157,6 +165,8 @@ def test_api_documents_index_single_ensure_index(settings):
},
"numchild": {"type": "integer"},
"content": {"type": "text"},
"content_uri": {"type": "text", "index": False},
"content_status": {"type": "keyword", "index": False},
"created_at": {"type": "date"},
"updated_at": {"type": "date"},
"size": {"type": "long"},
@@ -164,6 +174,7 @@ def test_api_documents_index_single_ensure_index(settings):
"groups": {"type": "keyword"},
"reach": {"type": "keyword"},
"is_active": {"type": "boolean"},
"mimetype": {"type": "keyword", "index": False},
"embedding": {
"type": "knn_vector",
"dimension": settings.EMBEDDING_DIMENSION,
@@ -289,6 +300,12 @@ def test_api_documents_index_single_ensure_index(settings):
"bool_parsing",
"Input should be a valid boolean, unable to interpret input",
),
(
"content_uri",
"notanurl",
"url_parsing",
"Input should be a valid URL, relative URL without a base",
),
],
)
def test_api_documents_index_single_invalid_document(
@@ -445,3 +462,29 @@ def test_api_documents_index_empty_content_check():
== "Value error, Either title or content should have at least 1 character"
)
assert response.data[0]["type"] == "value_error"
def test_api_documents_index_content_uri():
"""Test document indexing with content_uri."""
service = factories.ServiceFactory()
document = factories.DocumentSchemaFactory.build()
document["content_uri"] = "http://localhost/mydoc"
response = APIClient().post(
"/api/v1.0/documents/index/",
document,
HTTP_AUTHORIZATION=f"Bearer {service.token:s}",
format="json",
)
assert response.status_code == 201
assert response.json()["_id"] == str(document["id"])
new_indexed_document = opensearch.opensearch_client().get(
index=service.index_name, id=str(document["id"])
)
assert new_indexed_document["_version"] == 1
assert new_indexed_document["_source"]["title"] == document["title"].strip().lower()
assert new_indexed_document["_source"]["content"] == document["content"]
assert new_indexed_document["_source"]["content_uri"] == "http://localhost/mydoc"
@@ -0,0 +1,986 @@
"""Test index task service for documents"""
from dataclasses import asdict as dataasdict
from io import StringIO
from operator import attrgetter
from unittest import mock
import pytest
import responses
from core import enums, factories
from core.services import opensearch
from core.services.indexer_services import (
IndexBulkError,
IndexContentError,
IndexerTaskService,
hit_to_doc,
is_allowed_mimetype,
openbulk,
)
from core.tasks.indexer import dispatch_indexing_tasks
from core.tests.mock import albert_embedding_response
from core.tests.utils import (
check_hybrid_search_enabled,
enable_hybrid_search,
refresh_index,
)
pytestmark = pytest.mark.django_db
def search_all_docs(client, index_name):
"""Helper that returns all IndexDocument available on opensearch database"""
data = client.search(index=index_name, body={"query": {"match_all": {}}})
return [hit_to_doc(hit) for hit in data["hits"]["hits"]]
def test_services_is_allowed_mimetype():
"""
Check if the given mimetype matches with the allowed ones.
"""
assert is_allowed_mimetype("", ["text/"]) is False
assert is_allowed_mimetype("text/plain", []) is False
assert is_allowed_mimetype("text/plain", ["text/html"]) is False
assert is_allowed_mimetype("text/plain", ["text/"]) is True
assert is_allowed_mimetype("application/pdf", ["text/", "application/pdf"]) is True
assert (
is_allowed_mimetype("application/pdf+bin", ["text/", "application/pdf"])
is False
)
assert is_allowed_mimetype("application/pdf+bin", ["text/", "application/"]) is True
def test_services_openbulk():
"""Bulk opensearch actions."""
service = factories.ServiceFactory()
doc_create = factories.IndexDocumentFactory()
doc_update = factories.IndexDocumentFactory()
doc_index = factories.IndexDocumentFactory()
opensearch_client_ = opensearch.opensearch_client()
with pytest.raises(opensearch.NotFoundError):
opensearch_client_.indices.get(index=service.index_name)
with openbulk(index_name=service.index_name) as actions:
actions.create(doc_create.id, dataasdict(doc_create))
actions.index_document(doc_index)
actions.index(doc_update.id, data={})
actions.update(doc_update.id, dataasdict(doc_update))
# now the index exists
opensearch_client_.indices.get(index=service.index_name)
# no errors
assert len(actions.errors()) == 0
# no data too... the index is not refreshed
assert search_all_docs(opensearch_client_, service.index_name) == []
refresh_index(service.index_name)
# Now the docs are here
assert sorted(
search_all_docs(opensearch_client_, service.index_name), key=attrgetter("id")
) == sorted([doc_create, doc_update, doc_index], key=attrgetter("id"))
def test_services_openbulk__refresh():
"""Bulk opensearch actions and refresh the index"""
service = factories.ServiceFactory()
doc_create = factories.IndexDocumentFactory()
doc_update = factories.IndexDocumentFactory()
doc_index = factories.IndexDocumentFactory()
opensearch_client_ = opensearch.opensearch_client()
with pytest.raises(opensearch.NotFoundError):
opensearch_client_.indices.get(index=service.index_name)
with openbulk(index_name=service.index_name, refresh=True) as actions:
actions.create(doc_create.id, dataasdict(doc_create))
actions.index_document(doc_index)
actions.index(doc_update.id, data={})
actions.update(doc_update.id, dataasdict(doc_update))
# now the index exists
opensearch_client_.indices.get(index=service.index_name)
# no errors
assert len(actions.errors()) == 0
# Auto refreshed !
assert sorted(
search_all_docs(opensearch_client_, service.index_name), key=attrgetter("id")
) == sorted([doc_create, doc_update, doc_index], key=attrgetter("id"))
def test_services_openbulk__errors():
"""Should return all bulk errors"""
service = factories.ServiceFactory()
doc = factories.IndexDocumentFactory()
with openbulk(index_name=service.index_name) as actions:
actions.create(doc.id, dataasdict(doc))
actions.create(doc.id, dataasdict(doc))
assert len(actions.errors()) == 1
error = actions.errors()[0]
assert error.id == doc.id
assert error.action == "create"
assert (
error.message
== f"[{doc.id}]: version conflict, document already exists (current version [1])"
)
def test_services_openbulk__raise_on_status():
"""Should raise on bulk error"""
service = factories.ServiceFactory()
doc = factories.IndexDocumentFactory()
with pytest.raises(IndexBulkError) as err:
with openbulk(index_name=service.index_name, raise_on_status=True) as actions:
actions.create(doc.id, dataasdict(doc))
actions.create(doc.id, dataasdict(doc))
error = err.value
assert len(error.error_list) == 1
error = error.error_list[0]
assert error.id == doc.id
assert error.action == "create"
assert (
error.message
== f"[{doc.id}]: version conflict, document already exists (current version [1])"
)
def test_services_convert_content():
"""Should convert the document content with the according converter"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
indexer.converters = {"application/pdf": mock_pdf}
assert (
indexer.convert_content(
StringIO("This is a test"),
document=factories.IndexDocumentFactory(mimetype="text/plain"),
)
== "This is a test"
)
mock_pdf.assert_not_called()
assert (
indexer.convert_content(
"This is a test",
document=factories.IndexDocumentFactory(mimetype="text/plain"),
)
== "This is a test"
)
mock_pdf.assert_not_called()
mock_pdf.return_value = "This a converted PDF test"
assert (
indexer.convert_content(
"This is a test",
document=factories.IndexDocumentFactory(mimetype="application/pdf"),
)
== "This a converted PDF test"
)
mock_pdf.assert_called_once()
def test_services_convert_content__error():
"""Document content process should raise an error"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
with mock.patch("core.services.converters.pdf_to_markdown") as mock_pdf:
indexer.converters = {"application/pdf": mock_pdf}
mock_pdf.side_effect = KeyError()
with pytest.raises(IndexContentError):
indexer.convert_content(
StringIO("This is a test"),
document=factories.IndexDocumentFactory(mimetype="application/pdf"),
)
def test_services_process_content():
"""Should process the document content"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
def _processor_a(document):
document.content = f"A({document.content})"
return document
def _processor_b(document):
document.content = f"B({document.content})"
return document
indexer.processors = [_processor_a, _processor_b]
doc = indexer.process_document(
factories.IndexDocumentFactory(mimetype="text/plain", content="This is a test"),
)
assert doc.content == "B(A(This is a test))"
def test_services_process_content__error():
"""Document content process should raise an error"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
def _processor(document):
raise ValueError("Cannot preprocess it")
indexer.processors = [_processor]
with pytest.raises(IndexContentError):
indexer.process_document(
document=factories.IndexDocumentFactory(
mimetype="application/pdf", content="This is a test"
),
)
def test_services_search_documents():
"""Search document as an iterable of IndexDocument"""
service = factories.ServiceFactory()
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
assert not list(indexer.search_documents({"match_all": {}}))
indexer.index(docs)
assert list(indexer.search_documents({"match_all": {}})) == sorted(
docs, key=attrgetter("id")
)
assert list(
indexer.search_documents({"bool": {"must": {"term": {"is_active": True}}}})
) == sorted(active_docs, key=attrgetter("id"))
assert list(
indexer.search_documents({"bool": {"must": {"term": {"is_active": False}}}})
) == sorted(inactive_docs, key=attrgetter("id"))
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, force_refresh=True)
docs = [
factories.IndexDocumentFactory(
id=f"00000000-0000-0000-0000-0000000000{index:02d}" # fake uuid to know order
)
for index in range(1, 10)
]
assert not list(indexer.search_documents({"match_all": {}}, as_batch=True))
indexer.index(docs)
assert list(indexer.search_documents({"match_all": {}}, as_batch=True)) == [
docs[:3],
docs[3:6],
docs[6:],
]
@pytest.mark.parametrize(
"content_uri, content, mimetype, processed, status",
(
("", "", "text/plain", False, enums.ContentStatusEnum.READY),
("", "", "application/pdf", True, enums.ContentStatusEnum.READY),
(
"http://localhost/mydoc",
"",
"application/pdf",
False,
enums.ContentStatusEnum.WAIT,
),
(
"http://localhost/mydoc",
"This is a test",
"application/pdf",
True,
enums.ContentStatusEnum.READY,
),
("", "This is a test", "application/pdf", True, enums.ContentStatusEnum.READY),
("", "This is a test", "text/plain", False, enums.ContentStatusEnum.READY),
),
)
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, force_refresh=True)
doc = factories.IndexDocumentFactory(
mimetype=mimetype, content_uri=content_uri, content=content
)
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])
assert len(errors) == 0
assert mock_pdf.called == processed
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert len(indexed_docs) == 1
indexed_doc = indexed_docs[0]
assert indexed_doc.content_status == status
def test_services_index__convert_errors():
"""
Documents with unsupported mimetypes or conversion issues should appear in
the index() error list
"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf", content="this is a pdf"
)
text_doc = factories.IndexDocumentFactory(
mimetype="text/plain", content="this is a text"
)
unknown_doc = factories.IndexDocumentFactory(
mimetype="application/unknown", content="this is a ???"
)
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])
assert [(e.id, e.message) for e in errors] == [
(pdf_doc.id, "Unable to convert content with the mimetype application/pdf"),
(
unknown_doc.id,
"No such converter for the unindexable mimetype application/unknown",
),
]
# we only add the text doc to the index
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert len(indexed_docs) == 1
assert indexed_docs[0].id == text_doc.id
def test_services_index__process_errors():
"""
Documents with failed preprocessing should appear in the index() error list
"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf", content="this is a pdf"
)
text_doc = factories.IndexDocumentFactory(
mimetype="text/plain", content="this is a text"
)
def _preprocess(document):
if document.mimetype == "application/pdf":
raise ValueError("Not supported")
return document
indexer.processors = [_preprocess]
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([pdf_doc, text_doc])
assert [(e.id, e.message) for e in errors] == [
(pdf_doc.id, "Unable to process content : Not supported"),
]
# we only add the text doc to the index
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert len(indexed_docs) == 1
assert indexed_docs[0].id == text_doc.id
def test_services_index__bulk_errors():
"""
Opensearch bulk errors should appear in the index() error list
"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
doc_a, doc_b = factories.IndexDocumentFactory.create_batch(
2, mimetype="text/plain", content="this is a text"
)
with mock.patch.object(indexer.client, "bulk") as mock_client_bulk:
mock_client_bulk.return_value = {
"items": [
{"index": {"status": 500, "_id": doc_a.id}},
{
"index": {
"status": 400,
"_id": doc_b.id,
"error": {"reason": "Invalid arguments"},
}
},
]
}
errors = indexer.index([doc_a, doc_b])
assert [(e.id, e.message) for e in errors] == [
(doc_a.id, "Unknown error"),
(doc_b.id, "Invalid arguments"),
]
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert len(indexed_docs) == 0
@pytest.mark.parametrize(
"status, mimetype, is_active , expected",
(
(
enums.ContentStatusEnum.READY,
"text/plain",
True,
{
"content": "initial content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.LOADED,
"text/plain",
True,
{
"content": "initial content",
"status": enums.ContentStatusEnum.LOADED,
},
),
(
enums.ContentStatusEnum.WAIT,
"text/plain",
True,
{
"content": "loaded content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.WAIT,
"application/pdf",
True,
{
"content": "processed content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.WAIT,
"text/plain",
False,
{
"content": "initial content",
"status": enums.ContentStatusEnum.WAIT,
},
),
),
)
@responses.activate
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, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
mimetype=mimetype,
content_status=status,
content_uri="http://localhost/mydoc",
content="initial content",
is_active=is_active,
)
with openbulk(service.index_name, refresh=True) as actions:
for doc in docs:
actions.index_document(doc)
responses.add(
responses.GET,
"http://localhost/mydoc",
body="loaded content",
status=200,
)
indexed_docs = list(indexer.search_documents({"match_all": {}}))
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.converters.pdf_to_markdown") as mock_pdf:
mock_pdf.return_value = "processed content"
indexer.converters = {"application/pdf": mock_pdf}
errors = indexer.load_n_process_all()
assert len(errors) == 0
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert [d.content_status for d in indexed_docs] == [expected["status"]] * 3
assert [d.content for d in indexed_docs] == [expected["content"]] * 3
@responses.activate
def test_services_load_n_process_all__errors():
"""Should return download and processing errors"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf",
content_uri="http://localhost/mydoc",
content_status=enums.ContentStatusEnum.WAIT,
)
unknown_doc = factories.IndexDocumentFactory(
mimetype="application/unknown",
content_uri="http://localhost/mydoc",
content_status=enums.ContentStatusEnum.WAIT,
)
invalid_uri_doc = factories.IndexDocumentFactory(
mimetype="application/pdf",
content_uri="http://localhost/invalid",
content_status=enums.ContentStatusEnum.WAIT,
)
unsupported_doc = factories.IndexDocumentFactory(
mimetype="text/unsupported",
content_uri="http://localhost/mydoc",
content_status=enums.ContentStatusEnum.WAIT,
)
with openbulk(service.index_name, refresh=True) as actions:
actions.index_document(pdf_doc)
actions.index_document(unknown_doc)
actions.index_document(invalid_uri_doc)
actions.index_document(unsupported_doc)
responses.add(
responses.GET,
"http://localhost/mydoc",
body="loaded content",
status=200,
)
responses.add(
responses.GET,
"http://localhost/invalid",
body="not accessible content",
status=400,
)
def _invalid_process(document):
if document.mimetype == "text/unsupported":
raise ValueError("unsupported format")
indexer.processors = [_invalid_process]
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_n_process_all()
assert len(errors) == 4
assert sorted(((e.id, e.message) for e in errors)) == sorted(
[
(pdf_doc.id, "Unable to convert content with the mimetype application/pdf"),
(
unknown_doc.id,
"No such converter for the unindexable mimetype application/unknown",
),
(
invalid_uri_doc.id,
"400 Client Error: Bad Request for url: http://localhost/invalid",
),
(
unsupported_doc.id,
"Unable to process content : unsupported format",
),
]
)
indexed_docs = indexer.search_documents({"match_all": {}})
assert sorted((d.id, d.content_status) for d in indexed_docs) == sorted(
[
(pdf_doc.id, enums.ContentStatusEnum.WAIT),
(unknown_doc.id, enums.ContentStatusEnum.WAIT),
(invalid_uri_doc.id, enums.ContentStatusEnum.WAIT),
(unsupported_doc.id, enums.ContentStatusEnum.LOADED),
]
)
@responses.activate
def test_services_load_n_process_all__bulk_errors():
"""Should return bulk indexation errors"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
doc_a, doc_b = factories.IndexDocumentFactory.create_batch(
2,
mimetype="text/plain",
content="this is a text",
content_uri="http://localhost/mydoc",
content_status=enums.ContentStatusEnum.WAIT,
)
responses.add(
responses.GET,
"http://localhost/mydoc",
body="loaded content",
status=200,
)
with openbulk(service.index_name, refresh=True) as actions:
actions.index_document(doc_a)
actions.index_document(doc_b)
bulk_response = {
"items": [
{"index": {"status": 500, "_id": doc_a.id}},
{
"index": {
"status": 400,
"_id": doc_b.id,
"error": {"reason": "Invalid arguments"},
}
},
],
}
with mock.patch.object(indexer.client, "bulk") as mock_client_bulk:
mock_client_bulk.return_value = bulk_response
errors = indexer.load_n_process_all()
assert [(e.id, e.message) for e in errors] == [
(doc_a.id, "Unknown error"),
(doc_b.id, "Invalid arguments"),
]
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert sorted((d.id, d.content_status) for d in indexed_docs) == sorted(
[
(doc_a.id, enums.ContentStatusEnum.WAIT),
(doc_b.id, enums.ContentStatusEnum.WAIT),
]
)
@pytest.mark.parametrize(
"status, mimetype, is_active , expected",
(
(
enums.ContentStatusEnum.READY,
"text/plain",
True,
{
"content": "initial content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.WAIT,
"text/plain",
True,
{
"content": "initial content",
"status": enums.ContentStatusEnum.WAIT,
},
),
(
enums.ContentStatusEnum.LOADED,
"text/plain",
True,
{
"content": "data:text/plain;initial content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.LOADED,
"application/pdf",
True,
{
"content": "data:application/pdf;initial content",
"status": enums.ContentStatusEnum.READY,
},
),
(
enums.ContentStatusEnum.LOADED,
"text/plain",
False,
{
"content": "initial content",
"status": enums.ContentStatusEnum.LOADED,
},
),
),
)
@responses.activate
def test_services_process_all(status, mimetype, is_active, expected):
"""Documents content should be processed if needed"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
mimetype=mimetype,
content_status=status,
content_uri="http://localhost/mydoc",
content="initial content",
is_active=is_active,
)
with openbulk(service.index_name, refresh=True) as actions:
for doc in docs:
actions.index_document(doc)
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert [d.content_status for d in indexed_docs] == [status] * 3
assert [d.is_active for d in indexed_docs] == [is_active] * 3
def _process(document):
document.content = f"data:{document.mimetype};{document.content}"
return document
indexer.processors = [_process]
errors = indexer.process_all()
assert len(errors) == 0
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert [d.content_status for d in indexed_docs] == [expected["status"]] * 3
assert [d.content for d in indexed_docs] == [expected["content"]] * 3
def test_services_process_all__errors():
"""Documents content should be processed if needed"""
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
pdf_doc = factories.IndexDocumentFactory(
mimetype="application/pdf",
content_status=enums.ContentStatusEnum.LOADED,
)
unknown_doc = factories.IndexDocumentFactory(
mimetype="application/unknown",
content_status=enums.ContentStatusEnum.LOADED,
)
with openbulk(service.index_name, refresh=True) as actions:
actions.index_document(pdf_doc)
actions.index_document(unknown_doc)
def _process(document):
if document.mimetype == "application/unknown":
raise ValueError("unsupported format")
return document
indexer.processors = [_process]
errors = indexer.process_all()
assert len(errors) == 1
assert sorted(((e.id, e.message) for e in errors)) == sorted(
[
(unknown_doc.id, "Unable to process content : unsupported format"),
]
)
indexed_docs = indexer.search_documents({"match_all": {}})
assert sorted((d.id, d.content_status) for d in indexed_docs) == sorted(
[
(pdf_doc.id, enums.ContentStatusEnum.READY),
(unknown_doc.id, enums.ContentStatusEnum.LOADED),
]
)
@pytest.mark.parametrize(
"status, embedding_model, is_active , expected",
(
(
enums.ContentStatusEnum.WAIT,
None,
True,
{
"status": enums.ContentStatusEnum.WAIT,
"embedding_model": None,
"embedding": None,
},
),
(
enums.ContentStatusEnum.LOADED,
None,
True,
{
"status": enums.ContentStatusEnum.LOADED,
"embedding_model": None,
"embedding": None,
},
),
(
enums.ContentStatusEnum.READY,
None,
True,
{
"status": enums.ContentStatusEnum.READY,
"embedding_model": "mymodel",
"embedding": albert_embedding_response.response["data"][0]["embedding"],
},
),
(
enums.ContentStatusEnum.READY,
None,
False,
{
"status": enums.ContentStatusEnum.READY,
"embedding_model": None,
"embedding": None,
},
),
(
enums.ContentStatusEnum.READY,
"mymodel",
True,
{
"status": enums.ContentStatusEnum.READY,
"embedding_model": "mymodel",
"embedding": albert_embedding_response.response["data"][0]["embedding"],
},
),
),
)
@responses.activate
def test_services_embed_all(settings, status, embedding_model, is_active, expected):
"""Documents content should be processed if needed"""
enable_hybrid_search(settings)
assert check_hybrid_search_enabled()
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(
3,
content_status=status,
content_uri="http://localhost/mydoc",
embedding_model=embedding_model,
is_active=is_active,
)
with openbulk(service.index_name, refresh=True) as actions:
for doc in docs:
actions.index_document(doc)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
json=albert_embedding_response.response,
status=200,
)
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert [d.content_status for d in indexed_docs] == [status] * 3
assert [d.is_active for d in indexed_docs] == [is_active] * 3
errors = indexer.embed_all(model_name="mymodel")
assert len(errors) == 0
indexed_docs = list(indexer.search_documents({"match_all": {}}))
assert [d.content_status for d in indexed_docs] == [expected["status"]] * 3
assert [d.embedding_model for d in indexed_docs] == [
expected["embedding_model"]
] * 3
assert [d.embedding for d in indexed_docs] == [expected["embedding"]] * 3
@responses.activate
def test_services_embed_all__errors(settings):
"""Documents content should be processed if needed"""
enable_hybrid_search(settings)
service = factories.ServiceFactory()
indexer = IndexerTaskService(service, force_refresh=True)
docs = factories.IndexDocumentFactory.create_batch(3)
with openbulk(service.index_name, refresh=True) as actions:
for doc in docs:
actions.index_document(doc)
responses.add(
responses.POST,
settings.EMBEDDING_API_PATH,
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 : Invalid request")
for d in docs
]
)
indexed_docs = indexer.search_documents({"match_all": {}})
assert sorted((d.id, d.embedding_model) for d in indexed_docs) == sorted(
[(d.id, None) for d in docs]
)
@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_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
enable_hybrid_search(settings)
service = factories.ServiceFactory()
waiting_doc = factories.IndexDocumentFactory(
content_status=enums.ContentStatusEnum.WAIT
)
loaded_doc = factories.IndexDocumentFactory(
content_status=enums.ContentStatusEnum.LOADED
)
ready_doc = factories.IndexDocumentFactory(
content_status=enums.ContentStatusEnum.READY
)
assert waiting_doc.is_waiting
assert loaded_doc.is_loaded
assert ready_doc.is_ready
with openbulk(service.index_name, refresh=True, raise_on_status=True) as actions:
actions.index_document(waiting_doc)
actions.index_document(loaded_doc)
actions.index_document(ready_doc)
dispatch_indexing_tasks(service, [waiting_doc, loaded_doc, ready_doc])
mock_load_n_process_all.assert_called()
mock_process_all.assert_called()
mock_embed_all.assert_called()
+6 -1
View File
@@ -42,6 +42,11 @@ def enable_hybrid_search(settings):
ensure_search_pipeline_exists()
def refresh_index(index_name):
"""Force refresh again so all changes are visible to search"""
opensearch.opensearch_client().indices.refresh(index=index_name)
def bulk_create_documents(document_payloads):
"""Create documents in bulk from payloads"""
return [
@@ -88,7 +93,7 @@ def prepare_index(index_name, documents: List):
bulk(opensearch.opensearch_client(), actions)
# Force refresh again so all changes are visible to search
opensearch.opensearch_client().indices.refresh(index=index_name)
refresh_index(index_name)
count = opensearch.opensearch_client().count(index=index_name)["count"]
assert count == len(documents), f"Expected {len(documents)}, got {count}"
+61 -58
View File
@@ -2,8 +2,7 @@
import logging
from django.conf import settings
from django.core.exceptions import SuspiciousOperation
from django.core.exceptions import BadRequest, SuspiciousOperation
from lasuite.oidc_resource_server.authentication import ResourceServerAuthentication
from lasuite.oidc_resource_server.mixins import ResourceServerMixin
@@ -13,15 +12,14 @@ from rest_framework.response import Response
from . import schemas
from .authentication import ServiceTokenAuthentication
from .models import Service, get_opensearch_index_name
from .models import IndexDocument, Service, get_opensearch_index_name
from .permissions import IsAuthAuthenticated
from .services.indexer_services import IndexerTaskService
from .services.opensearch import (
check_hybrid_search_enabled,
embed_document,
ensure_index_exists,
opensearch_client,
search,
)
from .tasks.indexer import dispatch_indexing_tasks
logger = logging.getLogger(__name__)
@@ -37,6 +35,22 @@ class IndexDocumentView(views.APIView):
authentication_classes = [ServiceTokenAuthentication]
permission_classes = [IsAuthAuthenticated]
def clean_document(self, data):
"""
Returns a IndexDocument from request data.
"""
try:
document = schemas.DocumentSchema(**data)
except PydanticValidationError as err:
raise BadRequest(
[
{key: error[key] for key in ("msg", "type", "loc")}
for error in err.errors()
]
) from err
return IndexDocument(**document.model_dump())
# pylint: disable=too-many-locals
def post(self, request, *args, **kwargs):
"""
@@ -86,78 +100,67 @@ class IndexDocumentView(views.APIView):
- Returns a list of results for all documents, with details of success and indexing
errors.
"""
index_name = request.auth.index_name
service = request.auth
opensearch_client_ = opensearch_client()
indexer = IndexerTaskService(service, client=opensearch_client_)
# Bulk indexing several documents
if isinstance(request.data, list):
# Bulk indexing several documents
is_valid = True
results = []
actions = []
has_errors = False
documents = []
# Parse request data and raise on any validation error
for i, document_data in enumerate(request.data):
try:
document = schemas.DocumentSchema(**document_data)
except PydanticValidationError as excpt:
errors = [
{key: error[key] for key in ("msg", "type", "loc")}
for error in excpt.errors()
]
results.append({"index": i, "status": "error", "errors": errors})
has_errors = True
else:
document_dict = {
**document.model_dump(),
"embedding": embed_document(document)
if check_hybrid_search_enabled()
else None,
"embedding_model": settings.EMBEDDING_API_MODEL_NAME
if check_hybrid_search_enabled()
else None,
}
_id = document_dict.pop("id")
actions.append({"index": {"_id": _id}})
actions.append(document_dict)
results.append({"index": i, "_id": _id, "status": "valid"})
doc = self.clean_document(document_data)
documents.append(doc)
results.append({"index": i, "status": "valid", "_id": str(doc.id)})
except BadRequest as e:
results.append({"index": i, "status": "error", "errors": e.args[0]})
is_valid = False
if has_errors:
if not is_valid:
return Response(results, status=status.HTTP_400_BAD_REQUEST)
# Build index if needed.
ensure_index_exists(index_name)
# Indexing all documents
errors = indexer.index(documents)
response = opensearch_client_.bulk(index=index_name, body=actions)
for i, item in enumerate(response["items"]):
if item["index"]["status"] != 201:
results[i]["status"] = "error"
results[i]["message"] = (
item["index"].get("error", {}).get("reason", "Unknown error")
)
# Dispatch deferred indexation tasks for the documents
dispatch_indexing_tasks(service, documents)
# Update error status of documents
errors = {e.id: e for e in errors}
for result in results:
error = errors.get(result["_id"])
if error is None:
result["status"] = "success"
else:
results[i]["status"] = "success"
result["status"] = "error"
result["message"] = error.message
return Response(results, status=status.HTTP_201_CREATED)
try:
document = self.clean_document(request.data)
except BadRequest as e:
return Response(e.args[0], status=status.HTTP_400_BAD_REQUEST)
# Indexing a single document
document = schemas.DocumentSchema(**request.data)
document_dict = {
**document.model_dump(),
"embedding": embed_document(document)
if check_hybrid_search_enabled()
else None,
"embedding_model": settings.EMBEDDING_API_MODEL_NAME
if check_hybrid_search_enabled()
else None,
}
_id = document_dict.pop("id")
errors = indexer.index((document,))
# Build index if needed.
ensure_index_exists(index_name)
if errors:
return Response(
{"status": "error", **errors[0]}, status=status.HTTP_400_BAD_REQUEST
)
opensearch_client_.index(index=index_name, body=document_dict, id=_id)
# Dispatch deferred indexation tasks for the document
dispatch_indexing_tasks(service, (document,))
return Response(
{"status": "created", "_id": _id}, status=status.HTTP_201_CREATED
{"status": "created", "_id": document.id}, status=status.HTTP_201_CREATED
)
@@ -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
@@ -1,8 +1,7 @@
# ruff: noqa: S311
"""create_demo management command"""
import logging
import random
import secrets
import time
from uuid import uuid4
@@ -131,10 +130,10 @@ def generate_document():
"content": "\n".join(fake.paragraphs(nb=5)),
"created_at": created_at,
"updated_at": updated_at,
"size": random.randint(0, 100 * 1024**2),
"size": secrets.randbelow(100 * 1024**2),
"users": [str(uuid4()) for _ in range(3)],
"groups": [slugify(fake.word()) for _ in range(3)],
"reach": random.choice(list(enums.ReachEnum)).value,
"reach": secrets.choice(list(enums.ReachEnum)).value,
}
@@ -151,29 +150,29 @@ def create_demo(stdout):
)
for service in services:
ensure_index_exists(service.name)
opensearch_client_.indices.refresh(index=service.name)
ensure_index_exists(service.index_name)
opensearch_client_.indices.refresh(index=service.index_name)
with Timeit(stdout, "Creating documents"):
actions = BulkIndexing(stdout)
for _ in range(defaults.NB_OBJECTS["documents"]):
service = random.choice(services)
service = secrets.choice(services)
document = generate_document()
actions.push(service.name, uuid4(), document)
actions.push(service.index_name, uuid4(), document)
actions.flush()
with Timeit(stdout, "Creating dev services"):
for conf in defaults.DEV_SERVICES:
service = factories.ServiceFactory(**conf)
ensure_index_exists(service.name)
opensearch_client_.indices.refresh(index=service.name)
ensure_index_exists(service.index_name)
opensearch_client_.indices.refresh(index=service.index_name)
# Check and report on indexed documents
total_indexed = 0
for service in services:
opensearch_client_.indices.refresh(index=service.name)
indexed = opensearch_client_.count(index=service.name)["count"]
stdout.write(f" - {service.name:s}: {indexed:d} documents")
opensearch_client_.indices.refresh(index=service.index_name)
indexed = opensearch_client_.count(index=service.index_name)["count"]
stdout.write(f" - {service.index_name:s}: {indexed:d} documents")
total_indexed += indexed
stdout.write(f" TOTAL: {total_indexed:d} documents")
+17
View File
@@ -247,6 +247,21 @@ class Base(Configuration):
default="find", environ_name="OPENSEARCH_INDEX_PREFIX", environ_prefix=None
)
# IndexerTaskService settings
INDEXER_FORCE_REFRESH = values.BooleanValue(
default=False, environ_name="INDEXER_FORCE_REFRESH", environ_prefix=None
)
INDEXER_DOWNLOAD_TIMEOUT = values.IntegerValue(
default=100, environ_name="INDEXER_DOWNLOAD_TIMEOUT", environ_prefix=None
)
INDEXER_TASK_COUNTDOWN = values.IntegerValue(
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.",
@@ -562,6 +577,8 @@ class Test(Base):
USE_SWAGGER = True
CELERY_TASK_ALWAYS_EAGER = values.BooleanValue(True)
INDEXER_FORCE_REFRESH = True
HYBRID_SEARCH_ENABLED = False
def __init__(self):
# pylint: disable=invalid-name