utils.py

"""

Contains Utility functions for LLM and Database module. Along with some other misllaneous functions.

"""

from turtle import clear
from pymupdf import pymupdf
#from docx import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
#import tiktoken
import base64
import hashlib
from typing import List
from openai import OpenAI
#from dotenv import load_dotenv
import os
import hashlib
from datetime import datetime
from typing import List, Optional, Dict, Any, Tuple

def generate_file_id(file_bytes: bytes) -> str:
    """Generate a 4-character unique file ID for given file."""
    hash_obj = hashlib.sha256()
    hash_obj.update(file_bytes[:4096])  # Still hash the first 4096 bytes
    # Take first 2 bytes (16 bits) and convert to base36 (alphanumeric)
    file_id = hex(int.from_bytes(hash_obj.digest()[:2], 'big'))[2:].zfill(4)
    return file_id


def process_pdf_to_chunks(

    pdf_content: bytes,

    file_name: str,

    chunk_size: int = 512,

    chunk_overlap: int = 20

) -> Tuple[List[Dict[str, Any]], str]:
    """

    Process PDF content into chunks with column layout detection and proper image handling

    """
    doc = pymupdf.open(stream=pdf_content, filetype="pdf")
    document_text = ""
    all_images = []
    image_positions = []
    char_to_page_map = []
    layout_info = {}
    
    doc_id = generate_file_id(pdf_content)

    def detect_columns(blocks):
        """Detect if page has multiple columns based on text block positions"""
        if not blocks:
            return 1
            
        x_positions = [block[0] for block in blocks]
        x_positions.sort()
        
        if len(x_positions) > 1:
            gaps = [x_positions[i+1] - x_positions[i] for i in range(len(x_positions)-1)]
            significant_gaps = [gap for gap in gaps if gap > page.rect.width * 0.15]
            return len(significant_gaps) + 1
        return 1

    def sort_blocks_by_position(blocks, num_columns):
        """Sort blocks by column and vertical position"""
        if num_columns == 1:
            return sorted(blocks, key=lambda b: b[0][1])  # b[0] is the bbox tuple, b[0][1] is y coordinate
        
        page_width = page.rect.width
        column_width = page_width / num_columns
        
        def get_column(block):
            bbox = block[0]  # Get the bounding box tuple
            x_coord = bbox[0]  # Get the x coordinate (first element)
            return int(x_coord // column_width)
            
        return sorted(blocks, key=lambda b: (get_column(b), b[0][1]))

    # Process each page
    for page_num, page in enumerate(doc, 1):
        blocks = page.get_text_blocks()
        images = page.get_images()
        
        # Detect layout
        num_columns = detect_columns(blocks)
        layout_info[page_num] = {
            "columns": num_columns,
            "width": page.rect.width,
            "height": page.rect.height
        }
        
        # Create elements list with both text and images
        elements = [(block[:4], block[4], "text") for block in blocks]
        
        # Add images to elements

        for img in images:
            try:
                img_rects = page.get_image_rects(img[0])
                if img_rects and len(img_rects) > 0:
                    img_bbox = img_rects[0]
                    if img_bbox:
                        img_data = (img_bbox, img[0], "image")
                        elements.append(img_data)
            except Exception as e:
                print(f"Error processing image: {e}")
                continue
        
        # Sort elements by position
        sorted_elements = sort_blocks_by_position(elements, num_columns)
        
        # Process elements in order
        page_text = ""
        for element in sorted_elements:
            if element[2] == "text":
                text_content = element[1]
                page_text += text_content
                char_to_page_map.extend([page_num] * len(text_content))
            else:
                xref = element[1]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                # Convert image bytes to base64
                image_base64 = base64.b64encode(image_bytes).decode('utf-8')
                all_images.append(image_base64)  # Store base64 encoded image
                
                image_marker = f"\n<img_{len(all_images)-1}>\n"
                image_positions.append((len(all_images)-1, len(document_text) + len(page_text)))
                page_text += image_marker
                char_to_page_map.extend([page_num] * len(image_marker))
        
        document_text += page_text

    # Create chunks
    splitter = RecursiveCharacterTextSplitter(
        #separators=["\n\n", "\n", " ", ""],
        #keep_separator=True
    ).from_tiktoken_encoder(
        encoding_name="cl100k_base",
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )
    
    text_chunks = splitter.split_text(document_text)
    
    # Process chunks with metadata
    processed_chunks = []
    for chunk_idx, chunk in enumerate(text_chunks):
        chunk_start = document_text.find(chunk)
        chunk_end = chunk_start + len(chunk)
        
        # Get page range and layout info
        chunk_pages = sorted(set(char_to_page_map[chunk_start:chunk_end]))
        chunk_layouts = {page: layout_info[page] for page in chunk_pages}
        
        # Get images for this chunk
        chunk_images = []
        for img_idx, img_pos in image_positions:
            if chunk_start <= img_pos <= chunk_end:
                chunk_images.append(all_images[img_idx])  # Already base64 encoded
        
        # Clean the chunk text
        #cleaned_chunk = clean_text_for_llm(chunk)

        chunk_dict = {
            "text": chunk,
            "metadata": {
                "created_date": datetime.now().isoformat(),
                "file_name": file_name,
                "images": chunk_images,
                "document_id": doc_id,
                "location": {
                    "char_start": chunk_start,
                    "char_end": chunk_end,
                    "pages": chunk_pages,
                    "chunk_index": chunk_idx,
                    "total_chunks": len(text_chunks),
                    "layout": chunk_layouts
                }
            }
        }
        processed_chunks.append(chunk_dict)
    
    return processed_chunks, doc_id



# import re
# import unicodedata
# from typing import Optional

# # Compile regex patterns once
# HTML_TAG_PATTERN = re.compile(r'<[^>]+>')
# MULTIPLE_NEWLINES = re.compile(r'\n\s*\n')
# MULTIPLE_SPACES = re.compile(r'\s+')

# def clean_text_for_llm(text: Optional[str]) -> str:
#     """
#     Efficiently clean and normalize text for LLM processing.
#     """
#     # Early returns
#     if not text:
#         return ""
#     if not isinstance(text, str):
#         try:
#             text = str(text)
#         except Exception:
#             return ""

#     # Single-pass character filtering
#     chars = []
#     prev_char = ''
#     space_pending = False
    
    # for char in text:
    #     # Skip null bytes and most control characters
    #     if char == '\0' or unicodedata.category(char).startswith('C'):
    #         if char not in '\n\t':
    #             continue
        
    #     # Convert escaped sequences
    #     if prev_char == '\\':
    #         if char == 'n':
    #             chars[-1] = '\n'
    #             continue
    #         if char == 't':
    #             chars[-1] = '\t'
    #             continue
            
    #     # Handle whitespace
    #     if char.isspace():
    #         if not space_pending:
    #             space_pending = True
    #         continue
            
    #     if space_pending:
    #         chars.append(' ')
    #         space_pending = False
            
    #     chars.append(char)
    #     prev_char = char

    # # Join characters and perform remaining operations
    # text = ''.join(chars)
    
    # # Remove HTML tags
    # #text = HTML_TAG_PATTERN.sub('', text)
    
    # # Normalize Unicode in a single pass
    # text = unicodedata.normalize('NFKC', text)
    
    # # Clean up newlines
    # text = MULTIPLE_NEWLINES.sub('\n', text)
    
    # Final trim
    # return text.strip()