ppanchigar
423 lines
15 KiB
Python
423 lines
15 KiB
Python
"""
|
|
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 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 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(
|
|
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)}")
|
|
|
|
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)
|
|
|
|
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
|
|
answers_dict = string_utils.normalize_state_field(
|
|
answers_dict, field_name="PAYER_STATE"
|
|
)
|
|
|
|
# Normalize all 1:1 fields: convert single-element lists to strings
|
|
# Preserves legitimate multi-value lists (e.g., LOB, PROGRAM when passed to 1:1)
|
|
answers_dict = string_utils.normalize_one_to_one_answers_dict(answers_dict)
|
|
|
|
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:
|
|
# 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, _parser = 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
|
|
):
|
|
llm_answer_raw = llm_utils.invoke_claude(
|
|
prompt, "sonnet_latest", filename, max_tokens=8192
|
|
)
|
|
try:
|
|
llm_answer_final = _parser(llm_answer_raw) # returns dict
|
|
|
|
field_value = llm_answer_final.get(field.field_name)
|
|
if not isinstance(field_value, list):
|
|
field_value = [field_value]
|
|
|
|
if field_value[0] == "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)
|