Realise processing of HTML and improve both HTML & PDF processing giving new tenant information.

eveai_workers/tasks.py

@@ -1,7 +1,8 @@
 from datetime import datetime as dt, timezone as tz
 from flask import current_app
 from sqlalchemy.exc import SQLAlchemyError
-
+from celery import states
+from celery.exceptions import Ignore
 import os
 
 # Unstructured commercial client imports
@@ -18,6 +19,7 @@ from langchain_core.exceptions import LangChainException
 
 from common.utils.database import Database
 from common.models.document import DocumentVersion, EmbeddingMistral, EmbeddingSmallOpenAI
+from common.models.user import Tenant
 from common.extensions import db
 from common.utils.celery_utils import current_celery
 
@@ -25,14 +27,21 @@ from bs4 import BeautifulSoup
 
 
 @current_celery.task(name='create_embeddings', queue='embeddings')
-def create_embeddings(tenant_id, document_version_id, default_embedding_model):
+def create_embeddings(tenant_id, document_version_id):
     # Setup Remote Debugging only if PYCHARM_DEBUG=True
     if current_app.config['PYCHARM_DEBUG']:
         import pydevd_pycharm
         pydevd_pycharm.settrace('localhost', port=50170, stdoutToServer=True, stderrToServer=True)
 
-    current_app.logger.info(f'Creating embeddings for tenant {tenant_id} on document version {document_version_id} '
-                            f'with model {default_embedding_model}')
+    current_app.logger.info(f'Creating embeddings for tenant {tenant_id} on document version {document_version_id}.')
+
+    # Retrieve Tenant for which we are processing
+    tenant = Tenant.query.get(tenant_id)
+    if tenant is None:
+        current_app.logger.error(f'Cannot create embeddings for tenant {tenant_id}. '
+                                 f'Tenant not found')
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
 
     # Ensure we are working in the correct database schema
     Database(tenant_id).switch_schema()
@@ -42,7 +51,9 @@ def create_embeddings(tenant_id, document_version_id, default_embedding_model):
     if document_version is None:
         current_app.logger.error(f'Cannot create embeddings for tenant {tenant_id}. '
                                  f'Document version {document_version_id} not found')
-        return
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
+
     db.session.add(document_version)
 
     # start processing
@@ -52,29 +63,22 @@ def create_embeddings(tenant_id, document_version_id, default_embedding_model):
         db.session.commit()
     except SQLAlchemyError as e:
         current_app.logger.error(f'Error saving document version {document_version_id} to database '
-                                 f'for tenant {tenant_id} when creating embeddings. '
+                                 f'for tenant {tenant_id} when starting creating of embeddings. '
                                  f'error: {e}')
-        return
-
-    embed_provider = default_embedding_model.rsplit('.', 1)[0]
-    embed_model = default_embedding_model.rsplit('.', 1)[1]
-    # define embedding variables
-    embedding_function = None
-    match (embed_provider, embed_model):
-        case ('openai', 'text-embedding-3-small'):
-            embedding_function = embed_chunks_for_text_embedding_3_small
-        case ('mistral', 'mistral.mistral-embed'):
-            embedding_function = embed_chunks_for_mistral_embed
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
 
     match document_version.file_type:
         case 'pdf':
-            process_pdf(tenant_id, document_version, embedding_function, default_embedding_model)
+            process_pdf(tenant, document_version)
         case 'html':
-            process_html(tenant_id, document_version, embedding_function, default_embedding_model)
+            process_html(tenant, document_version)
         case _:
             current_app.logger.info(f'No functionality defined for file type {document_version.file_type} '
                                     f'for tenant {tenant_id} '
                                     f'while creating embeddings for document version {document_version_id}')
+            create_embeddings.update_state(state=states.FAILURE)
+            raise Ignore()
 
 
 @current_celery.task(name='ask_eve_ai', queue='llm_interactions')
@@ -83,7 +87,7 @@ def ask_eve_ai(query):
     pass
 
 
-def process_pdf(tenant_id, document_version, embedding_function, embedding_model):
+def process_pdf(tenant, document_version):
     file_path = os.path.join(current_app.config['UPLOAD_FOLDER'],
                              document_version.file_location,
                              document_version.file_name)
@@ -97,102 +101,162 @@ def process_pdf(tenant_id, document_version, embedding_function, embedding_model
                 coordinates=True,
                 extract_image_block_types=['Image', 'Table'],
                 chunking_strategy='by_title',
-                combine_under_n_chars=2000,
-                max_characters=3000,
+                combine_under_n_chars=current_app.config.get('MIN_CHUNK_SIZE'),
+                max_characters=current_app.config.get('MAX_CHUNK_SIZE'),
             )
-        try:
-            chunks = partition_doc_unstructured(tenant_id, document_version, req)
-            enriched_chunk_docs = enrich_chunks(tenant_id, document_version, chunks)
-            embeddings = embedding_function(tenant_id, document_version, enriched_chunk_docs)
-        except Exception as e:
-            current_app.logger.error(f'Unable to create Embeddings for tenant {tenant_id} '
-                                     f'on document version {document_version.id} '
-                                     f'with model {embedding_model} '
-                                     f'error: {e}')
-            raise
+    else:
+        current_app.logger.error(f'The physical file for document version {document_version.id} '
+                                 f'for tenant {tenant.id} '
+                                 f'at {file_path} does not exist')
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
 
-        # Save embeddings & processing information to the database
-        db.session.add_all(embeddings)
+    try:
+        chunks = partition_doc_unstructured(tenant, document_version, req)
+    except Exception as e:
+        current_app.logger.error(f'Unable to create Embeddings for tenant {tenant.id} '
+                                 f'while processing PDF on document version {document_version.id} '
+                                 f'error: {e}')
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
+
+    summary = summarize_chunk(tenant, document_version, chunks[0])
+    doc_lang = document_version.document_language
+    doc_lang.system_context = f'Summary: {summary}\n'
+    enriched_chunks = enrich_chunks(tenant, document_version, chunks)
+    embeddings = embed_chunks(tenant, document_version, enriched_chunks)
+
+    try:
+        db.session.add(doc_lang)
         db.session.add(document_version)
         document_version.processing_finished_at = dt.now(tz.utc)
         document_version.processing = False
-
-        try:
-            db.session.commit()
-        except SQLAlchemyError as e:
-            current_app.logger.error(f'Error saving embedding information for tenant {tenant_id} '
-                                     f'on document version {document_version.id}'
-                                     f'error: {e}')
-            db.session.rollback()
-            raise
-
-        current_app.logger.info(f'Embeddings created successfully for tenant {tenant_id} '
-                                f'on document version {document_version.id} :-)')
-    else:  # file exists
-        current_app.logger.error(f'The physical file for document version {document_version.id} '
-                                 f'at {file_path} does not exist')
+        db.session.add_all(embeddings)
+        db.session.commit()
+    except SQLAlchemyError as e:
+        current_app.logger.error(f'Error saving embedding information for tenant {tenant.id} '
+                                 f'on PDF, document version {document_version.id}'
+                                 f'error: {e}')
+        db.session.rollback()
         raise
 
+    current_app.logger.info(f'Embeddings created successfully for tenant {tenant.id} '
+                            f'on document version {document_version.id} :-)')
+
+
+def process_html(tenant, document_version):
+    # The tags to be considered can be dependent on the tenant
+    html_tags = tenant.html_tags
+    end_tags = tenant.html_end_tags
+    included_elements = tenant.html_included_elements
+    excluded_elements = tenant.html_excluded_elements
 
-def process_html(tenant_id, document_version, embedding_function, default_embedding_model):
     file_path = os.path.join(current_app.config['UPLOAD_FOLDER'],
                              document_version.file_location,
                              document_version.file_name)
+
     if os.path.exists(file_path):
         with open(file_path, 'rb') as f:
             html_content = f.read()
+    else:
+        current_app.logger.error(f'The physical file for document version {document_version.id} '
+                                 f'for tenant {tenant.id} '
+                                 f'at {file_path} does not exist')
+        create_embeddings.update_state(state=states.FAILURE)
+        raise Ignore()
+
+    extracted_data, title = parse_html(html_content, html_tags, included_elements=included_elements,
+                                       excluded_elements=excluded_elements)
+    potential_chunks = create_potential_chunks(extracted_data, end_tags)
+    chunks = combine_chunks(potential_chunks,
+                            current_app.config.get('MIN_CHUNK_SIZE'),
+                            current_app.config.get('MAX_CHUNK_SIZE')
+                            )
+    summary = summarize_chunk(tenant, document_version, chunks[0])
+    doc_lang = document_version.document_language
+    doc_lang.system_context = (f'Title: {title}\n'
+                               f'Summary: {summary}\n')
+
+    enriched_chunks = enrich_chunks(tenant, document_version, chunks)
+    embeddings = embed_chunks(tenant, document_version, enriched_chunks)
+
+    try:
+        db.session.add(doc_lang)
+        db.session.add(document_version)
+        document_version.processing_finished_at = dt.now(tz.utc)
+        document_version.processing = False
+        db.session.add_all(embeddings)
+        db.session.commit()
+    except SQLAlchemyError as e:
+        current_app.logger.error(f'Error saving embedding information for tenant {tenant.id} '
+                                 f'on HTML, document version {document_version.id}'
+                                 f'error: {e}')
+        raise
+
+    current_app.logger.info(f'Embeddings created successfully for tenant {tenant.id} '
+                            f'on document version {document_version.id} :-)')
 
 
+def enrich_chunks(tenant, document_version, chunks):
+    doc_lang = document_version.document_language
+    chunk_total_context = (f'Filename: {document_version.file_name}\n'
+                           f'{doc_lang.system_context}\n'
+                           f'User Context:\n{doc_lang.user_context}')
+    enriched_chunks = []
+    initial_chunk = f'Filename: {document_version.file_name}\n User Context:\n{doc_lang.user_context}\n{chunks[0]}'
+    enriched_chunks.append(initial_chunk)
+    for chunk in chunks[1:]:
+        enriched_chunk = f'{chunk_total_context}\n{chunk}'
+        enriched_chunks.append(enriched_chunk)
 
-def enrich_chunks(tenant_id, document_version, chunks):
-    # We're adding filename and a summary of the first chunk to all the chunks to create global context
-    # using openAI to summarise
-    api_key = current_app.config.get('OPENAI_API_KEY')
-    # TODO: model selection to be adapted to model approach
-    llm = ChatOpenAI(api_key=api_key, temperature=0, model='gpt-4-turbo')
+    return enriched_chunks
+
+
+def summarize_chunk(tenant, document_version, chunk):
+    llm_model = tenant.llm_model
+    llm_provider = llm_model.split('.', 1)[0]
+    llm_model = llm_model.split('.', 1)[1]
+
+    summary_template = ''
+    llm = None
+    match llm_provider:
+        case 'openai':
+            api_key = current_app.config.get('OPENAI_API_KEY')
+            llm = ChatOpenAI(api_key=api_key, temperature=0, model=llm_model)
+            match llm_model:
+                case 'gpt-4-turbo':
+                    summary_template = current_app.config.get('GPT4_SUMMARY_TEMPLATE')
+                case 'gpt-3.5-turbo':
+                    summary_template = current_app.config.get('GPT3_5_SUMMARY_TEMPLATE')
+                case _:
+                    current_app.logger.error(f'Error summarizing initial chunk for tenant {tenant.id} '
+                                             f'on document version {document_version.id} '
+                                             f'error: Invalid llm model')
+                    create_embeddings.update_state(state=states.FAILURE)
+                    raise Ignore()
+        case _:
+            current_app.logger.error(f'Error summarizing initial chunk for tenant {tenant.id} '
+                                     f'on document version {document_version.id} '
+                                     f'error: Invalid llm provider')
 
-    summary_template = current_app.config.get('GPT4_SUMMARY_TEMPLATE')
     prompt = ChatPromptTemplate.from_template(summary_template)
-
     chain = load_summarize_chain(llm, chain_type='stuff', prompt=prompt)
 
-    doc_creator = CharacterTextSplitter(chunk_size=9000, chunk_overlap=0)
-    text_to_summarize = doc_creator.create_documents(chunks[0])
+    doc_creator = CharacterTextSplitter(chunk_size=current_app.config.get('MAX_CHUNK_SIZE') * 2, chunk_overlap=0)
+    text_to_summarize = doc_creator.create_documents(chunk)
+
     try:
         summary = chain.run(text_to_summarize)
-        doc_lang = document_version.document_language
-        db.session.add(doc_lang)
-        doc_lang.system_context = f'Summary:\n {summary}'
-        try:
-            db.session.commit()
-        except SQLAlchemyError as e:
-            current_app.logger. error(f'Error saving summary to DocumentLanguage {doc_lang.id} '
-                                      f'while enriching chunks for tenant {tenant_id} '
-                                      f'on document version {document_version.id} '
-                                      f'error: {e}')
-            db.session.rollback()
-            raise
-
-        chunk_global_context = (f'Filename: {doc_lang.document.name}\n'
-                                f'User Context:\n{doc_lang.user_context}'
-                                f'System Context:\n{summary}')
-        enriched_chunks = []
-        initial_chunk = f'Filename: {document_version.file_name}\n User Context:\n{doc_lang.user_context}\n{chunks[0]}'
-        enriched_chunks.append(initial_chunk)
-        for chunk in chunks[1:]:
-            enriched_chunk = f'{chunk_global_context}\n{chunk}'
-            enriched_chunks.append(enriched_chunk)
-
-        return enriched_chunks
-
     except LangChainException as e:
-        current_app.logger.error(f'Error creating summary for chunk enrichment for tenant {tenant_id} '
+        current_app.logger.error(f'Error creating summary for chunk enrichment for tenant {tenant.id} '
                                  f'on document version {document_version.id} '
                                  f'error: {e}')
         raise
 
+    return summary
 
-def partition_doc_unstructured(tenant_id, document_version, unstructured_request):
+
+def partition_doc_unstructured(tenant, document_version, unstructured_request):
     # Initiate the connection to unstructured.io
     url = current_app.config.get('UNSTRUCTURED_FULL_URL')
     api_key = current_app.config.get('UNSTRUCTURED_API_KEY')
@@ -211,20 +275,41 @@ def partition_doc_unstructured(tenant_id, document_version, unstructured_request
                     chunks.append(el['metadata']['text_as_html'])
         return chunks
     except SDKError as e:
-        current_app.logger.error(f'Error creating embeddings for tenant {tenant_id} '
+        current_app.logger.error(f'Error creating embeddings for tenant {tenant.id} '
                                  f'on document version {document_version.id} while chuncking'
                                  f'error: {e}')
         raise
 
 
-def embed_chunks_for_text_embedding_3_small(tenant_id, document_version, chunks):
+def embed_chunks(tenant, document_version, chunks):
+    embedding_provider = tenant.embedding_model.rsplit('.', 1)[0]
+    embedding_model = tenant.embedding_model.rsplit('.', 1)[1]
+
+    match embedding_provider:
+        case 'openai':
+            match embedding_model:
+                case 'text-embedding-3-small':
+                    return embed_chunks_for_text_embedding_3_small(tenant, document_version, chunks)
+                case _:
+                    current_app.logger.error(f'Error creating embeddings for tenant {tenant.id} '
+                                             f'on document version {document_version.id} '
+                                             f'error: Invalid embedding model')
+                    create_embeddings.update_state(state=states.FAILURE)
+                    raise Ignore()
+        case _:
+            current_app.logger.error(f'Error creating embeddings for tenant {tenant.id} '
+                                     f'on document version {document_version.id} '
+                                     f'error: Invalid embedding provider')
+
+
+def embed_chunks_for_text_embedding_3_small(tenant, document_version, chunks):
     # Create embedding vectors using OpenAI
     api_key = current_app.config.get('OPENAI_API_KEY')
     embeddings_model = OpenAIEmbeddings(api_key=api_key, model='text-embedding-3-small')
     try:
         embeddings = embeddings_model.embed_documents(chunks)
     except LangChainException as e:
-        current_app.logger.error(f'Error creating embeddings for tenant {tenant_id} '
+        current_app.logger.error(f'Error creating embeddings for tenant {tenant.id} '
                                  f'on document version {document_version.id} while calling OpenAI API'
                                  f'error: {e}')
         raise
@@ -246,7 +331,7 @@ def embed_chunks_for_mistral_embed(tenant_id, document_version, chunks):
     pass
 
 
-def parse_html(html_content, included_elements=None, excluded_elements=None):
+def parse_html(html_content, tags, included_elements=None, excluded_elements=None):
     soup = BeautifulSoup(html_content, 'html.parser')
     extracted_content = []
 
@@ -263,4 +348,57 @@ def parse_html(html_content, included_elements=None, excluded_elements=None):
                 continue  # Skip this sub_element if it's within any of the excluded_elements
             extracted_content.append((sub_element.name, sub_element.get_text(strip=True)))
 
-    return extracted_content
+    title = soup.find('title').get_text(strip=True)
+
+    return extracted_content, title
+
+
+def create_potential_chunks(extracted_data, end_tags):
+    potential_chunks = []
+    current_chunk = []
+
+    for tag, text in extracted_data:
+        formatted_text = f"- {text}" if tag == 'li' else f"{text}\n"
+        if current_chunk and tag in end_tags and current_chunk[-1][0] in end_tags:
+            # Consecutive li and p elements stay together
+            current_chunk.append((tag, formatted_text))
+        else:
+            # End the current chunk if the last element was an end tag
+            if current_chunk and current_chunk[-1][0] in end_tags:
+                potential_chunks.append(current_chunk)
+                current_chunk = []
+            current_chunk.append((tag, formatted_text))
+
+    # Add the last chunk
+    if current_chunk:
+        potential_chunks.append(current_chunk)
+    return potential_chunks
+
+
+def combine_chunks(potential_chunks, min_chars, max_chars):
+    actual_chunks = []
+    current_chunk = ""
+    current_length = 0
+
+    for chunk in potential_chunks:
+        chunk_content = ''.join(text for _, text in chunk)
+        chunk_length = len(chunk_content)
+
+        if current_length + chunk_length > max_chars:
+            if current_length >= min_chars:
+                actual_chunks.append(current_chunk)
+                current_chunk = chunk_content
+                current_length = chunk_length
+            else:
+                # If the combined chunk is still less than max_chars, keep adding
+                current_chunk += chunk_content
+                current_length += chunk_length
+        else:
+            current_chunk += chunk_content
+            current_length += chunk_length
+
+    # Handle the last chunk
+    if current_chunk and current_length >= min_chars:
+        actual_chunks.append(current_chunk)
+
+    return actual_chunks