Merged in feature/lesser-table-caching-refactor-hybrid (pull request #847)
Feature/lesser table caching refactor hybrid * chore: Remove unused duplicate main.py from shared pipeline * fix: Correct crosswalk paths in aarete_derived.py * chore: Remove unused documentation files from fieldExtraction * docs: Add documentation files to documentation folder * docs: Update README with uv setup, expanded project structure, and branching conventions * docs: Add uv installation steps with Ubuntu/WSL emphasis * Enable prompt caching for all remaining LLM calls - Add _INSTRUCTION() functions for: EXHIBIT_HEADER, EXHIBIT_LINKAGE, EXHIBIT_TITLE_MATCH, DATE_FIX, DERIVED_TERM_DATE, CHECK_PROVIDER_NAME_MATCH, SPECIAL_CASE_ASSIGNMENT - Update all invoke_claude() calls in saas and clover pipelines to use cache=True with corresponding _INSTRUCTION() functions - Add new instructions to get_cacheable_instructions() for cache warming - Update tests for new instruction functions Functions now using caching: - prompt_exhibit_level - prompt_exhibit_lesser (EXHIBIT_LEVEL_LESSER_OF) - prompt_fee_schedule_breakout - prompt_grouper_breakout - prompt_special_case_assignment - prompt_exhibit_linkage - prompt_exhibit_header - prompt_smart_chunked (ONE_TO_ONE templates) - prompt_date_fix - prompt_derived_term_date - prompt_exhibit_title_match - provider_name_match_check 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Reorder * feat: Add bcbs_promise client pipeline with OFFSET_TERM extraction - Add new bcbs_promise client with HSC-based OFFSET_TERM field extraction - Extract full paragraph text of offset/recoupment provisions from contracts - Derive OFFSET_INDICATOR (Y/N) from OFFSET_TERM presence - Fix reorder_columns to preserve extra columns not in COLUMN_ORDER - Update QC/QA output path to outputs/qc_qa/ * fix: Update dev deps and test assertions for QC/QA output path - Add pytest/pytest-mock to dev dependencies for mypy type checking - Update test assertions to expect outputs/qc_qa instead of qa_qc_output * style: Apply black formatting to prompt_templates.py * Merge main, move scripts * Archive some scripts * update py version * remove .py version file * Remove ASCII characters * Restore testbed code * restore tracking * Update testbed metrics * Enable prompt caching for CODE_LAST_CHECK, FILL_BILL_TYPE, DUAL_LOB_CHECK, and GROUPER_BREAKOUT - Add CODE_LAST_CHECK_INSTRUCTION() for service specificity classification - Add FILL_BILL_TYPE_INSTRUCTION() for bill type code determination - Add DUAL_LOB_CHECK_INSTRUCTION() for Medicare/Medicaid classification - Update code_funcs.py to use caching for CODE_LAST_CHECK, FILL_BILL_TYPE, GROUPER_BREAKOUT - Update postprocessing_funcs.py to use caching for DUAL_LOB_CHECK - Add new instructions to get_cacheable_instructions() for cache warming - Add unit tests for new instruction functions 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com> * Fix postprocessing_funcs to remove invalid columns * Merge branch 'main' into feature/lesser-table-caching-refactor-hybrid * Revert prompt caching changes from aed1b73c * update formatting * Update imports Approved-by: Sha Brown Approved-by: Praneel Panchigar
This commit is contained in:
Katon Minhas
@@ -0,0 +1,390 @@
|
||||
"""
|
||||
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 {}
|
||||
Reference in New Issue
Block a user