docp 0.1.0b1__py3-none-any.whl

docp/__init__.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+:Purpose:   This module provides the project initilisation logic.
+
+:Platform:  Linux/Windows | Python 3.10+
+:Developer: J Berendt
+:Email:     development@s3dev.uk
+
+:Comments:  n/a
+
+"""
+
+import os
+import sys
+sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)))
+
+# Bring entry-points to the surface.
+try:
+    from loaders.chroma import ChromaLoader
+except ImportError as err:
+    # The chroma loader requires a lot of backend which is not required for the parser.
+    msg = f'An error occurred while importing the Chroma loader:\n- {err}'
+    raise ImportError(msg) from err
+
+try:
+    from .parsers.pdfparser import PDFParser
+    from ._version import __version__
+except ImportError:
+    from parsers.pdfparser import PDFParser
+    from _version import __version__

docp/_version.py

@@ -0,0 +1 @@
+__version__ = '0.1.0b1'

docp/dbs/chroma.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+:Purpose:   This module provides a localised wrapper and specialised
+            functionality around the
+            ``langchain_community.vectorstores.Chroma`` class, for
+            interacting with a Chroma database.
+
+:Platform:  Linux/Windows | Python 3.10+
+:Developer: J Berendt
+:Email:     development@s3dev.uk
+
+:Comments:  n/a
+
+"""
+# pylint: disable=wrong-import-order
+
+import chromadb
+import os
+import torch
+from glob import glob
+from hashlib import md5
+from langchain_huggingface import HuggingFaceEmbeddings
+# langchain's Chroma is used rather than the base chromadb as it provides
+# the add_texts method which support GPU processing and parallelisation.
+from langchain_community.vectorstores import Chroma as _Chroma
+
+
+class ChromaDB(_Chroma):
+    """Wrapper class around the ``chromadb`` library.
+
+    Args:
+        path (str): Path to the chroma database's *directory*.
+        collection (str): Collection name.
+        offline (bool, optional): Remain offline, used the cached
+            embedding function model rather than obtaining one online.
+            Defaults to False.
+    """
+    # pylint: disable=line-too-long
+
+    _MODEL_CACHE = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), '.cache')
+    # Installing torch is a huge overhead, just for this. However, torch
+    # will already be installed as part of the sentence-transformers library,
+    # so we'll use it here.
+    _MODEL_KWARGS = {'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+    # TODO: Add this to a config file.
+    _MODEL_NAME = 'all-MiniLM-L6-v2'
+
+    def __init__(self, path: str, collection: str, offline: bool=False):
+        """Chroma database class initialiser."""
+        self._path = os.path.realpath(path)
+        self._cname = collection
+        self._offline = offline
+        self._client = None         # Database 'client' object
+        self._dbc = None            # Database 'collection' object.
+        self._set_client()
+        self._set_embedding_fn()
+        super().__init__(client=self._client,
+                          collection_name=self._cname,
+                          embedding_function=self._embfn,
+                          persist_directory=self._path)
+        self._set_collection()
+
+    @property
+    def client(self):
+        """Accessor to the :class:`chromadb.PersistentClient` class."""
+        return self._client
+
+    @property
+    def collection(self):
+        """Accessor to the chromadb client's collection object."""
+        return self._dbc
+
+    @property
+    def embedding_function(self):
+        """Accessor to the embedding function used."""
+        return self._embfn
+
+    @property
+    def path(self) -> str:
+        """Accessor to the database's path."""
+        return self._path
+
+    def add_documents(self, docs: list):
+        """Add multiple documents to the collection.
+
+        This method wraps ``Chroma.add_texts`` method which supports GPU
+        processing and parallelisation. The ID is derived locally from
+        the file's basename, page number and page content.
+
+        Args:
+            docs (list): A list of ``langchain_core.documents.base.Document``
+                document objects.
+
+        """
+        # This method overrides the base class' add_documents method.
+        # pylint: disable=arguments-differ
+        # pylint: disable=arguments-renamed
+        if not isinstance(docs, list):
+            docs = [docs]
+        ids_, docs_, meta_ = self._preproc(docs=docs)
+        self.add_texts(ids=ids_, texts=docs_, metadatas=meta_)
+
+    def show_all(self):
+        """Return the entire contents of the collection.
+
+        This is an alias around ``.collection.get()``.
+
+        """
+        return self._dbc.get()
+
+    def _get_embedding_function_model(self) -> str:
+        """Derive the path to the embedding function model.
+
+        :Note:
+            If ``offline=True`` was passed into the class constructor,
+            the model cache is used, if available - otherwise the user
+            is warned.
+
+            If online usage is allowed, the model is obtained by the
+            means defined by the embedding function constructor.
+
+        Returns:
+            str: The name of the model. Or, if offline, the path to the
+            model's cache to be passed into the embedding function
+            constructor is returned.
+
+        """
+        if self._offline:
+            if not os.path.exists(self._MODEL_CACHE):
+                os.makedirs(self._MODEL_CACHE)
+                msg = ('Offline mode has been chosen, yet the embedding function model cache does not exist. '
+                       'Therefore, a model must be downloaded. Please enable online usage for the first run '
+                       'so a model can be downloaded and stored into the cache for future (offline) use.')
+                raise FileNotFoundError(msg)
+            # Find the cache directory containing the named model, this enables offline use.
+            model_loc = os.path.commonpath(filter(lambda x: 'config.json' in x,
+                                                  glob(os.path.join(self._MODEL_CACHE,
+                                                                    f'*{self._MODEL_NAME}*',
+                                                                    '**'),
+                                                        recursive=True)))
+            return model_loc
+        return self._MODEL_NAME
+
+    @staticmethod
+    def _preproc(docs: list):
+        """Pre-process the document objects to create the IDs.
+
+        Parse the ``Document`` object into its parts for storage.
+        Additionally, create the ID as a hash of the source document's
+        basename, page number and content.
+
+        """
+        ids = []
+        txts = []
+        metas = []
+        for doc in docs:
+            pc = doc.page_content
+            m = doc.metadata
+            pc_, src_ = map(str.encode, (pc, m['source']))
+            pg_ = str(m.get('pageno', 0)).zfill(4)
+            id_ = f'id_{md5(src_).hexdigest()}_{pg_}_{md5(pc_).hexdigest()}'
+            ids.append(id_)
+            txts.append(pc)
+            metas.append(m)
+        return ids, txts, metas
+
+    def _set_client(self):
+        """Set the database client object."""
+        settings = chromadb.Settings(anonymized_telemetry=False)
+        self._client = chromadb.PersistentClient(path=self._path,
+                                                 settings=settings)
+
+    def _set_collection(self):
+        """Set the database collection object."""
+        self._dbc = self._client.get_or_create_collection(self._cname,
+                                                          metadata={'hnsw:space': 'cosine'})
+
+    def _set_embedding_fn(self):
+        """Set the embeddings function object."""
+        model_name = self._get_embedding_function_model()
+        self._embfn = HuggingFaceEmbeddings(model_name=model_name,
+                                            model_kwargs=self._MODEL_KWARGS,
+                                            cache_folder=self._MODEL_CACHE)

docp/loaders/_chromabaseloader.py

@@ -0,0 +1,362 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+:Purpose:   This module provides functionality to parse and store
+            document data into a Chroma vector database.
+
+:Platform:  Linux/Windows | Python 3.10+
+:Developer: J Berendt
+:Email:     development@s3dev.uk
+
+:Comments:  n/a
+
+"""
+# pylint: disable=no-name-in-module  # langchain.chains.RetrievalQA
+
+import contextlib
+import os
+import re
+from chromadb.api.types import errors as chromadberrors
+from langchain.chains import RetrievalQA
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from utils4.reporterror import reporterror
+from utils4.user_interface import ui
+# locals
+try:
+    from .dbs.chroma import ChromaDB
+    from .parsers.pdfparser import PDFParser
+except ImportError:
+    from dbs.chroma import ChromaDB
+    from parsers.pdfparser import PDFParser
+
+_PRE_ERR = '\n[ERROR]:'
+_PRE_WARN = '\n[WARNING]:'
+
+
+class Tools:
+    """General tools used for loading documents."""
+
+    @staticmethod
+    def parse_to_keywords(resp: str) -> list:
+        """Parse the bot's response into a list of keywords.
+
+        Args:
+            resp (str): Text response directly from the bot.
+
+        Returns:
+            list: A list of keywords extracted from the response,
+            separated by asterisks as bullet points.
+
+        """
+        # Capture asterisk bullet points or a numbered list.
+        rexp = re.compile(r'(?:\*|[0-9]+\.)\s*(.*)\n')
+        trans = {45: ' ', 47: ' '}
+        resp_ = resp.translate(trans).lower()
+        kwds = rexp.findall(resp_)
+        if kwds:
+            return ', '.join(kwds)
+        return ''
+
+
+class _ChromaBaseLoader:
+    """Base class for loading documents into a Chroma vector database.
+
+    Args:
+        path (str): Full path to the file to be parsed and loaded.
+        dbpath (str | Chroma): Either the full path to the Chroma database
+            *directory*, or an instance of a :class:`~dbs.chroma.Chroma`
+            database. If the instance is passed, the ``collection``
+            argument is ignored.
+        collection (str, optional): Name of the Chroma database
+            collection. Only required if the ``db`` parameter is a path.
+            Defaults to None.
+        offline (bool, optional): Remain offline and use the locally
+            cached embedding function model. Defaults to False.
+
+    """
+
+    _PARSERS = {'.pdf': PDFParser}
+
+    def __init__(self,
+                 dbpath: str | ChromaDB,
+                 collection: str=None,
+                 *,
+                 load_keywords: bool=False,
+                 llm: object=None,
+                 offline: bool=False):
+        """Chroma database class initialiser."""
+        self._dbpath = dbpath
+        self._cname = collection
+        self._load_keywords = load_keywords
+        self._llm = llm
+        self._offline = offline
+        self._dbo = None            # Database object.
+        self._docs = []             # List of 'Document' objects.
+        self._docss = []            # List of 'Document' objects *with splits*.
+        self._fbase = None          # Basename of the document currently being loaded.
+        self._fpath = None          # Full path to the document currently being loaded.
+        self._p = None              # Document parser object.
+        self._splitter = None       # Text splitter.
+        self._set_db_client()
+        self._check_parameters()
+
+    @property
+    def chroma(self):
+        """Accessor to the database client object."""
+        return self._dbo
+
+    @property
+    def parser(self):
+        """Accessor to the document parser object."""
+        return self._p
+
+    def _check_parameters(self) -> None:
+        """Verify the class parameters are viable.
+
+        Raises:
+            ValueError: If the ``load_keywords`` argument is True and the
+                ``llm`` argument is None, or the inverse. Both arguments
+                must either sum to 0, or 2.
+
+        """
+        if sum((self._load_keywords, self._llm is not None)) not in (0, 2):
+            raise ValueError('For keyword loading, the load_keywords argument '
+                             'must be True and a model instance must be provided.')
+
+    def _create_documents(self) -> bool:
+        """Convert each extracted page into a ``Document`` object.
+
+        Returns:
+            bool: True of the pages are loaded as ``Document`` objects
+            successfully. Otherwise False.
+
+        """
+        self._docs = [Document(page_content=page.content,
+                               metadata={'source': self._p.doc.basename,
+                                         'pageno': page.pageno})
+                      for page in self._p.doc.pages if page.hastext]
+        if not self._docs:
+            msg = f'{_PRE_WARN} Text could not be parsed from {self._p.doc.basename}.'
+            ui.print_warning(msg)
+            return False
+        return True
+
+    def _get_keywords(self) -> str:
+        """Query the document (using the LLM) to extract the keywords."""
+        # pylint: disable=line-too-long
+        print('- Extracting keywords ...')
+        qry = ('List the important keywords which can be used to summarize this '
+               f'document: "{self._fbase}". Use only phrases which are found in the document.')
+        # Suppress stdout.
+        with contextlib.redirect_stdout(None):
+            nids = len(self._dbo.get(where={'source': self._fbase})['ids'])
+            # Max of 50, min n records; prefer n records or 10%.
+            filter_ = {'k': min(nids, max(25, min(nids//10, 50))),
+                       'filter': {'source': {'$eq': self._fbase}}}
+            # TODO: Replace this with the module.cless.method once created.
+            qa = RetrievalQA.from_chain_type(llm=self._llm,
+                                             chain_type="stuff",
+                                             retriever=self._dbo.as_retriever(search_kwargs=filter_),
+                                             return_source_documents=True,
+                                             verbose=True)
+            resp = qa.invoke(qry)
+        kwds = Tools.parse_to_keywords(resp=resp['result'])
+        return kwds
+
+    def _load(self, path: str, **kwargs):
+        """Load the selected files into the vector store.
+
+        Args:
+            path (str): Full path to the file to be loaded.
+
+        :Keyword Arguments:
+            Those passed from the loader-specific ``load`` method.
+
+        """
+        # pylint: disable=multiple-statements
+        self._fpath = path
+        self._fbase = os.path.basename(path)
+        s = self._set_parser()
+        if s: s = self._set_text_splitter()
+        if s: s = self._parse_text(**kwargs)
+        if s: s = self._create_documents()
+        if s: s = self._split_texts()
+        if s: s = self._load_worker()
+        if s and self._load_keywords and self._llm:
+            kwds = self._get_keywords()
+            s = self._store_keywords(kwds=kwds)
+        self._print_summary(success=s)
+
+    def _load_worker(self) -> bool:
+        """Load the split documents into the database collection.
+
+        Returns:
+            bool: True if loaded successfully, otherwise False. Success
+            is based on the number of records after the load being
+            greater than the number of records before the load, or not
+            exceptions being raised.
+
+        """
+        try:
+            print('- Loading the document into the database ...')
+            nrecs_b = self._dbo.collection.count()  # Count records before.
+            self._dbo.add_documents(self._docss)
+            nrecs_a = self._dbo.collection.count()  # Count records after.
+            return self._test_load(nrecs_b=nrecs_b, nrecs_a=nrecs_a)
+        except chromadberrors.DuplicateIDError:
+            print('-- Document already loaded; duplicate detected.')
+            return False  # Prevent from loading keywords.
+        except Exception as err:
+            reporterror(err)
+        return False
+
+    def _parse_text(self, **kwargs) -> bool:
+        """Parse text from the document.
+
+        :Keyword Arguments:
+            Those to be passed into the text extraction method.
+
+        Returns:
+            bool: True if the parser's 'text' object is populated,
+            otherwise False.
+
+        """
+        print('- Extracting text ...')
+        self._p.extract_text(**kwargs)
+        if len(self._p.doc.pages) < 2:
+            ui.print_warning(f'No text extracted from {self._p.doc.basename}')
+            return False
+        return True
+
+    @staticmethod
+    def _print_summary(success: bool):
+        """Print an end of processing summary.
+
+        Args:
+            success (bool): Success flag from the processor.
+
+        """
+        if success:
+            print('Processing complete. Success.')
+        else:
+            print('Processing aborted due to error. Failure.')
+
+    def _set_db_client(self) -> bool:
+        """Set the database client object.
+
+        If the ``_db`` object is a string, this is inferred as the *path*
+        to the database. Otherwise, it is inferred as the database object
+        itself.
+
+        Returns:
+            bool: True if the database object is set without error.
+            Otherwise False.
+
+        """
+        try:
+            if isinstance(self._dbpath, str):
+                self._dbo = ChromaDB(path=self._dbpath,
+                                     collection=self._cname,
+                                     offline=self._offline)
+            else:
+                self._dbo = self._dbpath
+        except Exception as err:
+            reporterror(err)
+            return False
+        return True
+
+    def _set_parser(self) -> bool:
+        """Set the appropriate document parser.
+
+        :Rationale:
+            The parser is set by the file extension. For example, a file
+            extension ``.pdf`` will set the
+            :class:`parsers.pdfparser.PDFParser` class.
+
+        Returns:
+            bool: True if a file extension appropriate parser was found.
+            Otherwise, False.
+
+        """
+        # pylint: disable=invalid-name  # OK as the variable (Parser) is a class.
+        # TODO: Updated this to use the (not-yet-available) ispdf utility
+        #       function, rather than relying on the file extension.
+        ext = os.path.splitext(self._fpath)[1]
+        Parser = self._PARSERS.get(ext)
+        if not Parser:
+            msg = f'{_PRE_WARN} Document parser not set for {os.path.basename(self._fpath)}.'
+            ui.print_warning(msg)
+            return False
+        self._p = Parser(path=self._fpath)
+        return True
+
+    # TODO: Add these to a config file.
+    def _set_text_splitter(self) -> bool:
+        """Define the text splitter to be used.
+
+        Returns:
+            bool: True, always.
+
+        """
+        self._splitter = RecursiveCharacterTextSplitter(chunk_size=256,
+                                                        chunk_overlap=25,
+                                                        separators=['\n\n\n', '\n\n', '\n', ' '])
+        return True
+
+    def _split_texts(self) -> bool:
+        """Split the document text using a recursive text splitter.
+
+        Returns:
+            bool: True if the text was split successfully, otherwise
+            False.
+
+        """
+        self._docss = self._splitter.split_documents(self._docs)
+        if not self._docss:
+            msg = (f'{_PRE_ERR} An error occurred while splitting the documents for '
+                   f'{self._p.doc.basename}.')
+            ui.print_warning(msg)
+            return False
+        return True
+
+    def _store_keywords(self, kwds: str) -> bool:
+        """Store the extracted keywords into the keywords collection.
+
+        Args:
+            kwds (str): A string containing the keywords extracted from
+                the document.
+
+        Returns:
+            bool: True if loaded successfully, otherwise False.
+
+        """
+        print('- Storing keywords ...')
+        db = ChromaDB(path=self._dbo.path, collection=f'{self._cname}-kwds', offline=self._offline)
+        nrecs_b = db.collection.count()  # Count records before.
+        docs = [Document(page_content=kwds, metadata={'source': self._fbase})]
+        db.add_documents(docs)
+        nrecs_a = db.collection.count()  # Count records after.
+        return 1 == nrecs_a - nrecs_b
+
+    def _test_load(self, nrecs_b: int, nrecs_a: int) -> bool:
+        """Test the document was loaded successfully.
+
+        :Test:
+            - Given a count of records before the load, verify the number
+              of records after the load is equal to the number of records
+              before, plus the number of split documents.
+
+        Args:
+            nrecs_b (int): Number of records *before* the load.
+            nrecs_a (int): Number of records *after* the load.
+
+        Returns:
+            bool: True if the number of records before the load plus the
+            number is splits is equal to the number of records after the
+            load.
+
+        """
+        if nrecs_a == nrecs_b:
+            ui.print_warning(f'{_PRE_WARN} No new documents added. Possibly already loaded?')
+        return nrecs_a == nrecs_b + len(self._docss)