Spaces:

IOE-IFAD
/

test

Paused

test / _backend /utils /document_management.py

Automated build & deploy

e454383 verified about 1 year ago

10.5 kB

import pandas as pd
import os
from utils.text_splitter import srt_chunker as chun
from utils.vector_db import VectorDB
import fitz
from datetime import datetime
import re

def load_document(file_path, raw_text=False):
    file_extension = os.path.splitext(file_path)[1].lower()
    if file_extension == '.pdf':
        return load_pdf(file_path)
    elif file_extension == '.csv':
        return load_csv(file_path)
    elif file_extension == '.srt':
        with open(file_path, 'r', encoding='utf-8') as file:
            content = file.read()
        paragraphs = content.split('\n\n')
        documents = [{"text": paragraph, "paragraph_number": index + 1}
                     for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
        return documents
    elif file_extension == '.md':
        if raw_text:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            return content
        else:
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            paragraphs = content.split('\n\n')
            documents = [{"text": paragraph, "paragraph_number": index + 1}
                         for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
            return documents
    else:
        raise ValueError(f"Unsupported file type: {file_extension}")

def load_pdf(file_path):
    """
    Load a PDF file and extract its text content into paragraphs.
    Args:
        file_path (str): The path to the PDF file.
    Returns:
        list[dict]: A list of dictionaries, each containing a paragraph of text and its paragraph number.
    """
    paragraphs = []
    pdf_document = fitz.open(file_path)

    for page in pdf_document:
        output = page.get_text("blocks")
        previous_block_id = None

        for block in output:
            if block[6] == 0:  # Check if the block is a text block
                cleaned_output = block[4].strip().replace("\n", " ")
                if cleaned_output and len(cleaned_output) >= 5:
                    if previous_block_id != block[5]:
                        paragraphs.append(cleaned_output)
                    else:
                        paragraphs[-1] += " " + cleaned_output
                    previous_block_id = block[5]

    # Adding the paragraph number to each paragraph
    documents = [{"text": paragraph, "paragraph_number": index + 1}
                 for index, paragraph in enumerate(paragraphs)]

    return documents


def load_csv(file_path):
    try:
        df = pd.read_csv(file_path)
        documents = []
        if "Paragraph text" in df.columns:
            df = df.rename(columns={"Paragraph text": "text"})
            df["page_number"] = df.index + 1
            documents = df.to_dict(orient="records")
            return documents
        else:
            for index, row in df.iterrows():
                content = row.to_json()
                documents.append({"text": content, "page_number": (index+1)})
            return documents
    except Exception as e:
        print(e)
        return []


def upload_document_to_vectorstore(chunks: list[dict], vector_db: VectorDB):
    """
    Uploads the chunks and their metadata to the vector database in batches.
    Args:
        chunks (list[dict]): A list of dictionaries, each containing a chunk of the document.
        vector_db (VectorDB): An instance of the VectorDB class to interact with the vector database.
    Returns:
        list: A list of points created in the vector store.
    """
    batch_size = 2000
    qdrant_batch_size = 500
    print(f"Uploading {len(chunks)} chunks to the vector store.")
    embeddings = []
    for batch_start in range(0, len(chunks), batch_size):
        batch_chunks = chunks[batch_start: batch_start + batch_size]
        batch_chunks = [chunk["text"] for chunk in batch_chunks]
        print(f"Encoding batch of {len(batch_chunks)} chunks.")
        batch_embeddings = vector_db.batch_encode_text(batch_chunks)

        embeddings.extend(batch_embeddings)

    points = vector_db.create_points(embeddings, chunks)
    print(f"Uploading {len(points)} points to the vector store.")

    for batch_start in range(0, len(points), qdrant_batch_size):
        batch_points = points[batch_start: batch_start + qdrant_batch_size]
        vector_db.insert_data(batch_points)
    print("Upload complete.")
    return points


def document_chunker(document: list[dict], chunking_function) -> list[dict]:
    """
    Chunk a document into smaller pieces using a provided chunking function.
    Args:
        document (list[dict]): A list of dictionaries, each containing a 'text' key with the document content.
        chunking_function (callable): A function that takes a string and returns a list of chunks.
    Returns:
        list[dict]: A list of dictionaries, each containing a chunk of the document.
    """
    chunks = []
    for page in document:
        text = page.pop("text")
        text_chunks = chunking_function(text)
        for chunk in text_chunks:
            chunks.append({"text": chunk, **page})

    return chunks


def transcript_chunker(document: list[dict], filename) -> list[dict]:
    """
    Chunk a transcript document into smaller pieces.
    Args:
        document (list[dict]): A list of dictionaries, each containing a 'text' key with the transcript content.
    Returns:
        list[dict]: A list of dictionaries, each containing a chunk of the transcript.
    """

    chunks = []

    def transcript_formatter(
        transcript): return f"{transcript['timestamp']} {transcript['speaker']}: {transcript['text']}"
    count = 0
    for page in document:
        text = page.pop("text")
        transcript = chun(text)

        if transcript is None:
            if "File name" in text:
                continue
            if chunks:
                chunks[-1]["text"] += text
            continue

        # Add the filename to the transcript
        transcript["filename"] = filename

        # Adding the index of the transcript to the transcript
        transcript["index"] = count
        count += 1

        # Add the current transcript to the chunks
        chunks.append(transcript)

        # Add the previous message to the current message
        if len(chunks) > 1:
            chunks[-1]["previous_message"] = transcript_formatter(chunks[-2])

        # Add the next message to the previous message
        if len(chunks) > 1:
            chunks[-2]["next_message"] = transcript_formatter(chunks[-1])

    return chunks

def merge_chunks(chunks: list[dict]) -> list[dict]:
    """
    Merge consecutive chunks that have the same speaker
    Args:
        chunks (list[dict]): A list of dictionaries, each containing a chunk of text and its metadata.
    Returns:
        list[dict]: A list of merged chunks.
    """
    if not chunks:
        return []

    merged_chunks = []
    current_chunk = chunks[0]
    for chunk in chunks[1:]:
        if chunk["speaker"] == current_chunk["speaker"]:
            # Merge the text of the current chunk with the next one
            current_chunk["text"] += " " + chunk["text"]
        else:
            # Add the current chunk to the merged list and start a new one
            merged_chunks.append(current_chunk)
            current_chunk = chunk
    # Add the last chunk
    merged_chunks.append(current_chunk)
    return merged_chunks


def generate_filename_from_details(minute_block_md: str) -> str:
    """
    Generates a filename from a minute block by extracting details 
    from a line like "| 24 March 2025, Guoqi Wu, AVP CSD | |".
    Format: YYYY-MM-DD_Name1Name2_RolePart1_RolePart2_CLE.md
    """
    lines = minute_block_md.splitlines()
    details_str = None

    if len(lines) > 1:
        # The line containing details for the filename is typically the second line of the block
        # e.g., | 24 March 2025, Guoqi Wu, AVP CSD | |
        header_info_line = lines[1] 
        
        # Extract content from between the first and second pipe characters
        line_parts = header_info_line.split('|')
        if len(line_parts) > 2 and line_parts[1].strip():
            details_str = line_parts[1].strip()
        else:
            raise ValueError(f"Could not parse details for filename from line: '{header_info_line}' in minute block.")
    else:
        raise ValueError("Minute block has insufficient lines to extract details for filename.")

    if not details_str:
        raise ValueError("Failed to extract details string from minute block.")

    # Regex to separate date from the rest of the details
    match = re.match(r"(\d{1,2}\s+[A-Za-z]+\s+\d{4}),\s*(.+)", details_str)
    if not match:
        raise ValueError(f"Details string '{details_str}' does not match expected format 'Date, Details...'")

    date_part_str = match.group(1) 
    name_and_role_parts_str = match.group(2)

    # Convert date to YYYY-MM-DD
    try:
        dt_object = datetime.strptime(date_part_str, "%d %B %Y")
        formatted_date = dt_object.strftime("%Y-%m-%d")
    except ValueError:
        try:
            dt_object = datetime.strptime(date_part_str, "%d %b %Y")
            formatted_date = dt_object.strftime("%Y-%m-%d")
        except ValueError as e:
            raise ValueError(f"Date string '{date_part_str}' could not be parsed: {e}")

    parts = [part.replace(" ", "") for part in name_and_role_parts_str.split(", ")]
    combined_name_role = "_".join(parts)
    combined_name_role_sanitized = re.sub(r'[^\w-]', '', combined_name_role.replace("(", "").replace(")", ""))

    return f"{formatted_date}_{combined_name_role_sanitized}_CLE.md", (formatted_date, combined_name_role_sanitized)

def load_prompt(prompt_file_name:str, working_directory) -> str:
    """
    Load a prompt from a file in the specified working directory's prompts folder.
    Args:
        prompt_file_name (str): The name of the prompt file to load.
        working_directory (str): The directory where the prompt file is located.
    Returns:
        str: The content of the prompt file.
    """
    prompt_file_path = os.path.join(working_directory, "prompts", prompt_file_name+".txt")
    if not os.path.exists(prompt_file_path):
        raise FileNotFoundError(f"Prompt file '{prompt_file_name}' at path '{prompt_file_path}' does not exist.")

    with open(prompt_file_path, 'r', encoding='utf-8') as file:
        return file.read()
    

	import pandas as pd
	import os
	from utils.text_splitter import srt_chunker as chun
	from utils.vector_db import VectorDB
	import fitz
	from datetime import datetime
	import re

	def load_document(file_path, raw_text=False):
	file_extension = os.path.splitext(file_path)[1].lower()
	if file_extension == '.pdf':
	return load_pdf(file_path)
	elif file_extension == '.csv':
	return load_csv(file_path)
	elif file_extension == '.srt':
	with open(file_path, 'r', encoding='utf-8') as file:
	content = file.read()
	paragraphs = content.split('\n\n')
	documents = [{"text": paragraph, "paragraph_number": index + 1}
	for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
	return documents
	elif file_extension == '.md':
	if raw_text:
	with open(file_path, 'r', encoding='utf-8') as file:
	content = file.read()
	return content
	else:
	with open(file_path, 'r', encoding='utf-8') as file:
	content = file.read()
	paragraphs = content.split('\n\n')
	documents = [{"text": paragraph, "paragraph_number": index + 1}
	for index, paragraph in enumerate(paragraphs) if paragraph.strip()]
	return documents
	else:
	raise ValueError(f"Unsupported file type: {file_extension}")

	def load_pdf(file_path):
	"""
	Load a PDF file and extract its text content into paragraphs.

	Args:
	file_path (str): The path to the PDF file.

	Returns:
	list[dict]: A list of dictionaries, each containing a paragraph of text and its paragraph number.
	"""
	paragraphs = []
	pdf_document = fitz.open(file_path)

	for page in pdf_document:
	output = page.get_text("blocks")
	previous_block_id = None

	for block in output:
	if block[6] == 0: # Check if the block is a text block
	cleaned_output = block[4].strip().replace("\n", " ")
	if cleaned_output and len(cleaned_output) >= 5:
	if previous_block_id != block[5]:
	paragraphs.append(cleaned_output)
	else:
	paragraphs[-1] += " " + cleaned_output
	previous_block_id = block[5]

	# Adding the paragraph number to each paragraph
	documents = [{"text": paragraph, "paragraph_number": index + 1}
	for index, paragraph in enumerate(paragraphs)]

	return documents


	def load_csv(file_path):
	try:
	df = pd.read_csv(file_path)
	documents = []
	if "Paragraph text" in df.columns:
	df = df.rename(columns={"Paragraph text": "text"})
	df["page_number"] = df.index + 1
	documents = df.to_dict(orient="records")
	return documents
	else:
	for index, row in df.iterrows():
	content = row.to_json()
	documents.append({"text": content, "page_number": (index+1)})
	return documents
	except Exception as e:
	print(e)
	return []


	def upload_document_to_vectorstore(chunks: list[dict], vector_db: VectorDB):
	"""
	Uploads the chunks and their metadata to the vector database in batches.

	Args:
	chunks (list[dict]): A list of dictionaries, each containing a chunk of the document.
	vector_db (VectorDB): An instance of the VectorDB class to interact with the vector database.

	Returns:
	list: A list of points created in the vector store.
	"""
	batch_size = 2000
	qdrant_batch_size = 500
	print(f"Uploading {len(chunks)} chunks to the vector store.")
	embeddings = []
	for batch_start in range(0, len(chunks), batch_size):
	batch_chunks = chunks[batch_start: batch_start + batch_size]
	batch_chunks = [chunk["text"] for chunk in batch_chunks]
	print(f"Encoding batch of {len(batch_chunks)} chunks.")
	batch_embeddings = vector_db.batch_encode_text(batch_chunks)

	embeddings.extend(batch_embeddings)

	points = vector_db.create_points(embeddings, chunks)
	print(f"Uploading {len(points)} points to the vector store.")

	for batch_start in range(0, len(points), qdrant_batch_size):
	batch_points = points[batch_start: batch_start + qdrant_batch_size]
	vector_db.insert_data(batch_points)
	print("Upload complete.")
	return points


	def document_chunker(document: list[dict], chunking_function) -> list[dict]:
	"""
	Chunk a document into smaller pieces using a provided chunking function.

	Args:
	document (list[dict]): A list of dictionaries, each containing a 'text' key with the document content.
	chunking_function (callable): A function that takes a string and returns a list of chunks.

	Returns:
	list[dict]: A list of dictionaries, each containing a chunk of the document.
	"""
	chunks = []
	for page in document:
	text = page.pop("text")
	text_chunks = chunking_function(text)
	for chunk in text_chunks:
	chunks.append({"text": chunk, **page})

	return chunks


	def transcript_chunker(document: list[dict], filename) -> list[dict]:
	"""
	Chunk a transcript document into smaller pieces.

	Args:
	document (list[dict]): A list of dictionaries, each containing a 'text' key with the transcript content.

	Returns:
	list[dict]: A list of dictionaries, each containing a chunk of the transcript.
	"""

	chunks = []

	def transcript_formatter(
	transcript): return f"{transcript['timestamp']} {transcript['speaker']}: {transcript['text']}"
	count = 0
	for page in document:
	text = page.pop("text")
	transcript = chun(text)

	if transcript is None:
	if "File name" in text:
	continue
	if chunks:
	chunks[-1]["text"] += text
	continue

	# Add the filename to the transcript
	transcript["filename"] = filename

	# Adding the index of the transcript to the transcript
	transcript["index"] = count
	count += 1

	# Add the current transcript to the chunks
	chunks.append(transcript)

	# Add the previous message to the current message
	if len(chunks) > 1:
	chunks[-1]["previous_message"] = transcript_formatter(chunks[-2])

	# Add the next message to the previous message
	if len(chunks) > 1:
	chunks[-2]["next_message"] = transcript_formatter(chunks[-1])

	return chunks

	def merge_chunks(chunks: list[dict]) -> list[dict]:
	"""
	Merge consecutive chunks that have the same speaker

	Args:
	chunks (list[dict]): A list of dictionaries, each containing a chunk of text and its metadata.

	Returns:
	list[dict]: A list of merged chunks.
	"""
	if not chunks:
	return []

	merged_chunks = []
	current_chunk = chunks[0]
	for chunk in chunks[1:]:
	if chunk["speaker"] == current_chunk["speaker"]:
	# Merge the text of the current chunk with the next one
	current_chunk["text"] += " " + chunk["text"]
	else:
	# Add the current chunk to the merged list and start a new one
	merged_chunks.append(current_chunk)
	current_chunk = chunk
	# Add the last chunk
	merged_chunks.append(current_chunk)
	return merged_chunks


	def generate_filename_from_details(minute_block_md: str) -> str:
	"""
	Generates a filename from a minute block by extracting details
	from a line like "\| 24 March 2025, Guoqi Wu, AVP CSD \| \|".
	Format: YYYY-MM-DD_Name1Name2_RolePart1_RolePart2_CLE.md
	"""
	lines = minute_block_md.splitlines()
	details_str = None

	if len(lines) > 1:
	# The line containing details for the filename is typically the second line of the block
	# e.g., \| 24 March 2025, Guoqi Wu, AVP CSD \| \|
	header_info_line = lines[1]

	# Extract content from between the first and second pipe characters
	line_parts = header_info_line.split('\|')
	if len(line_parts) > 2 and line_parts[1].strip():
	details_str = line_parts[1].strip()
	else:
	raise ValueError(f"Could not parse details for filename from line: '{header_info_line}' in minute block.")
	else:
	raise ValueError("Minute block has insufficient lines to extract details for filename.")

	if not details_str:
	raise ValueError("Failed to extract details string from minute block.")

	# Regex to separate date from the rest of the details
	match = re.match(r"(\d{1,2}\s+[A-Za-z]+\s+\d{4}),\s*(.+)", details_str)
	if not match:
	raise ValueError(f"Details string '{details_str}' does not match expected format 'Date, Details...'")

	date_part_str = match.group(1)
	name_and_role_parts_str = match.group(2)

	# Convert date to YYYY-MM-DD
	try:
	dt_object = datetime.strptime(date_part_str, "%d %B %Y")
	formatted_date = dt_object.strftime("%Y-%m-%d")
	except ValueError:
	try:
	dt_object = datetime.strptime(date_part_str, "%d %b %Y")
	formatted_date = dt_object.strftime("%Y-%m-%d")
	except ValueError as e:
	raise ValueError(f"Date string '{date_part_str}' could not be parsed: {e}")

	parts = [part.replace(" ", "") for part in name_and_role_parts_str.split(", ")]
	combined_name_role = "_".join(parts)
	combined_name_role_sanitized = re.sub(r'[^\w-]', '', combined_name_role.replace("(", "").replace(")", ""))

	return f"{formatted_date}_{combined_name_role_sanitized}_CLE.md", (formatted_date, combined_name_role_sanitized)

	def load_prompt(prompt_file_name:str, working_directory) -> str:
	"""
	Load a prompt from a file in the specified working directory's prompts folder.

	Args:
	prompt_file_name (str): The name of the prompt file to load.
	working_directory (str): The directory where the prompt file is located.

	Returns:
	str: The content of the prompt file.
	"""
	prompt_file_path = os.path.join(working_directory, "prompts", prompt_file_name+".txt")
	if not os.path.exists(prompt_file_path):
	raise FileNotFoundError(f"Prompt file '{prompt_file_name}' at path '{prompt_file_path}' does not exist.")

	with open(prompt_file_path, 'r', encoding='utf-8') as file:
	return file.read()