Merged in feature/saas-hybrid-smart-chunking (pull request #760)
Feature/saas hybrid smart chunking * WIP: Initial work on SaaS HSC integration * clean up hsc funcs script * fix merging dict issue * Merge remote-tracking branch 'origin/main' into feature/saas-hybrid-smart-chunking * added changes as per feedback comments * add tests and remove old smart chunking funcs * add more tests Approved-by: Katon Minhas
This commit is contained in:
Siddhant Medar
committed by
Katon Minhas
Katon Minhas
parent
992e891388
commit
55643b5a04
Generated
+2938
-2370
File diff suppressed because it is too large
Load Diff
@@ -107,6 +107,8 @@ ENABLE_RUNTIME_ROTATION = False # Set to True to enable multi-runtime failover
|
||||
FIELDS = get_arg_value("fields", "all") # Valid: one_to_one, one_to_n, all
|
||||
FIELD_JSON_PATH = "src/prompts/investment_prompts.json"
|
||||
|
||||
RAG_RETRIEVAL_QUESTIONS_PATH = "src/prompts/rag_retrieval_questions.json"
|
||||
|
||||
# Client-specific args
|
||||
STATE = [s.strip() for s in get_arg_value("state", "None").upper().split("-") if s]
|
||||
CLIENT = [s.strip() for s in get_arg_value("client", "None").split("-") if s]
|
||||
|
||||
@@ -9,6 +9,7 @@ import src.investment.one_to_one_funcs as one_to_one_funcs
|
||||
import src.investment.postprocess as postprocess
|
||||
import src.investment.postprocessing_funcs as postprocessing_funcs
|
||||
import src.investment.preprocess as preprocess
|
||||
import src.investment.hybrid_smart_chunking_funcs as hybrid_smart_chunking_funcs
|
||||
import src.investment.row_funcs as row_funcs
|
||||
import src.investment.tin_npi_funcs as tin_npi_funcs
|
||||
import src.utils.io_utils as io_utils
|
||||
@@ -132,24 +133,32 @@ def run_one_to_one_prompts(
|
||||
one_to_one_results, one_to_one_fields = tin_npi_funcs.run_provider_info_fields(
|
||||
contract_text, one_to_one_fields, text_dict, filename
|
||||
)
|
||||
|
||||
################## RUN SMART CHUNKED PROMPTS ##################
|
||||
smart_chunked_answers_dict, one_to_one_fields = one_to_one_funcs.run_smart_chunked_fields(
|
||||
one_to_one_fields, constants, contract_text, filename
|
||||
|
||||
################## RUN HYBRID SMART CHUNKED PROMPTS ##################
|
||||
# RAG function loads retrieval questions internally and matches with investment_prompts.json
|
||||
hybrid_smart_chunked_answers_dict = hybrid_smart_chunking_funcs.run_hybrid_smart_chunked_fields(
|
||||
one_to_one_fields, constants, contract_text, filename, text_dict
|
||||
)
|
||||
|
||||
|
||||
################## RUN FULL CONTEXT PROMPTS ##################
|
||||
full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
|
||||
one_to_one_fields, contract_text, constants, filename
|
||||
)
|
||||
|
||||
################## COMBINE ANSWERS ################
|
||||
one_to_one_results.update(smart_chunked_answers_dict)
|
||||
one_to_one_results.update(full_context_answers_dict)
|
||||
# Prefer HSC answers over full_context when HSC value is not N/A
|
||||
for key, value in full_context_answers_dict.items():
|
||||
one_to_one_results[key] = value
|
||||
|
||||
for key, value in hybrid_smart_chunked_answers_dict.items():
|
||||
if not string_utils.is_empty(value):
|
||||
one_to_one_results[key] = value
|
||||
elif key not in one_to_one_results:
|
||||
# Add HSC's N/A if field doesn't exist yet
|
||||
one_to_one_results[key] = value
|
||||
|
||||
################## ADD AD FIELDS ################
|
||||
one_to_one_results = one_to_one_funcs.get_aarete_derived_dates(one_to_one_results, text_dict, filename)
|
||||
|
||||
|
||||
################## Crosswalk Fields ##################
|
||||
one_to_one_results = aarete_derived.get_crosswalk_fields(
|
||||
[one_to_one_results], constants
|
||||
|
||||
@@ -0,0 +1,322 @@
|
||||
"""
|
||||
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 gc
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple, Dict
|
||||
|
||||
from langchain.schema 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.investment.preprocessing_funcs as preprocessing_funcs
|
||||
import src.investment.hybrid_smart_chunking_preprocessing as hybrid_smart_chunking_preprocessing
|
||||
from constants.constants import Constants
|
||||
from constants.delimiters import Delimiter
|
||||
from src import config
|
||||
from src.prompts import prompt_templates
|
||||
from src.prompts.fieldset import FieldSet, Field
|
||||
from src.investment.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 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.info(f"Context clipped from ~{token_count} to ~{max_tokens} tokens")
|
||||
return clipped_context
|
||||
else:
|
||||
logging.info(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.info(
|
||||
"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.info(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.retrieval_question is not None
|
||||
]
|
||||
|
||||
if not fields_to_process:
|
||||
logging.info(f"No fields with retrieval questions found | Filename: {filename}")
|
||||
return {}
|
||||
|
||||
try:
|
||||
# Setup: clients, splitter, chunks
|
||||
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", " ", ""],
|
||||
)
|
||||
|
||||
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
|
||||
vectorstore = create_vectorstore(chunks, embeddings)
|
||||
retriever = create_hybrid_retriever(chunks, vectorstore, rag_config)
|
||||
|
||||
# Process fields
|
||||
answers_dict = {}
|
||||
for field in fields_to_process:
|
||||
try:
|
||||
# Get retrieval question from field attribute
|
||||
retrieval_query = field.retrieval_question
|
||||
|
||||
if not retrieval_query:
|
||||
continue
|
||||
# Retrieve + rerank using retrieval question
|
||||
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"
|
||||
)
|
||||
continue
|
||||
|
||||
llm_prompt = field.get_prompt(
|
||||
constants
|
||||
) # Actual extraction prompt from investment_prompts.json
|
||||
prompt = prompt_templates.ONE_TO_ONE_SINGLE_FIELD_TEMPLATE(
|
||||
context, llm_prompt
|
||||
)
|
||||
|
||||
logging.info(f"Field: {field.field_name}")
|
||||
logging.info(f"Retrieval question: {retrieval_query}")
|
||||
logging.info(f"LLM prompt: {llm_prompt}")
|
||||
|
||||
# LLM call
|
||||
raw = llm_utils.invoke_claude(
|
||||
prompt, "sonnet_latest", filename, max_tokens=8192
|
||||
)
|
||||
|
||||
parsed = string_utils.extract_text_from_delimiters(
|
||||
raw, Delimiter.PIPE
|
||||
)
|
||||
|
||||
answers_dict[field.field_name] = parsed if parsed else "N/A"
|
||||
|
||||
if answers_dict[field.field_name] == "N/A":
|
||||
field.field_type = "full_context"
|
||||
logging.warning(
|
||||
f"Obtained N/A for {field.field_name}, marking as full_context"
|
||||
)
|
||||
|
||||
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"
|
||||
)
|
||||
answers_dict[field.field_name] = "N/A"
|
||||
|
||||
del retriever, vectorstore, chunks
|
||||
gc.collect()
|
||||
logging.info(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="AARETE_DERIVED_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 and PAYER_STATE to standardized two-letter abbreviation
|
||||
for state_field in ["PROVIDER_STATE", "PAYER_STATE"]:
|
||||
if state_field in answers_dict:
|
||||
answers_dict[state_field] = (
|
||||
string_utils.normalize_state_to_abbreviation(
|
||||
answers_dict[state_field]
|
||||
)
|
||||
)
|
||||
|
||||
return answers_dict
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"RAG pipeline failed: {e}")
|
||||
return {}
|
||||
@@ -0,0 +1,304 @@
|
||||
import re
|
||||
|
||||
|
||||
def split_contract_text_into_sections(text_content):
|
||||
"""
|
||||
Splits a text document into sections based on common heading patterns,
|
||||
ensuring 100% text preservation by tracking every single character.
|
||||
|
||||
Args:
|
||||
text_content (str): The text content to split into sections
|
||||
|
||||
Returns:
|
||||
tuple: (sections_dict, original_content, line_ending, verification_result)
|
||||
- sections_dict: Dictionary with section titles as keys and content as values
|
||||
- original_content: The original text content (for verification)
|
||||
- line_ending: The line ending type detected ('\n', '\r\n', or '')
|
||||
- verification_result: Tuple (is_preserved, message) indicating if content was preserved
|
||||
"""
|
||||
|
||||
original_content = text_content
|
||||
|
||||
# Split into lines but preserve the original line endings
|
||||
if "\r\n" in original_content:
|
||||
raw_lines = original_content.split("\r\n")
|
||||
line_ending = "\r\n"
|
||||
elif "\n" in original_content:
|
||||
raw_lines = original_content.split("\n")
|
||||
line_ending = "\n"
|
||||
else:
|
||||
raw_lines = [original_content]
|
||||
line_ending = ""
|
||||
|
||||
# --- Define all heading patterns ---
|
||||
main_arabic_heading_pattern = re.compile(
|
||||
r"^\s*(\d{1})\.?\s+([A-Za-z0-9\s&\/-]+?)\.?\s*$", re.IGNORECASE
|
||||
)
|
||||
sub_arabic_heading_pattern = re.compile(
|
||||
r"^\s*(\d+(?:\.\d+)+)\s+([A-Za-z0-9\s&\/-]+?)\.?\s*$", re.IGNORECASE
|
||||
)
|
||||
main_roman_heading_pattern = re.compile(
|
||||
r"^\s*([IVXLCDM]+)\.?\s+([A-Za-z0-9\s&\/-]+?)\.?\s*$", re.IGNORECASE
|
||||
)
|
||||
keyword_colon_pattern = re.compile(r"^\s*(RECITALS|AGREEMENT):\s*$", re.IGNORECASE)
|
||||
standalone_keyword_pattern = re.compile(
|
||||
r"^\s*(RECITALS|AGREEMENT)\s*$", re.IGNORECASE
|
||||
)
|
||||
article_roman_pattern = re.compile(
|
||||
r"^\s*ARTICLE\s+([IVXLCDM]+)\.?\s+([A-Za-z0-9\s&\/-]+?)\.?\s*$", re.IGNORECASE
|
||||
)
|
||||
exhibit_pattern = re.compile(r"^\s*EXHIBIT\s+([A-Z])(?:\s+(.*))?$", re.IGNORECASE)
|
||||
attachment_pattern = re.compile(
|
||||
r"^\s*Attachment\s+([A-Za-z0-9\s&\/.-]+?)\s*$", re.IGNORECASE
|
||||
)
|
||||
in_witness_whereof_pattern = re.compile(
|
||||
r"^\s*IN\s+WITNESS\s+WHEREOF,?$", re.IGNORECASE
|
||||
)
|
||||
signatory_name_pattern = re.compile(r"^\s*([A-Z][A-Z0-9\s&\/.-]+)$")
|
||||
specific_signatory_pattern = re.compile(r"^\s*(PROVIDER|ICM):\s*$", re.IGNORECASE)
|
||||
|
||||
# --- Identify all potential heading lines ---
|
||||
heading_events = []
|
||||
|
||||
i = 0
|
||||
while i < len(raw_lines):
|
||||
line_stripped = raw_lines[i].strip()
|
||||
next_line_raw = raw_lines[i + 1] if i + 1 < len(raw_lines) else ""
|
||||
next_line_stripped = next_line_raw.strip()
|
||||
|
||||
current_heading_title = None
|
||||
heading_start_line = i
|
||||
heading_end_line = i
|
||||
lines_consumed = 1
|
||||
|
||||
# Check for multi-line heading prefixes
|
||||
is_numeric_prefix_only = re.fullmatch(r"^\s*(\d{1})\.?\s*$", line_stripped)
|
||||
is_roman_prefix_only = re.fullmatch(r"^\s*([IVXLCDM]+)\.?\s*$", line_stripped)
|
||||
is_exhibit_prefix_only = re.fullmatch(
|
||||
r"^\s*EXHIBIT\s+([A-Z])\s*$", line_stripped, re.IGNORECASE
|
||||
)
|
||||
|
||||
# Check if next line is a primary heading
|
||||
primary_heading_patterns = [
|
||||
sub_arabic_heading_pattern,
|
||||
main_arabic_heading_pattern,
|
||||
main_roman_heading_pattern,
|
||||
keyword_colon_pattern,
|
||||
standalone_keyword_pattern,
|
||||
article_roman_pattern,
|
||||
exhibit_pattern,
|
||||
attachment_pattern,
|
||||
in_witness_whereof_pattern,
|
||||
specific_signatory_pattern,
|
||||
]
|
||||
|
||||
is_next_line_a_primary_heading = any(
|
||||
pattern.match(next_line_stripped) for pattern in primary_heading_patterns
|
||||
)
|
||||
|
||||
# Check if next line is a signatory block
|
||||
if (
|
||||
signatory_name_pattern.match(next_line_stripped)
|
||||
and len(next_line_stripped.split()) > 1
|
||||
and i + 2 < len(raw_lines)
|
||||
):
|
||||
line_after_next_stripped = raw_lines[i + 2].strip()
|
||||
if re.match(
|
||||
r"^\s*By:", line_after_next_stripped, re.IGNORECASE
|
||||
) or re.match(r"^\s*Name:", line_after_next_stripped, re.IGNORECASE):
|
||||
is_next_line_a_primary_heading = True
|
||||
|
||||
# Multi-line heading combination logic
|
||||
if (
|
||||
(is_numeric_prefix_only or is_roman_prefix_only)
|
||||
and next_line_stripped
|
||||
and not is_next_line_a_primary_heading
|
||||
):
|
||||
current_heading_title = f"{line_stripped} {next_line_stripped}"
|
||||
heading_end_line = i + 1
|
||||
lines_consumed = 2
|
||||
elif (
|
||||
is_exhibit_prefix_only
|
||||
and next_line_stripped
|
||||
and not is_next_line_a_primary_heading
|
||||
):
|
||||
current_heading_title = f"{line_stripped} {next_line_stripped}"
|
||||
heading_end_line = i + 1
|
||||
lines_consumed = 2
|
||||
else:
|
||||
# Single line heading checks
|
||||
if article_roman_pattern.match(line_stripped):
|
||||
match = article_roman_pattern.match(line_stripped)
|
||||
current_heading_title = (
|
||||
f"ARTICLE {match.group(1)} {match.group(2).strip().rstrip('.')}"
|
||||
)
|
||||
elif keyword_colon_pattern.match(line_stripped):
|
||||
match = keyword_colon_pattern.match(line_stripped)
|
||||
current_heading_title = match.group(1).strip().upper()
|
||||
elif standalone_keyword_pattern.match(line_stripped):
|
||||
match = standalone_keyword_pattern.match(line_stripped)
|
||||
current_heading_title = match.group(1).strip().upper()
|
||||
elif exhibit_pattern.match(line_stripped):
|
||||
match = exhibit_pattern.match(line_stripped)
|
||||
exhibit_title = f"EXHIBIT {match.group(1).strip()}"
|
||||
if match.group(2):
|
||||
exhibit_title += f" {match.group(2).strip()}"
|
||||
current_heading_title = exhibit_title
|
||||
elif attachment_pattern.match(line_stripped):
|
||||
match = attachment_pattern.match(line_stripped)
|
||||
current_heading_title = f"Attachment {match.group(1).strip()}"
|
||||
elif in_witness_whereof_pattern.match(line_stripped):
|
||||
current_heading_title = "IN WITNESS WHEREOF"
|
||||
elif sub_arabic_heading_pattern.match(line_stripped):
|
||||
match = sub_arabic_heading_pattern.match(line_stripped)
|
||||
current_heading_title = (
|
||||
f"{match.group(1)} {match.group(2).strip().rstrip('.')}"
|
||||
)
|
||||
elif main_arabic_heading_pattern.match(line_stripped):
|
||||
match = main_arabic_heading_pattern.match(line_stripped)
|
||||
current_heading_title = (
|
||||
f"{match.group(1)} {match.group(2).strip().rstrip('.')}"
|
||||
)
|
||||
elif main_roman_heading_pattern.match(line_stripped):
|
||||
match = main_roman_heading_pattern.match(line_stripped)
|
||||
current_heading_title = (
|
||||
f"{match.group(1)} {match.group(2).strip().rstrip('.')}"
|
||||
)
|
||||
elif specific_signatory_pattern.match(line_stripped):
|
||||
match = specific_signatory_pattern.match(line_stripped)
|
||||
current_heading_title = match.group(1).strip().upper()
|
||||
else:
|
||||
# Check for signatory block
|
||||
if (
|
||||
signatory_name_pattern.match(line_stripped)
|
||||
and len(line_stripped.split()) > 1
|
||||
):
|
||||
for j in range(1, min(5, len(raw_lines) - i)):
|
||||
next_line = raw_lines[i + j].strip()
|
||||
if re.match(r"^\s*By:", next_line, re.IGNORECASE) or re.match(
|
||||
r"^\s*Name:", next_line, re.IGNORECASE
|
||||
):
|
||||
current_heading_title = line_stripped
|
||||
break
|
||||
|
||||
if current_heading_title:
|
||||
heading_events.append(
|
||||
(heading_start_line, heading_end_line, current_heading_title)
|
||||
)
|
||||
|
||||
i += lines_consumed
|
||||
|
||||
# Sort heading events by their start line index
|
||||
heading_events.sort(key=lambda x: x[0])
|
||||
|
||||
# --- SECTION EXTRACTION ---
|
||||
# This ensures every single line is accounted for and handles duplicate section titles
|
||||
|
||||
sections = {}
|
||||
section_counter = {} # Track occurrences of each title
|
||||
current_start_line = 0
|
||||
current_section_title = "Preamble"
|
||||
|
||||
for start_line, end_line, title in heading_events:
|
||||
# Add content before this heading to current section
|
||||
if current_start_line < start_line:
|
||||
section_lines = raw_lines[current_start_line:start_line]
|
||||
section_content = line_ending.join(section_lines)
|
||||
if section_content: # Only add non-empty sections
|
||||
sections[current_section_title] = section_content
|
||||
|
||||
# Make section title unique if we've seen it before
|
||||
if title in section_counter:
|
||||
section_counter[title] += 1
|
||||
unique_title = f"{title} ({section_counter[title]})"
|
||||
else:
|
||||
section_counter[title] = 1
|
||||
unique_title = title
|
||||
|
||||
# Start new section with the heading itself
|
||||
current_section_title = unique_title
|
||||
current_start_line = start_line
|
||||
|
||||
# Add the final section (from last heading to end of file)
|
||||
if current_start_line < len(raw_lines):
|
||||
section_lines = raw_lines[current_start_line:]
|
||||
section_content = line_ending.join(section_lines)
|
||||
if section_content:
|
||||
sections[current_section_title] = section_content
|
||||
elif not sections and raw_lines:
|
||||
# Edge case: no headings found
|
||||
sections["Preamble"] = line_ending.join(raw_lines)
|
||||
|
||||
# Verify preservation
|
||||
verification_result = verify_preservation(original_content, sections, line_ending)
|
||||
|
||||
return sections, original_content, line_ending, verification_result
|
||||
|
||||
|
||||
def verify_preservation(original_content, sections, line_ending):
|
||||
"""
|
||||
Verification that ensures 100% text preservation.
|
||||
|
||||
Returns:
|
||||
tuple: (is_preserved, message) where is_preserved is True if content is perfectly preserved
|
||||
"""
|
||||
# Reconstruct the content exactly as it was, preserving order
|
||||
reconstructed_content = line_ending.join(sections.values())
|
||||
|
||||
# Character-by-character comparison
|
||||
if original_content == reconstructed_content:
|
||||
return True, "Perfect preservation: All content preserved exactly."
|
||||
|
||||
# Find the exact difference
|
||||
original_len = len(original_content)
|
||||
reconstructed_len = len(reconstructed_content)
|
||||
|
||||
# Find first difference
|
||||
min_len = min(original_len, reconstructed_len)
|
||||
for i in range(min_len):
|
||||
if original_content[i] != reconstructed_content[i]:
|
||||
context_start = max(0, i - 50)
|
||||
context_end = min(original_len, i + 50)
|
||||
return False, (
|
||||
f"First difference at position {i}:\n"
|
||||
f"Original: {repr(original_content[context_start:context_end])}\n"
|
||||
f"Reconstructed: {repr(reconstructed_content[context_start:context_end])}"
|
||||
)
|
||||
|
||||
# Same content but different lengths
|
||||
if original_len != reconstructed_len:
|
||||
shorter, longer = (
|
||||
(original_content, reconstructed_content)
|
||||
if original_len < reconstructed_len
|
||||
else (reconstructed_content, original_content)
|
||||
)
|
||||
return (
|
||||
False,
|
||||
f"Length mismatch. Extra content in longer version: {repr(longer[min_len:min_len+50])}",
|
||||
)
|
||||
|
||||
return False, "Unknown difference detected."
|
||||
|
||||
|
||||
def get_section_info(sections):
|
||||
"""
|
||||
Helper function to get information about the extracted sections.
|
||||
|
||||
Returns:
|
||||
dict: Information about sections including counts and character counts
|
||||
"""
|
||||
info = {
|
||||
"total_sections": len(sections),
|
||||
"section_details": {},
|
||||
"total_characters": 0,
|
||||
}
|
||||
|
||||
for title, content in sections.items():
|
||||
char_count = len(content)
|
||||
info["section_details"][title] = {
|
||||
"character_count": char_count,
|
||||
"line_count": content.count("\n") + 1 if content else 0,
|
||||
}
|
||||
info["total_characters"] += char_count
|
||||
|
||||
return info
|
||||
@@ -11,11 +11,6 @@ from constants.constants import Constants
|
||||
from constants.delimiters import Delimiter
|
||||
from src import config
|
||||
from src.investment import postprocessing_funcs
|
||||
from src.investment.smart_chunking_funcs import (
|
||||
field_context,
|
||||
group_fields,
|
||||
stitch_chunks,
|
||||
)
|
||||
from src.investment.vision_funcs import get_image_array_based_answer
|
||||
from src.prompts.fieldset import Field, FieldSet
|
||||
|
||||
@@ -63,72 +58,6 @@ def extract_global_lesser_of(contract_text: str, filename: str) -> dict:
|
||||
or "N/A"
|
||||
}
|
||||
|
||||
|
||||
def run_smart_chunked_fields(
|
||||
one_to_one_fields: FieldSet,
|
||||
constants,
|
||||
contract_text: str,
|
||||
filename: str,
|
||||
) -> tuple[dict[str, str], FieldSet]:
|
||||
"""Process fields that require smart chunking and return the extracted answers.
|
||||
|
||||
Args:
|
||||
one_to_one_fields (FieldSet): FieldSet containing the fields to be processed.
|
||||
contract_text (str): Full text of contract.
|
||||
filename (str): Name of the file being processed.
|
||||
|
||||
Raises: ValueError: If the field group has no defined fields.
|
||||
|
||||
Returns:
|
||||
dict[str, str]: Dictionary mapping field names to extracted answers.
|
||||
FieldSet: FieldSet containing the 1:1 fields
|
||||
"""
|
||||
|
||||
one_to_one_fields, splits = field_context(
|
||||
one_to_one_fields, constants, contract_text, filename
|
||||
)
|
||||
|
||||
smart_chunked_fields = one_to_one_fields.filter(
|
||||
field_type="smart_chunked"
|
||||
) # Only get the ones that are still smart_chunked after getting context
|
||||
groups = group_fields(smart_chunked_fields)
|
||||
|
||||
# create new field groups and corresponding chunks
|
||||
new_field_groups = {}
|
||||
one_to_one_chunks = {}
|
||||
for idx, group in enumerate(groups, 1):
|
||||
# convert dictionary of fields to list of fields
|
||||
group_list = list(group)
|
||||
new_field_groups[f"Group {idx}"] = {"fields": group_list}
|
||||
|
||||
# combine chunks of fields of one group into one list
|
||||
retrieved_text_list = [field.retrieved_text for field in group_list]
|
||||
|
||||
# Flatten list of lists into one list
|
||||
retrieved_text_list = [
|
||||
item for internal_list in retrieved_text_list for item in internal_list
|
||||
]
|
||||
retrieved_text_list = sorted(
|
||||
set(
|
||||
item
|
||||
for sublist in retrieved_text_list
|
||||
for item in (sublist if isinstance(sublist, (list, set)) else [sublist])
|
||||
)
|
||||
)
|
||||
|
||||
# join relevant chunks of all fields of the group into one text
|
||||
one_to_one_chunks[f"Group {idx}"] = stitch_chunks(retrieved_text_list)
|
||||
|
||||
field_groups = [field_group for field_group in new_field_groups.keys()]
|
||||
|
||||
answers_dict = {}
|
||||
for field_group in field_groups:
|
||||
answers_dict = prompt_calls.prompt_smart_chunked(answers_dict, new_field_groups, field_group, one_to_one_chunks, filename)
|
||||
|
||||
one_to_one_fields = pass_smart_chunked_to_full_context(answers_dict, one_to_one_fields)
|
||||
|
||||
return answers_dict, one_to_one_fields
|
||||
|
||||
def pass_smart_chunked_to_full_context(answers_dict, one_to_one_fields):
|
||||
# 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 "EFFECTIVE_DT" in answers_dict and string_utils.is_empty(
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
import logging
|
||||
|
||||
import src.investment.smart_chunking_funcs as smart_chunking_funcs
|
||||
import src.investment.table_funcs as table_funcs
|
||||
from src import config
|
||||
from src.investment import preprocessing_funcs
|
||||
|
||||
@@ -1,505 +0,0 @@
|
||||
"""
|
||||
Smart Chunking Functions for Contract Processing
|
||||
|
||||
This module provides intelligent text chunking and semantic search capabilities
|
||||
for contract processing. It includes:
|
||||
|
||||
- Lightweight in-memory vectorstore operations using FAISS
|
||||
- Multiple chunking strategies (hierarchical, OR, AND, regex)
|
||||
- Semantic and keyword-based chunk retrieval
|
||||
- Optimized for ephemeral, per-contract vectorstore usage
|
||||
|
||||
Key Components:
|
||||
- field_context(): Main function for semantic chunk retrieval using FAISS
|
||||
- group_fields(): Groups fields based on retrieved text overlap
|
||||
- Various chunk_*() functions: Different keyword matching strategies
|
||||
|
||||
Performance Notes:
|
||||
- Uses FAISS for fast, in-memory similarity search
|
||||
- No persistent storage - vectorstores are created and destroyed per contract
|
||||
- Designed for concurrent processing without resource leaks
|
||||
- Much lower memory footprint than persistent vectorstore solutions
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from itertools import combinations
|
||||
from typing import Any, Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
import nltk
|
||||
from constants.constants import Constants
|
||||
from langchain_aws.embeddings.bedrock import BedrockEmbeddings
|
||||
from langchain_community.retrievers import BM25Retriever
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from nltk.tokenize import word_tokenize
|
||||
from src import config
|
||||
from src.prompts.fieldset import FieldSet
|
||||
|
||||
# Append path of NLTK data directory (fieldExtraction/nltk_data)
|
||||
nltk.data.path.append(
|
||||
os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "nltk_data"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# Setup bedrock
|
||||
embeddings_model = BedrockEmbeddings(
|
||||
client=config.BEDROCK_RUNTIME, model_id="amazon.titan-embed-text-v1"
|
||||
)
|
||||
|
||||
# ===================================================================
|
||||
# CONFIGURATION CONSTANTS
|
||||
# ===================================================================
|
||||
|
||||
# Text chunking parameters
|
||||
CHUNK_PREFIX_LENGTH = 11 # "chunk 001: " length
|
||||
CHUNK_OVERLAP = 500
|
||||
CHUNK_SIZE = 2000
|
||||
TOP_K = 8
|
||||
|
||||
# Regular expressions
|
||||
DATE_PATTERN = r"(?i)\bday of\b|\b(?:jan(?:uary)?|feb(?:ruary)?|mar(?:ch)?|apr(?:il)?|may|jun(?:e)?|jul(?:y)?|aug(?:ust)?|sep(?:tember)?|oct(?:ober)?|nov(?:ember)?|dec(?:ember)?)\b|\b(?:\d{1,2}[\/\.-]\d{1,2}[\/\.-](?:\d{2}|\d{4})|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{8}|\d{2}\d{2}|\d{4}[\/\.-]\d{1,2}[\/\.-]\d{1,2}|\d{4}[\.]\d{1,2}[\.]\d{1,2})\b"
|
||||
|
||||
|
||||
def parse_chunk(chunk: str) -> tuple[int, str]:
|
||||
"""Extract chunk number and content from a chunk string
|
||||
|
||||
Args:
|
||||
chunk (str): Text chunk, prefixed with "Chunk NNN: "
|
||||
|
||||
Returns:
|
||||
tuple[int, str]: Tuple of (chunk_number, content)
|
||||
"""
|
||||
if not chunk.startswith("chunk ") or len(chunk) < CHUNK_PREFIX_LENGTH:
|
||||
raise ValueError("Invalid chunk format")
|
||||
try:
|
||||
chunk_number = int(chunk[6:9]) # Extract chunk number from "chunk NNN: "
|
||||
except ValueError:
|
||||
raise ValueError("Invalid chunk number format")
|
||||
content = chunk[CHUNK_PREFIX_LENGTH:]
|
||||
return chunk_number, content
|
||||
|
||||
|
||||
def find_overlap(prev_chunk, next_chunk, min_overlap=3):
|
||||
"""
|
||||
Finds the maximum overlap between the end of prev_chunk and the start of next_chunk.
|
||||
Returns the number of overlapping characters.
|
||||
"""
|
||||
max_overlap = min(len(prev_chunk), len(next_chunk))
|
||||
for i in range(max_overlap, min_overlap - 1, -1):
|
||||
if prev_chunk[-i:] == next_chunk[:i]:
|
||||
return i
|
||||
return 0 # No overlap found
|
||||
|
||||
|
||||
def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str:
|
||||
"""Stitch together a list of chunks, with an overlap of `overlap_size` characters
|
||||
|
||||
Args:
|
||||
chunks (list[str]): List of chunk strings, each prefixed with "chunk NNN: "
|
||||
overlap_size (int, optional): Overlap size between chunks. Defaults to CHUNK_OVERLAP.
|
||||
|
||||
Returns:
|
||||
str: Combined text with properly handled overlaps
|
||||
"""
|
||||
if not chunks:
|
||||
return ""
|
||||
|
||||
if overlap_size < 0:
|
||||
raise ValueError("Overlap size must be non-negative")
|
||||
|
||||
# Parse and sort chunks
|
||||
parsed_chunks = [parse_chunk(chunk) for chunk in chunks]
|
||||
sorted_chunks = sorted(parsed_chunks, key=lambda x: x[0])
|
||||
|
||||
# Start with the first chunk's content
|
||||
result = sorted_chunks[0][1]
|
||||
|
||||
# Process subsequent chunks
|
||||
for i in range(1, len(sorted_chunks)):
|
||||
current_chunk = sorted_chunks[i][1]
|
||||
# if chunks are contiguous
|
||||
if sorted_chunks[i][0] == sorted_chunks[i - 1][0] + 1:
|
||||
# Check for overlap
|
||||
overlap = find_overlap(result, current_chunk)
|
||||
if (
|
||||
overlap > 0 and len(result) >= overlap and len(current_chunk) >= overlap
|
||||
): # if overlap is enabled and there is enough text to overlap
|
||||
result = result[:-overlap] + current_chunk
|
||||
else:
|
||||
result += "\n" + current_chunk
|
||||
else: # if chunks are not contiguous, just append them
|
||||
result += "\n" + current_chunk
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def field_context(
|
||||
one_to_one_fields: FieldSet, constants: Constants, contract_text: str, filename: str
|
||||
) -> Tuple[FieldSet, List[str]]:
|
||||
"""Lightweight field context extraction using FAISS (in-memory only)."""
|
||||
|
||||
# Convert contract_text to lowercase for case-insensitive processing
|
||||
contract_text = contract_text.lower()
|
||||
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
|
||||
)
|
||||
splits = text_splitter.split_text(contract_text)
|
||||
|
||||
# add chunk number starting from 1
|
||||
splits = [
|
||||
"chunk " + str(idx + 1).zfill(3) + ": " + str(split)
|
||||
for idx, split in enumerate(splits)
|
||||
]
|
||||
|
||||
# Create lightweight FAISS vectorstore (in-memory only)
|
||||
vectorstore = FAISS.from_texts(texts=splits, embedding=embeddings_model)
|
||||
|
||||
vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K})
|
||||
keyword_retriever = BM25Retriever.from_texts(
|
||||
splits, k=TOP_K, preprocess_func=word_tokenize
|
||||
)
|
||||
|
||||
for field in one_to_one_fields.fields:
|
||||
if field.field_type == "smart_chunked":
|
||||
questions = {}
|
||||
questions[field.field_name] = field.get_prompt(constants)
|
||||
question = json.dumps(questions)
|
||||
|
||||
# Convert field keywords to lowercase for case-insensitive processing
|
||||
field_keywords = [keyword.lower() for keyword in field.keywords]
|
||||
|
||||
# get relevant chunks based on semantic search and keyword search and merge them
|
||||
semantic_relevant_docs = vectorstore_retriever.invoke(question)
|
||||
keyword_relevant_docs = keyword_retriever.invoke(" ;".join(field.keywords))
|
||||
semantic_relevant_chunks = [
|
||||
doc.page_content for doc in semantic_relevant_docs
|
||||
]
|
||||
keyword_relevant_chunks = [
|
||||
doc.page_content for doc in keyword_relevant_docs
|
||||
]
|
||||
if field.field_name == "CONTRACT_TITLE":
|
||||
keyword_relevant_chunks = keyword_relevant_chunks + splits[:5]
|
||||
field.retrieved_text = sorted(
|
||||
list(set(semantic_relevant_chunks + keyword_relevant_chunks))
|
||||
)
|
||||
|
||||
# Fallback: if no context or all context, treat as full_context
|
||||
if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len(
|
||||
splits
|
||||
):
|
||||
field.field_type = "full_context"
|
||||
|
||||
return one_to_one_fields, splits
|
||||
|
||||
|
||||
# Function to group fields based on overlap threshold
|
||||
def group_fields(field_set: FieldSet, threshold=3) -> List[Set]:
|
||||
"""
|
||||
Groups fields in the given FieldSet based on the overlap of their retrieved text.
|
||||
Args:
|
||||
field_set (FieldSet): A FieldSet object containing fields to be grouped.
|
||||
threshold (int, optional): The minimum overlap required to group fields together. Defaults to 3.
|
||||
Returns:
|
||||
List[Set[Field]]: A list of sets, where each set contains fields that have sufficient overlap.
|
||||
"""
|
||||
overlap_matrix = {
|
||||
(f1, f2): len(set(f1.retrieved_text) & set(f2.retrieved_text))
|
||||
for f1, f2 in combinations(field_set.fields, 2)
|
||||
if (overlap := len(set(f1.retrieved_text) & set(f2.retrieved_text)))
|
||||
>= threshold
|
||||
}
|
||||
|
||||
groups = []
|
||||
visited = set() # store field objects, not names
|
||||
|
||||
for field in field_set.fields:
|
||||
if field not in visited: # check if field is already in a group
|
||||
group = {field}
|
||||
stack = [field]
|
||||
while stack:
|
||||
curr = stack.pop()
|
||||
visited.add(curr) # add field to visited set
|
||||
for (f1, f2), overlap in overlap_matrix.items():
|
||||
if f1 == curr and f2 not in visited and overlap >= threshold:
|
||||
group.add(f2)
|
||||
stack.append(f2)
|
||||
elif f2 == curr and f1 not in visited and overlap >= threshold:
|
||||
group.add(f1)
|
||||
stack.append(f1)
|
||||
if group: # only add non-empty groups
|
||||
groups.append(group)
|
||||
return groups
|
||||
|
||||
|
||||
def chunk_hierarchical(
|
||||
text_dict: dict[str, str],
|
||||
keywords: list[str],
|
||||
regex: str = None,
|
||||
case_sensitive: bool = False,
|
||||
) -> list[str]:
|
||||
"""Look for the keywords, one-by-one, and take the first that satisfies the keyword.
|
||||
We pad responsive pages with a one-page buffer (before and after) when possible.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary keyed by string page number and valued
|
||||
by string page contents
|
||||
keywords (list[str]): A list of string keywords
|
||||
case_sensitive(bool): A flag for converting all keywords and text to lower case,
|
||||
in order to ignore case altogether.
|
||||
|
||||
Returns:
|
||||
list[str]: The sorted list of string page numbers responsive to the first
|
||||
keyword found
|
||||
"""
|
||||
if not case_sensitive:
|
||||
keywords = [keyword.lower() for keyword in keywords]
|
||||
text_dict = {key: value.lower() for key, value in text_dict.items()}
|
||||
|
||||
page_list = []
|
||||
for keyword in keywords:
|
||||
if len(page_list) == 0:
|
||||
for page in text_dict.keys():
|
||||
if keyword in text_dict[page]:
|
||||
if str(int(page) - 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) - 1))
|
||||
|
||||
page_list.append(page)
|
||||
|
||||
if str(int(page) + 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) + 1))
|
||||
|
||||
page_list_sorted = sorted([int(page_str) for page_str in set(page_list)])
|
||||
page_list_sorted_str = list(map(str, page_list_sorted))
|
||||
return page_list_sorted_str
|
||||
|
||||
|
||||
def chunk_or(
|
||||
text_dict: dict[str, str],
|
||||
keywords: list[str],
|
||||
regex: str = None,
|
||||
case_sensitive: bool = False,
|
||||
) -> list[str]:
|
||||
"""Look for any of the keywords and add any pages with any keywords to the chunk.
|
||||
We pad responsive pages with a one-page buffer (before and after) when possible.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary keyed by string page number and valued
|
||||
by string page contents
|
||||
keywords (list[str]): A list of string keywords
|
||||
case_sensitive(bool): A flag for converting all keywords and text to lower case,
|
||||
in order to ignore case altogether.
|
||||
|
||||
Returns:
|
||||
list[str]: The sorted list of string page numbers responsive to the first
|
||||
keyword found
|
||||
"""
|
||||
if not case_sensitive:
|
||||
keywords = [keyword.lower() for keyword in keywords]
|
||||
text_dict = {key: value.lower() for key, value in text_dict.items()}
|
||||
|
||||
page_list = []
|
||||
for page in text_dict.keys():
|
||||
for keyword in keywords:
|
||||
if keyword in text_dict[page]:
|
||||
# Add buffer of previous and next page, if we're not at the beginning nor end
|
||||
# If we are at the beginning (or the end) just take the next (or previous) page
|
||||
if str(int(page) - 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) - 1))
|
||||
|
||||
page_list.append(page)
|
||||
|
||||
if str(int(page) + 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) + 1))
|
||||
|
||||
page_list_sorted = sorted([int(page_str) for page_str in set(page_list)])
|
||||
page_list_sorted_str = list(map(str, page_list_sorted))
|
||||
return page_list_sorted_str
|
||||
|
||||
|
||||
# And: run each keyword and include a page if all are present
|
||||
def chunk_and(
|
||||
text_dict: dict[str, str],
|
||||
keywords: list[str],
|
||||
regex: str = None,
|
||||
case_sensitive: bool = False,
|
||||
) -> list[str]:
|
||||
"""Look for all of the keywords and add any pages with all keywords to the chunk.
|
||||
We pad responsive pages with a one-page buffer (before and after) when possible.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary keyed by string page number and valued
|
||||
by string page contents
|
||||
keywords (list[str]): A list of string keywords
|
||||
case_sensitive(bool): A flag for converting all keywords and text to lower case,
|
||||
in order to ignore case altogether.
|
||||
|
||||
Returns:
|
||||
list[str]: The sorted list of string page numbers responsive to the first
|
||||
keyword found
|
||||
"""
|
||||
if not case_sensitive:
|
||||
keywords = [keyword.lower() for keyword in keywords]
|
||||
text_dict = {key: value.lower() for key, value in text_dict.items()}
|
||||
|
||||
page_list = []
|
||||
for page in text_dict.keys():
|
||||
if all([keyword in text_dict[page] for keyword in keywords]):
|
||||
# Add buffer of previous and next page, if we're not at the beginning nor end
|
||||
# If we are at the beginning (or the end) just take the next (or previous) page
|
||||
if str(int(page) - 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) - 1))
|
||||
|
||||
page_list.append(page)
|
||||
|
||||
if str(int(page) + 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) + 1))
|
||||
|
||||
page_list_sorted = sorted([int(page_str) for page_str in set(page_list)])
|
||||
page_list_sorted_str = list(map(str, page_list_sorted))
|
||||
return page_list_sorted_str
|
||||
|
||||
|
||||
# Regex: Check for a regex and include a page if the regex matches within the page
|
||||
def chunk_regex(
|
||||
text_dict: dict[str, str], keywords: list[str], regex: str, case_sensitive: bool
|
||||
) -> list[str]:
|
||||
"""Look for regex match and add any pages with regex match to the chunk.
|
||||
We pad responsive pages with a one-page buffer (before and after) when possible.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary keyed by string page number and valued
|
||||
by string page contents
|
||||
regex (str): A regex expression
|
||||
|
||||
Returns:
|
||||
list[str]: The sorted list of string page numbers responsive to the first
|
||||
keyword found
|
||||
"""
|
||||
|
||||
page_list = []
|
||||
for page in text_dict.keys():
|
||||
if re.search(regex, text_dict[page]):
|
||||
# Add buffer of previous and next page, if we're not at the beginning nor end
|
||||
# If we are at the beginning (or the end) just take the next (or previous) page
|
||||
if str(int(page) - 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) - 1))
|
||||
|
||||
page_list.append(page)
|
||||
|
||||
if str(int(page) + 1) in text_dict.keys():
|
||||
page_list.append(str(int(page) + 1))
|
||||
|
||||
page_list_sorted = sorted([int(page_str) for page_str in set(page_list)])
|
||||
page_list_sorted_str = list(map(str, page_list_sorted))
|
||||
return page_list_sorted_str
|
||||
|
||||
|
||||
def keyword_search(
|
||||
text_dict: dict[str, str],
|
||||
keywords: list[str],
|
||||
case_sensitive: bool = False,
|
||||
regex: bool = False,
|
||||
) -> dict[str, list]:
|
||||
"""Look for keyword/regex match and add any pages with keyword/regex match to the chunk.
|
||||
|
||||
Args:
|
||||
text_dict (dict[str, str]): A dictionary keyed by string page number and valued
|
||||
by string page contents
|
||||
keywords (list[str]): A list of string keywords/regex
|
||||
case_sensitive(bool): A flag for converting all keywords and text to lower case,
|
||||
in order to ignore case altogether.
|
||||
regex (bool): A flag to identify if the passed keyword list is regex list or string list
|
||||
|
||||
Returns:
|
||||
list[str]: The dictionary containing keyword as keys and sorted list of string page numbers responsive to the
|
||||
keyword found
|
||||
"""
|
||||
if not case_sensitive and not regex:
|
||||
keywords = [keyword.lower() for keyword in keywords]
|
||||
text_dict = {key: value.lower() for key, value in text_dict.items()}
|
||||
|
||||
page_dict = {}
|
||||
for keyword in keywords:
|
||||
page_dict[keyword] = []
|
||||
|
||||
for page in text_dict.keys():
|
||||
if (not regex and keyword in text_dict[page]) or (
|
||||
regex and re.search(keyword, text_dict[page])
|
||||
):
|
||||
page_dict[keyword].append(page)
|
||||
|
||||
page_dict_sorted = {}
|
||||
for keyword in page_dict:
|
||||
page_dict_sorted[keyword] = sorted(
|
||||
[int(page_str) for page_str in set(page_dict[keyword])]
|
||||
)
|
||||
|
||||
return page_dict_sorted
|
||||
|
||||
|
||||
def smart_chunk_one_to_one(
|
||||
text_dict: dict, contract_text: str, one_to_one_fields: FieldSet
|
||||
) -> dict:
|
||||
"""Performs a keyword search on textract outputs to retrieve relevant chunks for a given field.
|
||||
|
||||
Args:
|
||||
text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page
|
||||
contract_text(str): Contract Text
|
||||
|
||||
Raises:
|
||||
ValueError: When the `key_dict` from keyword_mappings has a methodology that is not implemented.
|
||||
Current options are `or` and `hierarchy`.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary keyed by field (defined by the keys in keyword_mappings) valued by
|
||||
corresponding chunks of text (input strings delimited by newlines).
|
||||
"""
|
||||
|
||||
smart_chunk_fields = one_to_one_fields.filter(field_type="smart_chunked")
|
||||
|
||||
chunks = defaultdict(str)
|
||||
|
||||
for field in smart_chunk_fields.fields:
|
||||
# Determine chunking function based on methodology
|
||||
if field.methodology == "or":
|
||||
chunk_func = chunk_or
|
||||
elif field.methodology == "hierarchy":
|
||||
chunk_func = chunk_hierarchical
|
||||
elif field.methodology == "and":
|
||||
chunk_func = chunk_and
|
||||
elif field.methodology == "regex":
|
||||
chunk_func = chunk_regex
|
||||
else:
|
||||
chunk_func = None
|
||||
|
||||
if chunk_func:
|
||||
page_list = chunk_func(
|
||||
text_dict, field.keywords, field.regex_pattern, field.case_sensitive
|
||||
)
|
||||
keyword_page_dict = keyword_search(
|
||||
text_dict, field.keywords, field.case_sensitive, field.regex_pattern
|
||||
)
|
||||
else:
|
||||
page_list = list(text_dict.keys())
|
||||
keyword_page_dict = {"All pages": list(text_dict.keys())}
|
||||
raise ValueError(
|
||||
"methodology must be one of 'hierarchy', 'or', 'and', 'regex', 'include_exclude' - using entire contract as chunk"
|
||||
)
|
||||
|
||||
if len(page_list) == len(text_dict.keys()):
|
||||
chunks[field.field_name] = contract_text
|
||||
else:
|
||||
chunks[field.field_name] = "\n".join(
|
||||
[text_dict[str(page_num)] for page_num in page_list]
|
||||
)
|
||||
|
||||
chunks[field.field_name + "_pages"] = keyword_page_dict
|
||||
chunks[field.field_name + "_methodology"] = field.methodology
|
||||
chunks[field.field_name + "_case"] = field.case_sensitive
|
||||
return chunks
|
||||
@@ -75,6 +75,12 @@ class Field:
|
||||
if "breakout_template" in field_dict
|
||||
else None
|
||||
)
|
||||
# RAG Retrieval Attributes
|
||||
self.retrieval_question = (
|
||||
field_dict["retrieval_question"]
|
||||
if "retrieval_question" in field_dict
|
||||
else None
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
"""Returns a string representation of the Field object."""
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,169 @@
|
||||
"""
|
||||
RAG (Retrieval-Augmented Generation) Utilities
|
||||
|
||||
This module provides configuration and utility functions for RAG-based
|
||||
field extraction using hybrid retrieval (BM25 + semantic search) with reranking.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, List, Dict
|
||||
|
||||
import boto3
|
||||
from langchain.retrievers import EnsembleRetriever
|
||||
from langchain.schema import Document
|
||||
from langchain_aws import BedrockEmbeddings
|
||||
from langchain_community.retrievers import BM25Retriever
|
||||
from langchain_community.vectorstores import FAISS
|
||||
from langchain_community.vectorstores.utils import DistanceStrategy
|
||||
|
||||
|
||||
@dataclass
|
||||
class RAGConfig:
|
||||
"""Configuration for RAG pipeline"""
|
||||
|
||||
CHUNK_SIZE: int = 2000
|
||||
CHUNK_OVERLAP: int = 200
|
||||
SEMANTIC_TOP_K: int = 35
|
||||
KEYWORD_TOP_K: int = 25
|
||||
ENSEMBLE_TOP_K: int = 40
|
||||
RERANKER_TOP_K: int = 15
|
||||
BM25_WEIGHT: float = 0.40
|
||||
VECTOR_WEIGHT: float = 0.60
|
||||
RERANKER: str = "amazon" # "cohere" or "amazon"
|
||||
EMBEDDING_MODEL_ID: str = "amazon.titan-embed-text-v2:0"
|
||||
BEDROCK_REGION: str = "us-east-1"
|
||||
|
||||
|
||||
# Default configuration instance
|
||||
HSC_CONFIG = RAGConfig()
|
||||
|
||||
# Reranker model configurations
|
||||
RERANKER_CONFIGS = {
|
||||
"amazon": {
|
||||
"model_id": "amazon.rerank-v1:0",
|
||||
"name": "Amazon Rerank 1.0",
|
||||
"region": "us-west-2",
|
||||
},
|
||||
"cohere": {
|
||||
"model_id": "cohere.rerank-v3-5:0",
|
||||
"name": "Cohere Rerank 3.5",
|
||||
"region": "us-east-1",
|
||||
},
|
||||
}
|
||||
|
||||
# Initialize AWS Bedrock clients at module level
|
||||
_bedrock_client = None
|
||||
_reranker_clients = {} # Cache per region
|
||||
_embeddings = None
|
||||
|
||||
|
||||
def get_bedrock_client(region_name: str = HSC_CONFIG.BEDROCK_REGION):
|
||||
"""Get or create the Bedrock client."""
|
||||
global _bedrock_client
|
||||
if _bedrock_client is None:
|
||||
_bedrock_client = boto3.client("bedrock-runtime", region_name=region_name)
|
||||
return _bedrock_client
|
||||
|
||||
|
||||
def get_reranker_client(model: str = HSC_CONFIG.RERANKER):
|
||||
"""Get or create the reranker client for specific model."""
|
||||
global _reranker_clients
|
||||
cfg = RERANKER_CONFIGS.get(model, RERANKER_CONFIGS["amazon"])
|
||||
region = cfg["region"]
|
||||
|
||||
if region not in _reranker_clients:
|
||||
_reranker_clients[region] = boto3.client("bedrock-runtime", region_name=region)
|
||||
return _reranker_clients[region]
|
||||
|
||||
|
||||
def get_embeddings(
|
||||
model_id: str = HSC_CONFIG.EMBEDDING_MODEL_ID,
|
||||
region_name: str = HSC_CONFIG.BEDROCK_REGION,
|
||||
):
|
||||
"""Get or create the embeddings model."""
|
||||
global _embeddings
|
||||
if _embeddings is None:
|
||||
client = get_bedrock_client(region_name)
|
||||
_embeddings = BedrockEmbeddings(model_id=model_id, client=client)
|
||||
return _embeddings
|
||||
|
||||
|
||||
def create_vectorstore(chunks: List[Document], embeddings: BedrockEmbeddings) -> FAISS:
|
||||
"""Create FAISS vectorstore from chunks."""
|
||||
return FAISS.from_documents(
|
||||
chunks, embeddings, distance_strategy=DistanceStrategy.COSINE
|
||||
)
|
||||
|
||||
|
||||
def create_hybrid_retriever(
|
||||
chunks: List[Document], vectorstore: FAISS, cfg: RAGConfig = HSC_CONFIG
|
||||
) -> EnsembleRetriever:
|
||||
"""Create hybrid BM25 + semantic retriever."""
|
||||
bm25 = BM25Retriever.from_documents(chunks, k1=1.6, b=0.75)
|
||||
bm25.k = cfg.KEYWORD_TOP_K
|
||||
vector = vectorstore.as_retriever(
|
||||
search_type="similarity", search_kwargs={"k": cfg.SEMANTIC_TOP_K}
|
||||
)
|
||||
return EnsembleRetriever(
|
||||
retrievers=[vector, bm25], weights=[cfg.VECTOR_WEIGHT, cfg.BM25_WEIGHT]
|
||||
)
|
||||
|
||||
|
||||
def rerank_documents(
|
||||
query: str,
|
||||
docs: List[Document],
|
||||
client: Any,
|
||||
top_k: int = HSC_CONFIG.RERANKER_TOP_K,
|
||||
model: str = HSC_CONFIG.RERANKER,
|
||||
) -> List[Document]:
|
||||
"""Rerank documents using AWS Bedrock (Cohere/Amazon)."""
|
||||
if not docs:
|
||||
return docs
|
||||
|
||||
try:
|
||||
cfg = RERANKER_CONFIGS.get(model, RERANKER_CONFIGS["cohere"])
|
||||
|
||||
payload: Dict[str, Any] = (
|
||||
{"query": query, "documents": [{"text": d.page_content} for d in docs]}
|
||||
if model == "amazon"
|
||||
else {
|
||||
"query": query,
|
||||
"documents": [d.page_content for d in docs],
|
||||
"top_n": min(top_k, len(docs)),
|
||||
"api_version": 2,
|
||||
}
|
||||
)
|
||||
|
||||
result = json.loads(
|
||||
client.invoke_model(modelId=cfg["model_id"], body=json.dumps(payload))[
|
||||
"body"
|
||||
].read()
|
||||
)
|
||||
|
||||
scored = sorted(
|
||||
[
|
||||
{"score": item["relevance_score"], "doc": docs[item["index"]]}
|
||||
for item in result["results"]
|
||||
],
|
||||
key=lambda x: x["score"],
|
||||
reverse=True,
|
||||
)[:top_k]
|
||||
|
||||
return [
|
||||
Document(
|
||||
page_content=item["doc"].page_content,
|
||||
metadata={
|
||||
**item["doc"].metadata,
|
||||
"rerank_score": item["score"],
|
||||
"rerank_model": cfg["name"],
|
||||
},
|
||||
)
|
||||
for item in scored
|
||||
]
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
f"Reranking failed: {e} | Selecting {top_k} from base retrieved docs"
|
||||
)
|
||||
return docs[:top_k]
|
||||
@@ -0,0 +1,450 @@
|
||||
"""
|
||||
Unit tests for src.investment.hybrid_smart_chunking_funcs module
|
||||
|
||||
Tests hybrid smart chunking functions including:
|
||||
- Context clipping by token count
|
||||
- Document preparation for chunking
|
||||
- Chunk formatting for LLM
|
||||
- Page number extraction
|
||||
- Preprocessing pipeline
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
from langchain.schema import Document
|
||||
from src.investment.hybrid_smart_chunking_funcs import (
|
||||
clip_context_by_tokens,
|
||||
prepare_documents_for_chunking,
|
||||
format_chunks_for_llm,
|
||||
add_page_numbers_to_chunks,
|
||||
preprocessing_pipeline,
|
||||
)
|
||||
|
||||
|
||||
class TestClipContextByTokens:
|
||||
"""Test clip_context_by_tokens function"""
|
||||
|
||||
def test_context_under_limit(self):
|
||||
"""Test context that doesn't need clipping"""
|
||||
context = "Short text" * 100 # ~1100 chars = ~314 tokens
|
||||
max_tokens = 1000
|
||||
|
||||
result = clip_context_by_tokens(context, max_tokens)
|
||||
|
||||
assert result == context
|
||||
assert len(result) == len(context)
|
||||
|
||||
def test_context_over_limit(self):
|
||||
"""Test context that needs clipping"""
|
||||
context = "A" * 350000 # ~100k tokens
|
||||
max_tokens = 10000 # Clip to 10k tokens
|
||||
|
||||
result = clip_context_by_tokens(context, max_tokens)
|
||||
|
||||
# Should be clipped to approximately 10k tokens (35k chars)
|
||||
assert len(result) < len(context)
|
||||
assert len(result) <= int(max_tokens * 3.5)
|
||||
|
||||
def test_empty_context(self):
|
||||
"""Test empty context"""
|
||||
result = clip_context_by_tokens("")
|
||||
|
||||
assert result == ""
|
||||
|
||||
def test_none_context(self):
|
||||
"""Test None context"""
|
||||
result = clip_context_by_tokens(None)
|
||||
|
||||
assert result == ""
|
||||
|
||||
def test_exact_limit(self):
|
||||
"""Test context exactly at the limit"""
|
||||
max_tokens = 1000
|
||||
context = "A" * int(max_tokens * 3.5) # Exactly at limit
|
||||
|
||||
result = clip_context_by_tokens(context, max_tokens)
|
||||
|
||||
assert result == context
|
||||
|
||||
def test_custom_max_tokens(self):
|
||||
"""Test with custom max_tokens parameter"""
|
||||
context = "A" * 100000
|
||||
max_tokens = 5000
|
||||
|
||||
result = clip_context_by_tokens(context, max_tokens)
|
||||
|
||||
estimated_tokens = int(len(result) / 3.5)
|
||||
assert estimated_tokens <= max_tokens
|
||||
|
||||
def test_token_estimation_accuracy(self):
|
||||
"""Test that token estimation formula is applied correctly"""
|
||||
# 3500 chars should be ~1000 tokens
|
||||
context = "A" * 3500
|
||||
max_tokens = 1000
|
||||
|
||||
result = clip_context_by_tokens(context, max_tokens)
|
||||
|
||||
# Should not clip since it's right at the limit
|
||||
assert len(result) == 3500
|
||||
|
||||
|
||||
class TestPrepareDocumentsForChunking:
|
||||
"""Test prepare_documents_for_chunking function"""
|
||||
|
||||
def test_single_section(self):
|
||||
"""Test preparing a single section"""
|
||||
sections = {
|
||||
"Section 1": {
|
||||
"content": "Section content here",
|
||||
"pages": [1, 2]
|
||||
}
|
||||
}
|
||||
filename = "test_contract.pdf"
|
||||
|
||||
docs = prepare_documents_for_chunking(sections, filename)
|
||||
|
||||
assert len(docs) == 1
|
||||
assert docs[0].page_content == "Section content here"
|
||||
assert docs[0].metadata["section_header"] == "Section 1"
|
||||
assert docs[0].metadata["pages"] == [1, 2]
|
||||
assert docs[0].metadata["filename"] == filename
|
||||
|
||||
def test_multiple_sections(self):
|
||||
"""Test preparing multiple sections"""
|
||||
sections = {
|
||||
"Section 1": {"content": "Content 1", "pages": [1]},
|
||||
"Section 2": {"content": "Content 2", "pages": [2, 3]},
|
||||
"Section 3": {"content": "Content 3", "pages": [4]}
|
||||
}
|
||||
|
||||
docs = prepare_documents_for_chunking(sections)
|
||||
|
||||
assert len(docs) == 3
|
||||
assert all(isinstance(doc, Document) for doc in docs)
|
||||
|
||||
def test_empty_sections(self):
|
||||
"""Test with empty sections dictionary"""
|
||||
sections = {}
|
||||
|
||||
docs = prepare_documents_for_chunking(sections)
|
||||
|
||||
assert docs == []
|
||||
|
||||
def test_section_metadata(self):
|
||||
"""Test that metadata is correctly set"""
|
||||
sections = {
|
||||
"Test Section": {
|
||||
"content": "This is test content",
|
||||
"pages": [5, 6, 7]
|
||||
}
|
||||
}
|
||||
filename = "contract.pdf"
|
||||
|
||||
docs = prepare_documents_for_chunking(sections, filename)
|
||||
|
||||
metadata = docs[0].metadata
|
||||
assert metadata["section_header"] == "Test Section"
|
||||
assert metadata["section_length"] == len("This is test content")
|
||||
assert metadata["pages"] == [5, 6, 7]
|
||||
assert metadata["filename"] == filename
|
||||
|
||||
def test_default_filename(self):
|
||||
"""Test default filename when not provided"""
|
||||
sections = {
|
||||
"Section": {"content": "Content", "pages": [1]}
|
||||
}
|
||||
|
||||
docs = prepare_documents_for_chunking(sections)
|
||||
|
||||
assert docs[0].metadata["filename"] == "contract"
|
||||
|
||||
|
||||
class TestFormatChunksForLLM:
|
||||
"""Test format_chunks_for_llm function"""
|
||||
|
||||
def test_empty_docs(self):
|
||||
"""Test with empty document list"""
|
||||
text_dict = {"1": "Page 1 content"}
|
||||
|
||||
result = format_chunks_for_llm(text_dict, [])
|
||||
|
||||
assert result == ""
|
||||
|
||||
def test_single_document(self):
|
||||
"""Test with single document"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content",
|
||||
"2": "Page 2 content"
|
||||
}
|
||||
docs = [
|
||||
Document(
|
||||
page_content="chunk1",
|
||||
metadata={"pages": [1]}
|
||||
)
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
assert "Page 1 content" in result
|
||||
assert "Page 2 content" not in result
|
||||
|
||||
def test_multiple_documents_different_pages(self):
|
||||
"""Test with documents from different pages"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content",
|
||||
"2": "Page 2 content",
|
||||
"3": "Page 3 content"
|
||||
}
|
||||
docs = [
|
||||
Document(page_content="chunk1", metadata={"pages": [1]}),
|
||||
Document(page_content="chunk2", metadata={"pages": [3]})
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
assert "Page 1 content" in result
|
||||
assert "Page 3 content" in result
|
||||
assert "Page 2 content" not in result
|
||||
|
||||
def test_overlapping_pages(self):
|
||||
"""Test with documents that reference overlapping pages"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content",
|
||||
"2": "Page 2 content"
|
||||
}
|
||||
docs = [
|
||||
Document(page_content="chunk1", metadata={"pages": [1, 2]}),
|
||||
Document(page_content="chunk2", metadata={"pages": [2]})
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
# Both pages should appear only once (set deduplication)
|
||||
assert "Page 1 content" in result
|
||||
assert "Page 2 content" in result
|
||||
|
||||
def test_missing_page_in_text_dict(self):
|
||||
"""Test when document references page not in text_dict"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content"
|
||||
}
|
||||
docs = [
|
||||
Document(page_content="chunk", metadata={"pages": [1, 99]})
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
# Should only include page 1, skip page 99
|
||||
assert "Page 1 content" in result
|
||||
|
||||
def test_none_values_in_text_dict(self):
|
||||
"""Test handling of None values in text_dict"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content",
|
||||
"2": None,
|
||||
"3": "Page 3 content"
|
||||
}
|
||||
docs = [
|
||||
Document(page_content="chunk", metadata={"pages": [1, 2, 3]})
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
# Should skip None value for page 2
|
||||
assert "Page 1 content" in result
|
||||
assert "Page 3 content" in result
|
||||
|
||||
def test_document_without_pages_metadata(self):
|
||||
"""Test document without pages in metadata"""
|
||||
text_dict = {"1": "Page 1 content"}
|
||||
docs = [
|
||||
Document(page_content="chunk", metadata={})
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
# Should return empty or handle gracefully
|
||||
assert isinstance(result, str)
|
||||
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.clip_context_by_tokens')
|
||||
def test_calls_clip_context(self, mock_clip):
|
||||
"""Test that format_chunks_for_llm calls clip_context_by_tokens"""
|
||||
mock_clip.return_value = "clipped context"
|
||||
text_dict = {"1": "Page 1 content"}
|
||||
docs = [Document(page_content="chunk", metadata={"pages": [1]})]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
mock_clip.assert_called_once()
|
||||
assert result == "clipped context"
|
||||
|
||||
def test_integer_page_numbers(self):
|
||||
"""Test handling of integer page numbers (converted to strings)"""
|
||||
text_dict = {
|
||||
"1": "Page 1 content",
|
||||
"2": "Page 2 content"
|
||||
}
|
||||
docs = [
|
||||
Document(page_content="chunk", metadata={"pages": [1, 2]}) # Integers
|
||||
]
|
||||
|
||||
result = format_chunks_for_llm(text_dict, docs)
|
||||
|
||||
assert "Page 1 content" in result
|
||||
assert "Page 2 content" in result
|
||||
|
||||
|
||||
class TestAddPageNumbersToChunks:
|
||||
"""Test add_page_numbers_to_chunks function"""
|
||||
|
||||
def test_single_section_with_page_markers(self):
|
||||
"""Test extracting page numbers from section with markers"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No. = 1\nContent here\nStart of Page No. = 2\nMore content"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
assert "Section 1" in result
|
||||
assert result["Section 1"]["content"] == sections["Section 1"]
|
||||
assert 1 in result["Section 1"]["pages"]
|
||||
assert 2 in result["Section 1"]["pages"]
|
||||
|
||||
def test_multiple_sections(self):
|
||||
"""Test multiple sections with page markers"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No. = 5\nContent",
|
||||
"Section 2": "Start of Page No. = 6\nMore content\nStart of Page No. = 7\nEven more"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
assert 5 in result["Section 1"]["pages"]
|
||||
assert 6 in result["Section 2"]["pages"]
|
||||
assert 7 in result["Section 2"]["pages"]
|
||||
|
||||
def test_section_without_markers(self):
|
||||
"""Test section without page markers"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No. = 1\nContent",
|
||||
"Section 2": "Content without markers"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
# Section 2 should inherit last page from Section 1
|
||||
assert result["Section 2"]["pages"] == [1]
|
||||
|
||||
def test_empty_sections(self):
|
||||
"""Test with empty sections dictionary"""
|
||||
sections = {}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
assert result == {}
|
||||
|
||||
def test_page_number_deduplication(self):
|
||||
"""Test that duplicate page numbers are removed"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No. = 3\nContent\nStart of Page No. = 3\nMore"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
# Should have unique page numbers
|
||||
assert result["Section 1"]["pages"].count(3) == 1
|
||||
|
||||
def test_page_numbers_are_sorted(self):
|
||||
"""Test that page numbers are sorted"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No. = 5\nStart of Page No. = 2\nStart of Page No. = 8"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
pages = result["Section 1"]["pages"]
|
||||
assert pages == sorted(pages)
|
||||
|
||||
def test_alternative_marker_format(self):
|
||||
"""Test alternative page marker format without period"""
|
||||
sections = {
|
||||
"Section 1": "Start of Page No = 10\nContent"
|
||||
}
|
||||
|
||||
result = add_page_numbers_to_chunks(sections)
|
||||
|
||||
assert 10 in result["Section 1"]["pages"]
|
||||
|
||||
|
||||
class TestPreprocessingPipeline:
|
||||
"""Test preprocessing_pipeline function"""
|
||||
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
||||
def test_pipeline_calls_functions(self, mock_add_pages, mock_split):
|
||||
"""Test that pipeline calls the necessary functions"""
|
||||
contract_text = "Test contract text"
|
||||
|
||||
# Mock split_contract_text_into_sections return
|
||||
mock_sections = {"Section 1": "content"}
|
||||
mock_verification = (True, "Perfect preservation")
|
||||
mock_split.return_value = (mock_sections, contract_text, "\n", mock_verification)
|
||||
|
||||
# Mock add_page_numbers_to_chunks return
|
||||
mock_sections_with_pages = {
|
||||
"Section 1": {"content": "content", "pages": [1]}
|
||||
}
|
||||
mock_add_pages.return_value = mock_sections_with_pages
|
||||
|
||||
sections, is_preserved = preprocessing_pipeline(contract_text)
|
||||
|
||||
mock_split.assert_called_once_with(contract_text)
|
||||
mock_add_pages.assert_called_once()
|
||||
assert is_preserved is True
|
||||
assert sections == mock_sections_with_pages
|
||||
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
||||
def test_pipeline_returns_preservation_status(self, mock_add_pages, mock_split):
|
||||
"""Test that pipeline returns preservation status"""
|
||||
# Test successful preservation
|
||||
mock_split.return_value = ({}, "text", "\n", (True, "Success"))
|
||||
mock_add_pages.return_value = {}
|
||||
|
||||
sections, is_preserved = preprocessing_pipeline("text")
|
||||
assert is_preserved is True
|
||||
|
||||
# Test failed preservation
|
||||
mock_split.return_value = ({}, "text", "\n", (False, "Failed"))
|
||||
mock_add_pages.return_value = {}
|
||||
|
||||
sections, is_preserved = preprocessing_pipeline("text")
|
||||
assert is_preserved is False
|
||||
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.hybrid_smart_chunking_preprocessing.split_contract_text_into_sections')
|
||||
@patch('src.investment.hybrid_smart_chunking_funcs.add_page_numbers_to_chunks')
|
||||
def test_pipeline_with_real_text(self, mock_add_pages, mock_split):
|
||||
"""Test pipeline with realistic contract text"""
|
||||
contract_text = """1 Definitions
|
||||
Definition content
|
||||
2 Terms
|
||||
Terms content"""
|
||||
|
||||
mock_sections = {
|
||||
"1 Definitions": "Definition content",
|
||||
"2 Terms": "Terms content"
|
||||
}
|
||||
mock_split.return_value = (mock_sections, contract_text, "\n", (True, "OK"))
|
||||
|
||||
mock_sections_with_pages = {
|
||||
"1 Definitions": {"content": "Definition content", "pages": [1]},
|
||||
"2 Terms": {"content": "Terms content", "pages": [2]}
|
||||
}
|
||||
mock_add_pages.return_value = mock_sections_with_pages
|
||||
|
||||
sections, is_preserved = preprocessing_pipeline(contract_text)
|
||||
|
||||
assert len(sections) == 2
|
||||
assert "1 Definitions" in sections
|
||||
assert "2 Terms" in sections
|
||||
assert is_preserved is True
|
||||
@@ -0,0 +1,353 @@
|
||||
"""
|
||||
Unit tests for src.investment.hybrid_smart_chunking_preprocessing module
|
||||
|
||||
Tests contract text preprocessing and section splitting functionality including:
|
||||
- Section identification based on various heading patterns
|
||||
- Text preservation verification
|
||||
- Section metadata extraction
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from src.investment.hybrid_smart_chunking_preprocessing import (
|
||||
split_contract_text_into_sections,
|
||||
verify_preservation,
|
||||
get_section_info,
|
||||
)
|
||||
|
||||
|
||||
class TestSplitContractTextIntoSections:
|
||||
"""Test split_contract_text_into_sections function"""
|
||||
|
||||
def test_simple_text_no_headings(self):
|
||||
"""Test text with no recognizable headings"""
|
||||
text = "This is a simple contract text.\nWith no headings at all."
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "Preamble" in sections
|
||||
assert sections["Preamble"] == text
|
||||
assert original == text
|
||||
assert line_ending == "\n"
|
||||
assert verification[0] is True
|
||||
|
||||
def test_arabic_numbered_headings(self):
|
||||
"""Test single-digit arabic numbered headings (1. Title)"""
|
||||
text = "Preamble text\n1 Definitions\nDefinition content\n2 Terms\nTerms content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "Preamble" in sections
|
||||
assert "1 Definitions" in sections
|
||||
assert "2 Terms" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_sub_arabic_numbered_headings(self):
|
||||
"""Test multi-level arabic numbered headings (1.1, 1.2, etc.)"""
|
||||
text = "1.1 First Section\nContent A\n1.2 Second Section\nContent B"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "1.1 First Section" in sections
|
||||
assert "1.2 Second Section" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_roman_numeral_headings(self):
|
||||
"""Test roman numeral headings (I. Title, II. Title)"""
|
||||
text = "I Introduction\nIntro content\nII Background\nBackground content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "I Introduction" in sections
|
||||
assert "II Background" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_article_headings(self):
|
||||
"""Test ARTICLE headings (ARTICLE I Title)"""
|
||||
text = "ARTICLE I Definitions\nDef content\nARTICLE II Terms\nTerms content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "ARTICLE I Definitions" in sections
|
||||
assert "ARTICLE II Terms" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_exhibit_headings(self):
|
||||
"""Test EXHIBIT headings (EXHIBIT A, EXHIBIT B Title)"""
|
||||
text = "Main text\nEXHIBIT A\nExhibit A content\nEXHIBIT B Schedule\nExhibit B content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "EXHIBIT A" in sections
|
||||
assert "EXHIBIT B Schedule" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_attachment_headings(self):
|
||||
"""Test Attachment headings"""
|
||||
text = "Contract text\nAttachment 1\nAttachment content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "Attachment 1" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_recitals_and_agreement_keywords(self):
|
||||
"""Test RECITALS and AGREEMENT keyword headings"""
|
||||
text = "RECITALS:\nRecital text\nAGREEMENT\nAgreement text"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "RECITALS" in sections
|
||||
assert "AGREEMENT" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_in_witness_whereof_heading(self):
|
||||
"""Test IN WITNESS WHEREOF heading"""
|
||||
text = "Contract body\nIN WITNESS WHEREOF\nSignature section"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "IN WITNESS WHEREOF" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_signatory_headings(self):
|
||||
"""Test signatory block headings (PROVIDER:, ICM:)"""
|
||||
text = "Agreement text\nPROVIDER:\nBy: John Doe\nICM:\nBy: Jane Smith"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "PROVIDER" in sections or "ICM" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_multi_line_headings(self):
|
||||
"""Test headings that span multiple lines"""
|
||||
text = "1\nDefinitions and Terms\nContent here\n2\nOther Section\nMore content"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
# Should combine multi-line headings
|
||||
assert any("Definitions" in key for key in sections.keys())
|
||||
assert verification[0] is True
|
||||
|
||||
def test_duplicate_section_titles(self):
|
||||
"""Test handling of duplicate section titles"""
|
||||
text = "1 Terms\nFirst terms\n1 Terms\nSecond terms"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
# Should have unique keys (second occurrence gets suffix)
|
||||
assert "1 Terms" in sections
|
||||
assert "1 Terms (2)" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_windows_line_endings(self):
|
||||
"""Test text with Windows line endings (\\r\\n)"""
|
||||
text = "Section 1\r\nContent A\r\nSection 2\r\nContent B"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert line_ending == "\r\n"
|
||||
assert verification[0] is True
|
||||
|
||||
def test_unix_line_endings(self):
|
||||
"""Test text with Unix line endings (\\n)"""
|
||||
text = "Section 1\nContent A\nSection 2\nContent B"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert line_ending == "\n"
|
||||
assert verification[0] is True
|
||||
|
||||
def test_no_line_endings(self):
|
||||
"""Test text with no line endings"""
|
||||
text = "Single line of text with no breaks"
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert line_ending == ""
|
||||
assert "Preamble" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
def test_empty_text(self):
|
||||
"""Test empty text input"""
|
||||
text = ""
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
# Should handle empty text gracefully
|
||||
assert isinstance(sections, dict)
|
||||
assert original == ""
|
||||
|
||||
def test_mixed_heading_types(self):
|
||||
"""Test document with multiple types of headings"""
|
||||
text = """RECITALS:
|
||||
Whereas clause
|
||||
1 Definitions
|
||||
Term definitions
|
||||
ARTICLE I Scope
|
||||
Scope content
|
||||
EXHIBIT A
|
||||
Exhibit content"""
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(text)
|
||||
|
||||
assert "RECITALS" in sections
|
||||
assert "1 Definitions" in sections
|
||||
assert "ARTICLE I Scope" in sections
|
||||
# EXHIBIT A might combine with next line, so check for both possibilities
|
||||
assert "EXHIBIT A" in sections or "EXHIBIT A Exhibit content" in sections
|
||||
assert verification[0] is True
|
||||
|
||||
|
||||
class TestVerifyPreservation:
|
||||
"""Test verify_preservation function"""
|
||||
|
||||
def test_perfect_preservation(self):
|
||||
"""Test when content is perfectly preserved"""
|
||||
original = "Line 1\nLine 2\nLine 3"
|
||||
sections = {"Section1": "Line 1\nLine 2", "Section2": "Line 3"}
|
||||
line_ending = "\n"
|
||||
|
||||
is_preserved, message = verify_preservation(original, sections, line_ending)
|
||||
|
||||
assert is_preserved is True
|
||||
assert "Perfect preservation" in message
|
||||
|
||||
def test_content_mismatch(self):
|
||||
"""Test when content differs"""
|
||||
original = "Line 1\nLine 2\nLine 3"
|
||||
sections = {"Section1": "Line 1\nLine X"} # Different content
|
||||
line_ending = "\n"
|
||||
|
||||
is_preserved, message = verify_preservation(original, sections, line_ending)
|
||||
|
||||
assert is_preserved is False
|
||||
assert "First difference at position" in message or "Length mismatch" in message
|
||||
|
||||
def test_length_mismatch(self):
|
||||
"""Test when reconstructed content has different length"""
|
||||
original = "Short text"
|
||||
sections = {"Section1": "Much longer text than original"}
|
||||
line_ending = "\n"
|
||||
|
||||
is_preserved, message = verify_preservation(original, sections, line_ending)
|
||||
|
||||
assert is_preserved is False
|
||||
|
||||
def test_empty_sections(self):
|
||||
"""Test with empty sections dict"""
|
||||
original = "Some content"
|
||||
sections = {}
|
||||
line_ending = "\n"
|
||||
|
||||
is_preserved, message = verify_preservation(original, sections, line_ending)
|
||||
|
||||
assert is_preserved is False
|
||||
|
||||
def test_windows_line_endings_preservation(self):
|
||||
"""Test preservation with Windows line endings"""
|
||||
original = "Line 1\r\nLine 2\r\nLine 3"
|
||||
sections = {"Section1": "Line 1\r\nLine 2", "Section2": "Line 3"}
|
||||
line_ending = "\r\n"
|
||||
|
||||
is_preserved, message = verify_preservation(original, sections, line_ending)
|
||||
|
||||
assert is_preserved is True
|
||||
|
||||
|
||||
class TestGetSectionInfo:
|
||||
"""Test get_section_info function"""
|
||||
|
||||
def test_basic_section_info(self):
|
||||
"""Test extracting basic section information"""
|
||||
sections = {
|
||||
"Section 1": "This is section one content",
|
||||
"Section 2": "This is section two\nWith multiple lines"
|
||||
}
|
||||
|
||||
info = get_section_info(sections)
|
||||
|
||||
assert info["total_sections"] == 2
|
||||
assert "Section 1" in info["section_details"]
|
||||
assert "Section 2" in info["section_details"]
|
||||
assert info["section_details"]["Section 1"]["character_count"] == 27
|
||||
assert info["section_details"]["Section 2"]["line_count"] == 2
|
||||
|
||||
def test_empty_sections(self):
|
||||
"""Test with empty sections dict"""
|
||||
sections = {}
|
||||
|
||||
info = get_section_info(sections)
|
||||
|
||||
assert info["total_sections"] == 0
|
||||
assert info["section_details"] == {}
|
||||
assert info["total_characters"] == 0
|
||||
|
||||
def test_section_with_empty_content(self):
|
||||
"""Test section with empty string content"""
|
||||
sections = {"Empty Section": ""}
|
||||
|
||||
info = get_section_info(sections)
|
||||
|
||||
assert info["total_sections"] == 1
|
||||
assert info["section_details"]["Empty Section"]["character_count"] == 0
|
||||
|
||||
def test_total_character_count(self):
|
||||
"""Test total character count calculation"""
|
||||
sections = {
|
||||
"Section 1": "12345", # 5 chars
|
||||
"Section 2": "67890", # 5 chars
|
||||
"Section 3": "ABCDE" # 5 chars
|
||||
}
|
||||
|
||||
info = get_section_info(sections)
|
||||
|
||||
assert info["total_characters"] == 15
|
||||
|
||||
def test_line_count_calculation(self):
|
||||
"""Test line count is calculated correctly"""
|
||||
sections = {
|
||||
"Single Line": "One line only",
|
||||
"Multi Line": "Line 1\nLine 2\nLine 3"
|
||||
}
|
||||
|
||||
info = get_section_info(sections)
|
||||
|
||||
assert info["section_details"]["Single Line"]["line_count"] == 1
|
||||
assert info["section_details"]["Multi Line"]["line_count"] == 3
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for the full preprocessing workflow"""
|
||||
|
||||
def test_realistic_contract_structure(self):
|
||||
"""Test with a realistic contract structure"""
|
||||
contract_text = """PROVIDER AGREEMENT
|
||||
|
||||
RECITALS:
|
||||
This Agreement is made between the parties.
|
||||
|
||||
AGREEMENT
|
||||
|
||||
1 Definitions
|
||||
For purposes of this Agreement, the following terms shall have the meanings set forth below.
|
||||
|
||||
1.1 Provider
|
||||
Provider means the healthcare provider.
|
||||
|
||||
1.2 Services
|
||||
Services means healthcare services.
|
||||
|
||||
2 Terms and Conditions
|
||||
The parties agree to the following terms.
|
||||
|
||||
EXHIBIT A Payment Schedule
|
||||
Payment details here.
|
||||
|
||||
IN WITNESS WHEREOF
|
||||
The parties have executed this Agreement.
|
||||
|
||||
PROVIDER:
|
||||
By: _______________
|
||||
Name: John Doe
|
||||
"""
|
||||
|
||||
sections, original, line_ending, verification = split_contract_text_into_sections(contract_text)
|
||||
|
||||
# Verify major sections are identified
|
||||
assert "RECITALS" in sections
|
||||
assert "AGREEMENT" in sections
|
||||
assert any("Definitions" in key for key in sections.keys())
|
||||
assert "EXHIBIT A Payment Schedule" in sections
|
||||
assert "IN WITNESS WHEREOF" in sections
|
||||
|
||||
# Verify preservation
|
||||
assert verification[0] is True
|
||||
|
||||
# Verify section info
|
||||
info = get_section_info(sections)
|
||||
assert info["total_sections"] > 0
|
||||
# Total characters should match (allow for slight differences in section joining)
|
||||
assert abs(info["total_characters"] - len(contract_text)) < 20
|
||||
@@ -0,0 +1,332 @@
|
||||
"""
|
||||
Unit tests for src.utils.rag_utils module
|
||||
|
||||
Tests RAG (Retrieval-Augmented Generation) utility functions including:
|
||||
- RAGConfig configuration
|
||||
- Client initialization and caching
|
||||
- Vectorstore creation
|
||||
- Hybrid retriever creation
|
||||
- Document reranking
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
from langchain.schema import Document
|
||||
from src.utils.rag_utils import (
|
||||
RAGConfig,
|
||||
HSC_CONFIG,
|
||||
RERANKER_CONFIGS,
|
||||
get_bedrock_client,
|
||||
get_reranker_client,
|
||||
get_embeddings,
|
||||
create_vectorstore,
|
||||
create_hybrid_retriever,
|
||||
rerank_documents,
|
||||
)
|
||||
|
||||
|
||||
class TestRAGConfig:
|
||||
"""Test RAGConfig dataclass"""
|
||||
|
||||
def test_default_config(self):
|
||||
"""Test that HSC_CONFIG has expected default values"""
|
||||
assert HSC_CONFIG.CHUNK_SIZE == 2000
|
||||
assert HSC_CONFIG.CHUNK_OVERLAP == 200
|
||||
assert HSC_CONFIG.SEMANTIC_TOP_K == 35
|
||||
assert HSC_CONFIG.KEYWORD_TOP_K == 25
|
||||
assert HSC_CONFIG.ENSEMBLE_TOP_K == 40
|
||||
assert HSC_CONFIG.RERANKER_TOP_K == 15
|
||||
assert HSC_CONFIG.BM25_WEIGHT == 0.40
|
||||
assert HSC_CONFIG.VECTOR_WEIGHT == 0.60
|
||||
assert HSC_CONFIG.RERANKER == "amazon"
|
||||
assert HSC_CONFIG.EMBEDDING_MODEL_ID == "amazon.titan-embed-text-v2:0"
|
||||
assert HSC_CONFIG.BEDROCK_REGION == "us-east-1"
|
||||
|
||||
def test_custom_config(self):
|
||||
"""Test creating custom RAGConfig"""
|
||||
custom_config = RAGConfig(
|
||||
CHUNK_SIZE=1000,
|
||||
CHUNK_OVERLAP=100,
|
||||
RERANKER="cohere"
|
||||
)
|
||||
assert custom_config.CHUNK_SIZE == 1000
|
||||
assert custom_config.CHUNK_OVERLAP == 100
|
||||
assert custom_config.RERANKER == "cohere"
|
||||
# Defaults should still apply
|
||||
assert custom_config.SEMANTIC_TOP_K == 35
|
||||
|
||||
def test_reranker_configs(self):
|
||||
"""Test RERANKER_CONFIGS dictionary structure"""
|
||||
assert "amazon" in RERANKER_CONFIGS
|
||||
assert "cohere" in RERANKER_CONFIGS
|
||||
|
||||
# Amazon config
|
||||
assert RERANKER_CONFIGS["amazon"]["model_id"] == "amazon.rerank-v1:0"
|
||||
assert RERANKER_CONFIGS["amazon"]["name"] == "Amazon Rerank 1.0"
|
||||
assert RERANKER_CONFIGS["amazon"]["region"] == "us-west-2"
|
||||
|
||||
# Cohere config
|
||||
assert RERANKER_CONFIGS["cohere"]["model_id"] == "cohere.rerank-v3-5:0"
|
||||
assert RERANKER_CONFIGS["cohere"]["name"] == "Cohere Rerank 3.5"
|
||||
assert RERANKER_CONFIGS["cohere"]["region"] == "us-east-1"
|
||||
|
||||
|
||||
class TestClientFunctions:
|
||||
"""Test client initialization and caching functions"""
|
||||
|
||||
@patch('src.utils.rag_utils.boto3.client')
|
||||
def test_get_bedrock_client(self, mock_boto_client):
|
||||
"""Test Bedrock client creation and caching"""
|
||||
# Reset module-level cache
|
||||
import src.utils.rag_utils as rag_utils
|
||||
rag_utils._bedrock_client = None
|
||||
|
||||
mock_client = Mock()
|
||||
mock_boto_client.return_value = mock_client
|
||||
|
||||
# First call should create client
|
||||
client1 = get_bedrock_client()
|
||||
assert client1 == mock_client
|
||||
mock_boto_client.assert_called_once_with("bedrock-runtime", region_name="us-east-1")
|
||||
|
||||
# Second call should return cached client
|
||||
client2 = get_bedrock_client()
|
||||
assert client2 == mock_client
|
||||
assert mock_boto_client.call_count == 1 # Should not create new client
|
||||
|
||||
@patch('src.utils.rag_utils.boto3.client')
|
||||
def test_get_reranker_client_amazon(self, mock_boto_client):
|
||||
"""Test Amazon reranker client creation"""
|
||||
# Reset module-level cache
|
||||
import src.utils.rag_utils as rag_utils
|
||||
rag_utils._reranker_clients = {}
|
||||
|
||||
mock_client = Mock()
|
||||
mock_boto_client.return_value = mock_client
|
||||
|
||||
client = get_reranker_client("amazon")
|
||||
assert client == mock_client
|
||||
mock_boto_client.assert_called_with("bedrock-runtime", region_name="us-west-2")
|
||||
|
||||
@patch('src.utils.rag_utils.boto3.client')
|
||||
def test_get_reranker_client_cohere(self, mock_boto_client):
|
||||
"""Test Cohere reranker client creation"""
|
||||
# Reset module-level cache
|
||||
import src.utils.rag_utils as rag_utils
|
||||
rag_utils._reranker_clients = {}
|
||||
|
||||
mock_client = Mock()
|
||||
mock_boto_client.return_value = mock_client
|
||||
|
||||
client = get_reranker_client("cohere")
|
||||
assert client == mock_client
|
||||
mock_boto_client.assert_called_with("bedrock-runtime", region_name="us-east-1")
|
||||
|
||||
@patch('src.utils.rag_utils.get_bedrock_client')
|
||||
@patch('src.utils.rag_utils.BedrockEmbeddings')
|
||||
def test_get_embeddings(self, mock_embeddings_class, mock_get_client):
|
||||
"""Test embeddings model creation and caching"""
|
||||
# Reset module-level cache
|
||||
import src.utils.rag_utils as rag_utils
|
||||
rag_utils._embeddings = None
|
||||
|
||||
mock_client = Mock()
|
||||
mock_get_client.return_value = mock_client
|
||||
mock_embeddings = Mock()
|
||||
mock_embeddings_class.return_value = mock_embeddings
|
||||
|
||||
# First call should create embeddings
|
||||
embeddings1 = get_embeddings()
|
||||
assert embeddings1 == mock_embeddings
|
||||
mock_embeddings_class.assert_called_once_with(
|
||||
model_id="amazon.titan-embed-text-v2:0",
|
||||
client=mock_client
|
||||
)
|
||||
|
||||
# Second call should return cached embeddings
|
||||
embeddings2 = get_embeddings()
|
||||
assert embeddings2 == mock_embeddings
|
||||
assert mock_embeddings_class.call_count == 1
|
||||
|
||||
|
||||
class TestVectorstoreFunctions:
|
||||
"""Test vectorstore and retriever creation functions"""
|
||||
|
||||
@patch('src.utils.rag_utils.FAISS.from_documents')
|
||||
def test_create_vectorstore(self, mock_faiss):
|
||||
"""Test FAISS vectorstore creation"""
|
||||
mock_vectorstore = Mock()
|
||||
mock_faiss.return_value = mock_vectorstore
|
||||
mock_embeddings = Mock()
|
||||
|
||||
chunks = [
|
||||
Document(page_content="test content 1"),
|
||||
Document(page_content="test content 2")
|
||||
]
|
||||
|
||||
result = create_vectorstore(chunks, mock_embeddings)
|
||||
|
||||
assert result == mock_vectorstore
|
||||
mock_faiss.assert_called_once()
|
||||
# Verify chunks and embeddings were passed
|
||||
call_args = mock_faiss.call_args
|
||||
assert call_args[0][0] == chunks
|
||||
assert call_args[0][1] == mock_embeddings
|
||||
|
||||
@patch('src.utils.rag_utils.BM25Retriever.from_documents')
|
||||
@patch('src.utils.rag_utils.EnsembleRetriever')
|
||||
def test_create_hybrid_retriever(self, mock_ensemble, mock_bm25):
|
||||
"""Test hybrid retriever creation with BM25 and semantic search"""
|
||||
# Setup mocks
|
||||
mock_bm25_retriever = Mock()
|
||||
mock_bm25.return_value = mock_bm25_retriever
|
||||
mock_ensemble_retriever = Mock()
|
||||
mock_ensemble.return_value = mock_ensemble_retriever
|
||||
|
||||
mock_vectorstore = Mock()
|
||||
mock_vector_retriever = Mock()
|
||||
mock_vectorstore.as_retriever.return_value = mock_vector_retriever
|
||||
|
||||
chunks = [Document(page_content="test")]
|
||||
config = RAGConfig()
|
||||
|
||||
result = create_hybrid_retriever(chunks, mock_vectorstore, config)
|
||||
|
||||
assert result == mock_ensemble_retriever
|
||||
# Verify BM25 was created
|
||||
mock_bm25.assert_called_once()
|
||||
# Verify vectorstore retriever was created
|
||||
mock_vectorstore.as_retriever.assert_called_once()
|
||||
# Verify ensemble was created with correct weights
|
||||
call_args = mock_ensemble.call_args[1]
|
||||
assert call_args['weights'] == [config.VECTOR_WEIGHT, config.BM25_WEIGHT]
|
||||
|
||||
|
||||
class TestRerankDocuments:
|
||||
"""Test document reranking functionality"""
|
||||
|
||||
def test_rerank_empty_docs(self):
|
||||
"""Test reranking with empty document list"""
|
||||
result = rerank_documents("query", [], Mock())
|
||||
assert result == []
|
||||
|
||||
@patch('src.utils.rag_utils.json.loads')
|
||||
def test_rerank_amazon_model(self, mock_json_loads):
|
||||
"""Test reranking with Amazon model"""
|
||||
# Setup
|
||||
mock_client = Mock()
|
||||
mock_response = {
|
||||
"body": Mock(read=Mock(return_value=b'{}'))
|
||||
}
|
||||
mock_client.invoke_model.return_value = mock_response
|
||||
mock_json_loads.return_value = {
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
{"index": 1, "relevance_score": 0.7}
|
||||
]
|
||||
}
|
||||
|
||||
docs = [
|
||||
Document(page_content="doc1", metadata={"source": "test"}),
|
||||
Document(page_content="doc2", metadata={"source": "test"})
|
||||
]
|
||||
|
||||
result = rerank_documents(
|
||||
query="test query",
|
||||
docs=docs,
|
||||
client=mock_client,
|
||||
top_k=2,
|
||||
model="amazon"
|
||||
)
|
||||
|
||||
# Verify results
|
||||
assert len(result) == 2
|
||||
assert result[0].page_content == "doc1"
|
||||
assert result[0].metadata["rerank_score"] == 0.9
|
||||
assert result[1].page_content == "doc2"
|
||||
assert result[1].metadata["rerank_score"] == 0.7
|
||||
|
||||
@patch('src.utils.rag_utils.json.loads')
|
||||
def test_rerank_cohere_model(self, mock_json_loads):
|
||||
"""Test reranking with Cohere model"""
|
||||
mock_client = Mock()
|
||||
mock_response = {
|
||||
"body": Mock(read=Mock(return_value=b'{}'))
|
||||
}
|
||||
mock_client.invoke_model.return_value = mock_response
|
||||
mock_json_loads.return_value = {
|
||||
"results": [
|
||||
{"index": 1, "relevance_score": 0.8},
|
||||
{"index": 0, "relevance_score": 0.6}
|
||||
]
|
||||
}
|
||||
|
||||
docs = [
|
||||
Document(page_content="doc1"),
|
||||
Document(page_content="doc2")
|
||||
]
|
||||
|
||||
result = rerank_documents(
|
||||
query="test query",
|
||||
docs=docs,
|
||||
client=mock_client,
|
||||
top_k=2,
|
||||
model="cohere"
|
||||
)
|
||||
|
||||
# Verify reranked order (doc2 should be first due to higher score)
|
||||
assert len(result) == 2
|
||||
assert result[0].page_content == "doc2"
|
||||
assert result[1].page_content == "doc1"
|
||||
|
||||
@patch('src.utils.rag_utils.json.loads', side_effect=Exception("API Error"))
|
||||
@patch('src.utils.rag_utils.logging')
|
||||
def test_rerank_error_handling(self, mock_logging, mock_json_loads):
|
||||
"""Test that reranking failures return top_k documents"""
|
||||
docs = [
|
||||
Document(page_content=f"doc{i}") for i in range(10)
|
||||
]
|
||||
mock_client = Mock()
|
||||
|
||||
result = rerank_documents(
|
||||
query="test",
|
||||
docs=docs,
|
||||
client=mock_client,
|
||||
top_k=3,
|
||||
model="amazon"
|
||||
)
|
||||
|
||||
# Should return first top_k docs on error
|
||||
assert len(result) == 3
|
||||
assert result[0].page_content == "doc0"
|
||||
# Verify error was logged
|
||||
mock_logging.error.assert_called_once()
|
||||
|
||||
def test_rerank_top_k_limit(self):
|
||||
"""Test that reranking respects top_k limit"""
|
||||
mock_client = Mock()
|
||||
mock_response = {
|
||||
"body": Mock(read=Mock(return_value=b'{"results": [{"index": 0, "relevance_score": 0.9}, {"index": 1, "relevance_score": 0.8}, {"index": 2, "relevance_score": 0.7}]}'))
|
||||
}
|
||||
mock_client.invoke_model.return_value = mock_response
|
||||
|
||||
docs = [Document(page_content=f"doc{i}") for i in range(3)]
|
||||
|
||||
with patch('src.utils.rag_utils.json.loads') as mock_json:
|
||||
mock_json.return_value = {
|
||||
"results": [
|
||||
{"index": 0, "relevance_score": 0.9},
|
||||
{"index": 1, "relevance_score": 0.8},
|
||||
{"index": 2, "relevance_score": 0.7}
|
||||
]
|
||||
}
|
||||
|
||||
result = rerank_documents(
|
||||
query="test",
|
||||
docs=docs,
|
||||
client=mock_client,
|
||||
top_k=2,
|
||||
model="amazon"
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
@@ -1,72 +0,0 @@
|
||||
import pytest
|
||||
from src.investment.smart_chunking_funcs import parse_chunk, stitch_chunks
|
||||
|
||||
|
||||
def test_parse_chunk_valid():
|
||||
chunk = "chunk 001: This is a test chunk."
|
||||
chunk_number, content = parse_chunk(chunk)
|
||||
assert chunk_number == 1
|
||||
assert content == "This is a test chunk."
|
||||
|
||||
|
||||
def test_parse_chunk_valid_with_large_number():
|
||||
chunk = "chunk 123: Another test chunk with a larger number."
|
||||
chunk_number, content = parse_chunk(chunk)
|
||||
assert chunk_number == 123
|
||||
assert content == "Another test chunk with a larger number."
|
||||
|
||||
|
||||
def test_parse_chunk_empty_content():
|
||||
chunk = "chunk 005: "
|
||||
chunk_number, content = parse_chunk(chunk)
|
||||
assert chunk_number == 5
|
||||
assert content == ""
|
||||
|
||||
|
||||
def test_parse_chunk_invalid_format():
|
||||
chunk = "invalid chunk format"
|
||||
with pytest.raises(ValueError):
|
||||
parse_chunk(chunk)
|
||||
|
||||
|
||||
def test_parse_chunk_short_prefix():
|
||||
chunk = "chunk 01: Short prefix"
|
||||
with pytest.raises(ValueError):
|
||||
parse_chunk(chunk)
|
||||
|
||||
|
||||
def test_stitch_chunks_empty_list():
|
||||
assert stitch_chunks([]) == ""
|
||||
|
||||
|
||||
def test_stitch_chunks_single_chunk():
|
||||
assert stitch_chunks(["chunk 001: hello world"]) == "hello world"
|
||||
|
||||
|
||||
def test_stitch_chunks_two_chunks():
|
||||
chunks = ["chunk 001: hello world", "chunk 002: world and more"]
|
||||
assert stitch_chunks(chunks, overlap_size=5) == "hello world and more"
|
||||
|
||||
|
||||
def test_stitch_chunks_out_of_order():
|
||||
chunks = ["chunk 002: world", "chunk 001: hello"]
|
||||
assert stitch_chunks(chunks, overlap_size=0) in ["helloworld", "hello\nworld"]
|
||||
|
||||
|
||||
def test_stitch_noncontiguous_chunks():
|
||||
chunks = ["chunk 001: hello", "chunk 003: world"]
|
||||
assert stitch_chunks(chunks, overlap_size=10) in [
|
||||
"helloworld",
|
||||
"hello\nworld",
|
||||
] # the overlap size will be ignored in the noncontiguous case
|
||||
|
||||
|
||||
def test_stitch_long_overlap(): # not enough text to overlap, so we just append them
|
||||
chunks = ["chunk 001: hello", "chunk 002: world"]
|
||||
assert stitch_chunks(chunks, overlap_size=10) in ["helloworld", "hello\nworld"]
|
||||
|
||||
|
||||
def test_stitch_chunks_negative_overlap():
|
||||
chunks = ["chunk 001: hello", "chunk 002: world"]
|
||||
with pytest.raises(ValueError):
|
||||
stitch_chunks(chunks, overlap_size=-1)
|
||||
Reference in New Issue
Block a user