doczyai-pipelines/src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py

"""
RAG-Based Field Extraction Functions

This module provides RAG (Retrieval-Augmented Generation) based field extraction
for contract processing using hybrid retrieval (BM25 + semantic search) with reranking.
"""

import concurrent.futures
import gc
import logging
import re
from typing import List, Tuple, Dict

from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from src.utils import (
    instrumentation_context,
    llm_utils,
    string_utils,
    timing_utils,
    rag_utils,
)
import src.pipelines.shared.preprocessing.preprocessing_funcs as preprocessing_funcs
import src.pipelines.shared.preprocessing.hybrid_smart_chunking_preprocessing as hybrid_smart_chunking_preprocessing
from src import config
from src.prompts import prompt_templates
from src.prompts.fieldset import FieldSet
from src.pipelines.shared.extraction.one_to_one_funcs import (
    check_and_update_effective_date,
)
from src.utils.rag_utils import (
    RAGConfig,
    HSC_CONFIG,
    get_reranker_client,
    get_embeddings,
    create_vectorstore,
    create_hybrid_retriever,
)


def clip_context_by_tokens(context: str, max_tokens: int = 150000) -> str:
    """
    Clip context text if it exceeds max_tokens.
    Uses closed formula: char_count / 3.5 for Claude tokenization.

    Args:
        context: Text to potentially clip
        max_tokens: Maximum token limit (default 150k for Claude 200k window)

    Returns:
        Clipped or original context text
    """
    if not context:
        return ""

    # Estimate tokens: char_count / 3.5 (accurate for Claude)
    token_count = int(len(context) / 3.5)

    # Only clip if exceeds max_tokens
    if token_count > max_tokens:
        approx_chars = int(max_tokens * 3.5)
        clipped_context = context[:approx_chars]
        logging.debug(f"Context clipped from ~{token_count} to ~{max_tokens} tokens")
        return clipped_context
    else:
        logging.debug(f"Context size ~{token_count} tokens - no clipping needed")
        return context


def prepare_documents_for_chunking(
    sections: dict, filename: str = "contract"
) -> List[Document]:
    """Prepare contract sections as documents for chunking."""
    documents = []
    for section, section_data in sections.items():
        section_header = section
        page_numbers = section_data["pages"]
        section_text = section_data["content"]
        doc = Document(
            page_content=section_text,
            metadata={
                "section_header": section_header,
                "section_length": len(section_text),
                "pages": page_numbers,
                "filename": filename,
            },
        )
        documents.append(doc)

    return documents


def format_chunks_for_llm(text_dict: Dict, docs: List[Document]) -> str:
    """Format retrieved documents for LLM context."""
    if not docs:
        return ""

    pages = set()

    for doc in docs:
        if "pages" in doc.metadata:
            pages.update(doc.metadata["pages"])

    pages = set(map(str, pages))

    # Filter out None values and pages that don't exist in text_dict
    context = "\n\n".join(
        [
            text_dict[page]
            for page in pages
            if page in text_dict and text_dict[page] is not None
        ]
    )

    return clip_context_by_tokens(context)


def add_page_numbers_to_chunks(sections):
    prev = None
    updated_sections = {}
    for section, section_data in sections.items():
        matches = re.findall(r"Start of Page No\.?\s*=\s*(\d+)", section_data)
        prev_last_page = int(prev[-1]) if prev else None
        matches = [int(m) for m in matches]
        all_matches = (
            matches + [prev_last_page] if prev_last_page is not None else matches
        )
        all_matches = sorted(set(all_matches))
        updated_sections[section] = {"content": section_data, "pages": all_matches}
        prev = all_matches
    return updated_sections


def preprocessing_pipeline(contract_text):
    cleaned_text = preprocessing_funcs.clean_newlines(contract_text)
    cleaned_text = preprocessing_funcs.clean_law_symbols(cleaned_text)
    cleaned_text = re.sub(
        r"^(.*?)(Start of Page No\.?\s*=\s*1)", r"\2", cleaned_text, flags=re.DOTALL
    )
    sections, original_content, line_ending, verification_result = (
        hybrid_smart_chunking_preprocessing.split_contract_text_into_sections(
            cleaned_text
        )
    )
    is_preserved = verification_result[0]  # get status

    if not is_preserved:
        logging.debug(
            "Failed to preserve content during section split, reverting to full contract"
        )
        sections = {}
        chunk_size = 9000  # keep bigger chunks to have enough coverage as it would be split further by text splitter
        for i in range(0, len(cleaned_text), chunk_size):
            section_name = f"Section {i // chunk_size}"
            sections[section_name] = cleaned_text[i : i + chunk_size]

    sections_with_page_numbers = add_page_numbers_to_chunks(sections)
    return sections_with_page_numbers, is_preserved


def update_provider_name(
    answers_dict: dict,
    filename: str,
    one_to_one_fields: FieldSet,
    retriever,
    reranker,
    rag_config: RAGConfig,
    text_dict: dict,
    constants,
) -> dict:
    """Update provider name fields using enhanced RAG if initial extraction failed.

    Args:
        answers_dict: Dictionary containing extracted field values
        filename: Name of the file being processed
        one_to_one_fields: FieldSet containing field definitions
        retriever: Hybrid retriever instance
        reranker: Reranker client
        rag_config: RAG configuration
        text_dict: Dictionary mapping page numbers to text content
        constants: Constants object

    Returns:
        Updated answers_dict with provider name fields populated
    """
    provider_name = answers_dict.get("PROVIDER_NAME")

    # if (provider name is a str and non empty) or (provider name is a list and has only one non empty item),
    # then populate AARETE_DERIVED_PROVIDER_NAME with that provider name
    if provider_name and (
        (isinstance(provider_name, str) and not string_utils.is_empty(provider_name))
        or (
            isinstance(provider_name, list)
            and len(provider_name) == 1
            and not string_utils.is_empty(provider_name)
        )
    ):
        answers_dict["AARETE_DERIVED_PROVIDER_NAME"] = (
            provider_name if isinstance(provider_name, str) else provider_name[0]
        )
        logging.info(
            f"AARETE_DERIVED_PROVIDER_NAME set to: {answers_dict['AARETE_DERIVED_PROVIDER_NAME']} | Filename: {filename}"
        )
    else:
        logging.info(
            f"Attempting enhanced RAG extraction for AARETE_DERIVED_PROVIDER_NAME | Filename: {filename}"
        )

        # Create custom RAG config with more aggressive retrieval for AARETE_DERIVED_PROVIDER_NAME
        enhanced_config = RAGConfig(
            CHUNK_SIZE=rag_config.CHUNK_SIZE,
            CHUNK_OVERLAP=rag_config.CHUNK_OVERLAP,
            SEMANTIC_TOP_K=35,
            KEYWORD_TOP_K=25,
            ENSEMBLE_TOP_K=40,
            RERANKER_TOP_K=30,  # increase reranker top k to have more candidates for LLM to extract from
            BM25_WEIGHT=0.5,  # Balanced for better keyword matching
            VECTOR_WEIGHT=0.5,  # Balanced for better semantic matching
            RERANKER=rag_config.RERANKER,
            EMBEDDING_MODEL_ID=rag_config.EMBEDDING_MODEL_ID,
            BEDROCK_REGION=rag_config.BEDROCK_REGION,
        )
        hsc_provider_name = prompt_hsc_single_field(
            one_to_one_fields.get_field("AARETE_DERIVED_PROVIDER_NAME"),
            retriever,
            reranker,
            enhanced_config,  # Use enhanced config instead of default
            text_dict,
            constants,
            filename,
        )[
            1
        ]  # Get the field value from the returned tuple
        if not string_utils.is_empty(hsc_provider_name):
            answers_dict["AARETE_DERIVED_PROVIDER_NAME"] = (
                hsc_provider_name[0]
                if isinstance(hsc_provider_name, list)
                else hsc_provider_name
            )
            logging.info(
                f"AARETE_DERIVED_PROVIDER_NAME extracted via enhanced RAG: {answers_dict['AARETE_DERIVED_PROVIDER_NAME']} | Filename: {filename}"
            )

        if string_utils.is_empty(answers_dict["PROVIDER_NAME"]):
            answers_dict["PROVIDER_NAME"] = answers_dict.get(
                "AARETE_DERIVED_PROVIDER_NAME", "N/A"
            )
            logging.info(
                f"Empty PROVIDER_NAME updated to: {answers_dict['PROVIDER_NAME']} | Filename: {filename}"
            )

    return answers_dict


def run_hybrid_smart_chunked_fields(
    one_to_one_fields: FieldSet,
    constants,
    contract_text: str,
    filename: str,
    text_dict: dict,
    rag_config: RAGConfig = HSC_CONFIG,
) -> dict[str, str]:
    """Process fields using RAG: hybrid retrieval (BM25+semantic) + reranking.

    Uses retrieval questions from the field's retrieval_question attribute for chunk retrieval,
    then actual LLM prompts from the field's prompt attribute for field extraction.
    """

    logging.debug(f"Starting RAG extraction for {filename}")

    # Collect fields that have retrieval questions defined
    # Only process fields that have the retrieval_question attribute set
    fields_to_process = [
        field
        for field in one_to_one_fields.fields
        if field.field_type == "smart_chunked"
        and not string_utils.is_empty(field.retrieval_question)
    ]

    if not fields_to_process:
        logging.debug(
            f"No fields with retrieval questions found | Filename: {filename}"
        )
        return {}

    try:
        # Setup: clients, splitter, chunks
        with timing_utils.timed_block("hsc.setup", context=filename):
            reranker = get_reranker_client(rag_config.RERANKER)
            embeddings = get_embeddings(
                rag_config.EMBEDDING_MODEL_ID, rag_config.BEDROCK_REGION
            )
            splitter = RecursiveCharacterTextSplitter(
                chunk_size=rag_config.CHUNK_SIZE,
                chunk_overlap=rag_config.CHUNK_OVERLAP,
                add_start_index=True,
                strip_whitespace=True,
                separators=["\n\n", "\n", " ", ""],
            )

        with timing_utils.timed_block("hsc.preprocessing", context=filename):
            sections, is_preserved = preprocessing_pipeline(contract_text)
            docs = prepare_documents_for_chunking(sections, filename)
            chunks = splitter.split_documents(docs)
            del docs
            gc.collect()

        if not chunks:
            return {}

        # Build retriever
        with timing_utils.timed_block("hsc.vectorstore_creation", context=filename):
            vectorstore = create_vectorstore(chunks, embeddings)
            retriever = create_hybrid_retriever(chunks, vectorstore, rag_config)

        answers_dict = {}
        with timing_utils.timed_block(
            "hsc.parallel_field_processing", context=filename
        ):
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=min(len(fields_to_process), 20)
            ) as executor:
                futures = [
                    instrumentation_context.submit_with_context(
                        executor,
                        prompt_hsc_single_field,
                        field,
                        retriever,
                        reranker,
                        rag_config,
                        text_dict,
                        constants,
                        filename,
                    )
                    for field in fields_to_process
                ]

                for future in concurrent.futures.as_completed(futures):
                    try:
                        field_name, field_value, field = future.result()
                        answers_dict[field_name] = field_value
                    except Exception as e:
                        logging.error(f"Error processing field: {str(e)}")

        # Update provider name fields with enhanced RAG if needed
        answers_dict = update_provider_name(
            answers_dict=answers_dict,
            filename=filename,
            one_to_one_fields=one_to_one_fields,
            retriever=retriever,
            reranker=reranker,
            rag_config=rag_config,
            text_dict=text_dict,
            constants=constants,
        )

        del retriever, vectorstore, chunks
        gc.collect()
        logging.debug(f"RAG extraction done: {len(answers_dict)} fields")

        # FILENAME_AMENDMENT_NUM
        answers_dict = extract_amendment_num_from_filename(answers_dict, filename)

        # Normalize PROVIDER_STATE
        answers_dict = string_utils.normalize_state_field(
            answers_dict, field_name="PROVIDER_STATE"
        )

        # Normalize PAYER_STATE
        answers_dict = string_utils.normalize_state_field(
            answers_dict, field_name="PAYER_STATE"
        )

        return answers_dict

    except Exception as e:
        logging.error(f"RAG pipeline failed: {e}")
        return {}


def extract_amendment_num_from_filename(answer_dict: dict, filename: str) -> dict:
    """Extract amendment number from filename if not found in contract content.

    Args:
        answer_dict (dict): Dictionary containing CONTRACT_AMENDMENT_NUM key.
        filename (str): Filename to extract amendment number from.

    Returns:
        dict: Updated answer_dict with FILENAME_AMENDMENT_NUM populated.
    """
    if (
        string_utils.is_empty(answer_dict.get("CONTRACT_AMENDMENT_NUM"))
        or answer_dict.get("CONTRACT_AMENDMENT_NUM") == "UNKNOWN"
    ):

        prompt, _parser = prompt_templates.EXTRACT_AMENDMENT_NUM_FROM_FILENAME(filename)
        llm_answer_raw = llm_utils.invoke_claude(
            prompt,
            "sonnet_latest",
            filename,
            cache=True,
            instruction=prompt_templates.EXTRACT_AMENDMENT_NUM_FROM_FILENAME_INSTRUCTION(),
            usage_label="EXTRACT_AMENDMENT_NUM_FROM_FILENAME",
        )

        try:
            llm_answer_final = _parser(llm_answer_raw)
            answer_dict["FILENAME_AMENDMENT_NUM"] = llm_answer_final.get(
                "amendment_number", "N/A"
            )
        except (ValueError, KeyError) as e:
            logging.error(
                f"Failed to parse amendment number from filename {filename}: {e}"
            )
            answer_dict["FILENAME_AMENDMENT_NUM"] = "N/A"

    return answer_dict


def prompt_hsc_single_field(
    field, retriever, reranker, rag_config, text_dict, constants, filename
):
    """Process a single field in parallel for HSC."""
    try:
        retrieval_query = field.retrieval_question

        if not retrieval_query:
            return (field.field_name, "N/A", field)

        # Retrieve + rerank using retrieval question
        with timing_utils.timed_block(
            f"hsc.field.{field.field_name}.retrieval", context=filename
        ):
            retrieved = retriever.invoke(retrieval_query)[: rag_config.ENSEMBLE_TOP_K]
            reranked = rag_utils.rerank_documents(
                retrieval_query,
                retrieved,
                reranker,
                rag_config.RERANKER_TOP_K,
                rag_config.RERANKER,
            )

        # Build context and use actual LLM prompt for extraction
        context = format_chunks_for_llm(text_dict, reranked)

        if len(context) == 0:
            logging.warning(
                f"No context retrieved for {field.field_name}, marking as Null value"
            )
            return (field.field_name, "N/A", field)

        llm_prompt = field.get_prompt_dict(
            constants
        )  # Returns dict of {field_name : prompt}
        if field.field_name in ["CONTRACT_TITLE", "PAYER_NAME"]:
            # concatenate values from first 3 items in text_dict to provide more context for contract title extraction, as it is often found in the beginning of the contract
            context = "\n\n".join(list(text_dict.values())[:3])
        if field.field_name == "EFFECTIVE_DT":
            # for effective date always provide first 3 pages as context as effective date is often found in the beginning of the contract and it is important to provide enough context for LLM to extract it correctly
            context = "\n\n".join(list(text_dict.values())[:3]) + "\n\n" + context
        prompt, _parser = prompt_templates.ONE_TO_ONE_SINGLE_FIELD_TEMPLATE(
            context, llm_prompt, field_name=field.field_name
        )

        logging.debug(f"Field: {field.field_name}")
        logging.debug(f"Retrieval question: {retrieval_query}")
        logging.debug(f"LLM prompt: {llm_prompt}")
        logging.debug(
            f"context with len {len(context)}: {context[:500]}........{context[-500:]}"
        )  # log only the beginning and end of the context to avoid clutter

        # LLM call
        with timing_utils.timed_block(
            f"hsc.field.{field.field_name}.llm_call", context=filename
        ):
            llm_answer_raw = llm_utils.invoke_claude(
                prompt, "sonnet_latest", filename, max_tokens=8192
            )
        logging.debug(f"Raw LLM answer for {field.field_name}: {llm_answer_raw}")
        try:
            llm_answer_final = _parser(llm_answer_raw)  # returns dict
            field_value = llm_answer_final.get(field.field_name)
            return (field.field_name, field_value, field)

        except Exception as e:
            logging.warning(
                f"Failed to parse JSON response for field {field.field_name}: "
                f"{type(e).__name__}: {str(e)}. Setting field value to N/A."
            )
            return (field.field_name, "N/A", field)

    except Exception as e:
        logging.error(f"Error on {field.field_name}: {type(e).__name__}: {str(e)}")
        return (field.field_name, "N/A", field)