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

import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
import src.utils.timing_utils as timing_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.constants.constants import Constants
from src.constants.delimiters import Delimiter
from src import config
from src.prompts import prompt_templates
from src.prompts.fieldset import FieldSet, Field
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_bedrock_client,
    get_reranker_client,
    get_embeddings,
    create_vectorstore,
    create_hybrid_retriever,
    rerank_documents,
)


def process_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 = 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:
            # Fallback: mark field for full_context processing
            field.field_type = "full_context"
            logging.warning(
                f"No context retrieved for {field.field_name}, marking as full_context"
            )
            return (field.field_name, "N/A", field)

        llm_prompt = field.get_prompt_dict(
            constants
        )  # Returns dict of {field_name : prompt}
        prompt = prompt_templates.ONE_TO_ONE_SINGLE_FIELD_TEMPLATE(context, llm_prompt)

        logging.debug(f"Field: {field.field_name}")
        logging.debug(f"Retrieval question: {retrieval_query}")
        logging.debug(f"LLM prompt: {llm_prompt}")

        # LLM call
        with timing_utils.timed_block(
            f"hsc.field.{field.field_name}.llm_call", context=filename
        ):
            raw = llm_utils.invoke_claude(
                prompt, "sonnet_latest", filename, max_tokens=8192
            )
        try:
            parsed = string_utils.universal_json_load(raw)  # returns dict

            val = parsed.get(field.field_name)
            if isinstance(val, list):
                parsed[field.field_name] = "|".join(map(str, val))
            elif isinstance(val, str):
                parsed[field.field_name] = val
            field_value = parsed[field.field_name]

            if field_value == "N/A":
                field.field_type = "full_context"
                if field.field_name == "PAYER_NAME":
                    field.prompt = (
                        field.prompt
                        + prompt_templates.FULL_CONTEXT_PAYER_NAME_ADDITIONAL_INSTRUCTION()
                    )
                logging.warning(
                    f"Obtained N/A for {field.field_name}, marking as full_context"
                )

            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:
        # Mark field for full_context retry and preserve error info
        field.field_type = "full_context"
        logging.error(
            f"Error on {field.field_name}: {type(e).__name__}: {str(e)}, marking as full_context"
        )
        return (field.field_name, "N/A", field)


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 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 = [
                    executor.submit(
                        process_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)}")

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

        answers_dict = check_and_update_effective_date(
            answers_dict, text_dict, filename
        )

        # even after running the vision api, if the effective date is still N/A, we need to run the prompt again with full context
        if "AARETE_DERIVED_EFFECTIVE_DT" in answers_dict and string_utils.is_empty(
            answers_dict.get("AARETE_DERIVED_EFFECTIVE_DT")
        ):
            # if the effective date is N/A, we need to run the prompt again with full context
            one_to_one_fields.update_field_attr(
                field_name="EFFECTIVE_DT",
                attr="field_type",
                value="full_context",
            )
        # if the termination date is N/A, we need to run the prompt again with full context
        if "TERMINATION_DT" in answers_dict and string_utils.is_empty(
            answers_dict.get("TERMINATION_DT")
        ):
            # if the termination date is N/A, we need to run the prompt again with full context
            one_to_one_fields.update_field_attr(
                field_name="TERMINATION_DT",
                attr="field_type",
                value="full_context",
            )

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

        # Normalize PAYER_STATE to standardized two-letter abbreviation
        for state_field in ["PAYER_STATE"]:
            if state_field in answers_dict:
                answers_dict[state_field] = (
                    string_utils.normalize_state_to_abbreviation(
                        answers_dict[state_field]
                    )
                )
        # 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 {}