mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2025-04-02 13:45:10 -05:00
304 lines
8.7 KiB
Python
304 lines
8.7 KiB
Python
import logging
|
|
import os
|
|
from contextlib import contextmanager
|
|
|
|
import math
|
|
from dateutil.parser import isoparse
|
|
from django.conf import settings
|
|
from whoosh import highlight, classify, query
|
|
from whoosh.fields import Schema, TEXT, NUMERIC, KEYWORD, DATETIME, BOOLEAN
|
|
from whoosh.highlight import HtmlFormatter
|
|
from whoosh.index import create_in, exists_in, open_dir
|
|
from whoosh.qparser import MultifieldParser
|
|
from whoosh.qparser.dateparse import DateParserPlugin
|
|
from whoosh.searching import ResultsPage, Searcher
|
|
from whoosh.writing import AsyncWriter
|
|
|
|
from documents.models import Document
|
|
|
|
logger = logging.getLogger("paperless.index")
|
|
|
|
|
|
def get_schema():
|
|
return Schema(
|
|
id=NUMERIC(
|
|
stored=True,
|
|
unique=True
|
|
),
|
|
title=TEXT(
|
|
sortable=True
|
|
),
|
|
content=TEXT(),
|
|
asn=NUMERIC(
|
|
sortable=True
|
|
),
|
|
|
|
correspondent=TEXT(
|
|
sortable=True
|
|
),
|
|
correspondent_id=NUMERIC(),
|
|
has_correspondent=BOOLEAN(),
|
|
|
|
tag=KEYWORD(
|
|
commas=True,
|
|
scorable=True,
|
|
lowercase=True
|
|
),
|
|
tag_id=KEYWORD(
|
|
commas=True,
|
|
scorable=True
|
|
),
|
|
has_tag=BOOLEAN(),
|
|
|
|
type=TEXT(
|
|
sortable=True
|
|
),
|
|
type_id=NUMERIC(),
|
|
has_type=BOOLEAN(),
|
|
|
|
created=DATETIME(
|
|
sortable=True
|
|
),
|
|
modified=DATETIME(
|
|
sortable=True
|
|
),
|
|
added=DATETIME(
|
|
sortable=True
|
|
),
|
|
|
|
)
|
|
|
|
|
|
def open_index(recreate=False):
|
|
try:
|
|
if exists_in(settings.INDEX_DIR) and not recreate:
|
|
return open_dir(settings.INDEX_DIR, schema=get_schema())
|
|
except Exception:
|
|
logger.exception(f"Error while opening the index, recreating.")
|
|
|
|
if not os.path.isdir(settings.INDEX_DIR):
|
|
os.makedirs(settings.INDEX_DIR, exist_ok=True)
|
|
return create_in(settings.INDEX_DIR, get_schema())
|
|
|
|
|
|
@contextmanager
|
|
def open_index_writer(optimize=False):
|
|
writer = AsyncWriter(open_index())
|
|
|
|
try:
|
|
yield writer
|
|
except Exception as e:
|
|
logger.exception(str(e))
|
|
writer.cancel()
|
|
finally:
|
|
writer.commit(optimize=optimize)
|
|
|
|
|
|
@contextmanager
|
|
def open_index_searcher():
|
|
searcher = open_index().searcher()
|
|
|
|
try:
|
|
yield searcher
|
|
finally:
|
|
searcher.close()
|
|
|
|
|
|
def update_document(writer, doc):
|
|
tags = ",".join([t.name for t in doc.tags.all()])
|
|
tags_ids = ",".join([str(t.id) for t in doc.tags.all()])
|
|
writer.update_document(
|
|
id=doc.pk,
|
|
title=doc.title,
|
|
content=doc.content,
|
|
correspondent=doc.correspondent.name if doc.correspondent else None,
|
|
correspondent_id=doc.correspondent.id if doc.correspondent else None,
|
|
has_correspondent=doc.correspondent is not None,
|
|
tag=tags if tags else None,
|
|
tag_id=tags_ids if tags_ids else None,
|
|
has_tag=len(tags) > 0,
|
|
type=doc.document_type.name if doc.document_type else None,
|
|
type_id=doc.document_type.id if doc.document_type else None,
|
|
has_type=doc.document_type is not None,
|
|
created=doc.created,
|
|
added=doc.added,
|
|
asn=doc.archive_serial_number,
|
|
modified=doc.modified,
|
|
)
|
|
|
|
|
|
def remove_document(writer, doc):
|
|
remove_document_by_id(writer, doc.pk)
|
|
|
|
|
|
def remove_document_by_id(writer, doc_id):
|
|
writer.delete_by_term('id', doc_id)
|
|
|
|
|
|
def add_or_update_document(document):
|
|
with open_index_writer() as writer:
|
|
update_document(writer, document)
|
|
|
|
|
|
def remove_document_from_index(document):
|
|
with open_index_writer() as writer:
|
|
remove_document(writer, document)
|
|
|
|
|
|
class DelayedQuery:
|
|
|
|
def _get_query(self):
|
|
raise NotImplementedError()
|
|
|
|
def _get_query_filter(self):
|
|
criterias = []
|
|
for k, v in self.query_params.items():
|
|
if k == 'correspondent__id':
|
|
criterias.append(query.Term('correspondent_id', v))
|
|
elif k == 'tags__id__all':
|
|
for tag_id in v.split(","):
|
|
criterias.append(query.Term('tag_id', tag_id))
|
|
elif k == 'document_type__id':
|
|
criterias.append(query.Term('type_id', v))
|
|
elif k == 'correspondent__isnull':
|
|
criterias.append(query.Term("has_correspondent", v == "false"))
|
|
elif k == 'is_tagged':
|
|
criterias.append(query.Term("has_tag", v == "true"))
|
|
elif k == 'document_type__isnull':
|
|
criterias.append(query.Term("has_type", v == "false"))
|
|
elif k == 'created__date__lt':
|
|
criterias.append(
|
|
query.DateRange("created", start=None, end=isoparse(v)))
|
|
elif k == 'created__date__gt':
|
|
criterias.append(
|
|
query.DateRange("created", start=isoparse(v), end=None))
|
|
elif k == 'added__date__gt':
|
|
criterias.append(
|
|
query.DateRange("added", start=isoparse(v), end=None))
|
|
elif k == 'added__date__lt':
|
|
criterias.append(
|
|
query.DateRange("added", start=None, end=isoparse(v)))
|
|
if len(criterias) > 0:
|
|
return query.And(criterias)
|
|
else:
|
|
return None
|
|
|
|
def _get_query_sortedby(self):
|
|
if 'ordering' not in self.query_params:
|
|
return None, False
|
|
|
|
field: str = self.query_params['ordering']
|
|
|
|
sort_fields_map = {
|
|
"created": "created",
|
|
"modified": "modified",
|
|
"added": "added",
|
|
"title": "title",
|
|
"correspondent__name": "correspondent",
|
|
"document_type__name": "type",
|
|
"archive_serial_number": "asn"
|
|
}
|
|
|
|
if field.startswith('-'):
|
|
field = field[1:]
|
|
reverse = True
|
|
else:
|
|
reverse = False
|
|
|
|
if field not in sort_fields_map:
|
|
return None, False
|
|
else:
|
|
return sort_fields_map[field], reverse
|
|
|
|
def __init__(self, searcher: Searcher, query_params, page_size):
|
|
self.searcher = searcher
|
|
self.query_params = query_params
|
|
self.page_size = page_size
|
|
self.saved_results = dict()
|
|
self.first_score = None
|
|
|
|
def __len__(self):
|
|
page = self[0:1]
|
|
return len(page)
|
|
|
|
def __getitem__(self, item):
|
|
if item.start in self.saved_results:
|
|
return self.saved_results[item.start]
|
|
|
|
q, mask = self._get_query()
|
|
sortedby, reverse = self._get_query_sortedby()
|
|
|
|
page: ResultsPage = self.searcher.search_page(
|
|
q,
|
|
mask=mask,
|
|
filter=self._get_query_filter(),
|
|
pagenum=math.floor(item.start / self.page_size) + 1,
|
|
pagelen=self.page_size,
|
|
sortedby=sortedby,
|
|
reverse=reverse
|
|
)
|
|
page.results.fragmenter = highlight.ContextFragmenter(
|
|
surround=50)
|
|
page.results.formatter = HtmlFormatter(tagname="span", between=" ... ")
|
|
|
|
if (not self.first_score and
|
|
len(page.results) > 0 and
|
|
sortedby is None):
|
|
self.first_score = page.results[0].score
|
|
|
|
page.results.top_n = list(map(
|
|
lambda hit: (
|
|
(hit[0] / self.first_score) if self.first_score else None,
|
|
hit[1]
|
|
),
|
|
page.results.top_n
|
|
))
|
|
|
|
self.saved_results[item.start] = page
|
|
|
|
return page
|
|
|
|
|
|
class DelayedFullTextQuery(DelayedQuery):
|
|
|
|
def _get_query(self):
|
|
q_str = self.query_params['query']
|
|
qp = MultifieldParser(
|
|
["content", "title", "correspondent", "tag", "type"],
|
|
self.searcher.ixreader.schema)
|
|
qp.add_plugin(DateParserPlugin())
|
|
q = qp.parse(q_str)
|
|
|
|
corrected = self.searcher.correct_query(q, q_str)
|
|
if corrected.query != q:
|
|
corrected_query = corrected.string
|
|
|
|
return q, None
|
|
|
|
|
|
class DelayedMoreLikeThisQuery(DelayedQuery):
|
|
|
|
def _get_query(self):
|
|
more_like_doc_id = int(self.query_params['more_like_id'])
|
|
content = Document.objects.get(id=more_like_doc_id).content
|
|
|
|
docnum = self.searcher.document_number(id=more_like_doc_id)
|
|
kts = self.searcher.key_terms_from_text(
|
|
'content', content, numterms=20,
|
|
model=classify.Bo1Model, normalize=False)
|
|
q = query.Or(
|
|
[query.Term('content', word, boost=weight)
|
|
for word, weight in kts])
|
|
mask = {docnum}
|
|
|
|
return q, mask
|
|
|
|
|
|
def autocomplete(ix, term, limit=10):
|
|
with ix.reader() as reader:
|
|
terms = []
|
|
for (score, t) in reader.most_distinctive_terms(
|
|
"content", number=limit, prefix=term.lower()):
|
|
terms.append(t)
|
|
return terms
|