Merged in feature/dtc-clean (pull request #767)
Feature/dtc clean * dtc v1 initial commit * second layer of prompt * Added MAX_PAGES_TO_CHECK and page looping to get answer from llm * Edited contract_type prompt to get cleaner output with no additional text * Updated document_type_classification prompt * Added non-contract type classifcation * Added dtc_v1_test.py for testing only * Updated prompt * Updated prompt * at scale changes * add save feature for individual file output * update dockerfile * clean dtc to prepare for merge * clean up dtc script * Merge remote-tracking branch 'origin/main' into feature/dtc-clean * reafactor code * Merge remote-tracking branch 'origin/main' into feature/dtc-clean Approved-by: Katon Minhas
This commit is contained in:
Siddhant Medar
committed by
Katon Minhas
Katon Minhas
parent
7056c1c687
commit
f7a9d80162
@@ -329,3 +329,11 @@ DOCZY_PDF_FILES_BUCKET_S3_URL = f"https://{DOCZY_PDF_FILES_BUCKET_NAME}.s3.us-ea
|
||||
|
||||
############## VISION API SETTINGS ##############
|
||||
ENABLE_VISION = get_arg_value("enable_vision", "False") == "True" # False by default
|
||||
|
||||
|
||||
############## DOCUMENT TYPE CLASSIFICATION SETTINGS ##############
|
||||
PERFORM_DTC = get_arg_value("perform_dtc", "False") == "True" # Enable document type classification pre-filter
|
||||
DTC_PROMPTS_JSON_PATH = "src/prompts/document_classification_prompts.json"
|
||||
DTC_MAX_PAGES_TO_CHECK = 5 # Number of pages to check per document for contract detection
|
||||
DTC_OUTPUT_FILE = f"{BATCH_ID}_Document_Classification_Report.csv" # Output CSV file path
|
||||
DTC_JSON_OUTPUT_FOLDER = "dtc_json_results" # Folder for individual JSON results
|
||||
@@ -1,88 +0,0 @@
|
||||
|
||||
|
||||
def apply_rules_by_state(row, state):
|
||||
"""
|
||||
Check if final verdict should be set as false based on state-specific rules.
|
||||
Returns True if document should be excluded (final_verdict = False).
|
||||
"""
|
||||
if state is None:
|
||||
return False
|
||||
|
||||
filename = row['Contract Name'].upper() # Convert to uppercase for case-insensitive matching
|
||||
pages = row['total_pages']
|
||||
|
||||
# Define state-specific exclusion keywords
|
||||
state_rules = {
|
||||
'WI': {
|
||||
'max_pages': 75,
|
||||
'keywords': [
|
||||
'W9', 'COM - LTR', 'FORM - FORM', 'FORM - LIC', 'FORM - OWN',
|
||||
'FORM - DOCS', 'FORM - APP', 'FORM - CRED', 'FORM - PPA',
|
||||
'FORM - WPA', 'CRF', 'APPLICATION', 'CREDENTIALING',
|
||||
'CRED DOCUMENTS', 'CRED DOCUMENT', ' PIF ', ' COLL ', ' ROS ', ' HDO '
|
||||
]
|
||||
},
|
||||
'MA': {
|
||||
'max_pages': None,
|
||||
'keywords': [
|
||||
'CRF', 'FORM - CRED', 'CREDENTIALING', 'COM - LTR',
|
||||
'FORM - APP', 'APPLICATION'
|
||||
]
|
||||
},
|
||||
'MS': {
|
||||
'max_pages': None,
|
||||
'keywords': [
|
||||
'FORM - OWN', ' PIF ', 'FORM - LIC', 'FORM - DOCS', 'FORM PPA', ' COLL '
|
||||
],
|
||||
'special_rules': {
|
||||
'W9': 6 # W9 with exactly 6 pages
|
||||
}
|
||||
},
|
||||
'SC': {
|
||||
'max_pages': None,
|
||||
'keywords': [
|
||||
'FORM - FORM', 'FORM - OWN', 'FORM - LIC', 'COM - LTR'
|
||||
],
|
||||
'special_rules': {
|
||||
'W9': 6 # W9 with exactly 6 pages
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state = state.upper()
|
||||
if state not in state_rules:
|
||||
return False
|
||||
|
||||
rules = state_rules[state]
|
||||
|
||||
# Check if state has a max page limit
|
||||
if rules['max_pages'] is not None:
|
||||
# For states with page limits (like WI)
|
||||
if pages > rules['max_pages']:
|
||||
return True # Exclude if pages exceed limit
|
||||
|
||||
# For documents within page limit, check for exclusion keywords
|
||||
for keyword in rules['keywords']:
|
||||
if keyword in filename:
|
||||
return True
|
||||
|
||||
# Check special rules
|
||||
if 'special_rules' in rules:
|
||||
for keyword, required_pages in rules['special_rules'].items():
|
||||
if keyword in filename and pages == required_pages:
|
||||
return True
|
||||
else:
|
||||
# For states without page limits (like MA, MS, SC)
|
||||
# Check keywords regardless of page count
|
||||
for keyword in rules['keywords']:
|
||||
if keyword in filename:
|
||||
return True
|
||||
|
||||
# Check special rules
|
||||
if 'special_rules' in rules:
|
||||
for keyword, required_pages in rules['special_rules'].items():
|
||||
if keyword in filename and pages == required_pages:
|
||||
return True
|
||||
|
||||
# If no exclusion rules matched, don't exclude
|
||||
return False
|
||||
@@ -1,86 +1,328 @@
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime
|
||||
from threading import Lock
|
||||
|
||||
import pandas as pd
|
||||
import re
|
||||
from src.investment.preprocessing_funcs import clean_newlines, split_text, clean_law_symbols
|
||||
from src.document_classification import doc_classification_funcs
|
||||
from src.utils import io_utils, string_utils
|
||||
import src.config as config
|
||||
|
||||
OUTPUT_FILE = f'Document-Classification-{config.BATCH_ID}.xlsx'
|
||||
from src.investment.preprocessing_funcs import (
|
||||
clean_newlines,
|
||||
split_text,
|
||||
clean_law_symbols,
|
||||
)
|
||||
from src.prompts.fieldset import Field
|
||||
from src.utils import io_utils, llm_utils
|
||||
from src import config
|
||||
|
||||
def main(state=None):
|
||||
# Load prompts from JSON using Field class (with thread-safe caching)
|
||||
# Load prompts at module level (cached for reuse across threads)
|
||||
DOCUMENT_TYPE_CLASSIFICATION_PROMPT = Field.load_from_file(
|
||||
config.DTC_PROMPTS_JSON_PATH, "document_type_classification"
|
||||
).prompt
|
||||
|
||||
yes_keywords = ['Agreement', 'Amendment', 'Amend', 'Agmt', 'Agree', 'Agrmnt', 'Standard Clause Letter', 'BH Letter', 'Appendix', 'AMD', 'Notice', 'Contract']
|
||||
no_keywords = ['void', 'do not use', 'Internal memo', 'Disclosure', 'Disc of Owner', 'W9', 'W-9', 'Credential', 'Coversheet', 'Cvrsht', 'Thumbs.db', 'Questionnaire',' Cert ', 'Certificate', 'email', 'Exclusion', 'Change Letter', 'Termination Letter', 'Form', 'Provider Letter', '.xls', '.lnk', 'Taxpayer', 'Regulatory Deeming', 'Administrative Amendment', 'Administrative Agreement', 'Instructions', 'Adobe Sign', 'Notice of Legal', 'Receipt ',
|
||||
"Dear", "to whom it may concern"]
|
||||
CONTRACT_TYPE_PROMPT = Field.load_from_file(
|
||||
config.DTC_PROMPTS_JSON_PATH, "contract_type"
|
||||
).prompt
|
||||
|
||||
yes_patterns = [(kw, re.compile(re.escape(kw), re.IGNORECASE)) for kw in yes_keywords]
|
||||
no_patterns = [(kw, re.compile(re.escape(kw), re.IGNORECASE)) for kw in no_keywords]
|
||||
NON_CONTRACT_TYPE_PROMPT = Field.load_from_file(
|
||||
config.DTC_PROMPTS_JSON_PATH, "non_contract_type"
|
||||
).prompt
|
||||
|
||||
record = []
|
||||
# Thread-safe counter for progress tracking
|
||||
progress_lock = Lock()
|
||||
progress_counter = {"completed": 0, "total": 0}
|
||||
|
||||
files = io_utils.read_input()
|
||||
|
||||
for filename, contract_text in files.items():
|
||||
|
||||
|
||||
def call_llm(filename, context, question):
|
||||
"""Call Claude through llm_utils.invoke_claude (unified LLM interface)
|
||||
|
||||
Args:
|
||||
filename: Name of the file being processed (for logging/tracking)
|
||||
context: Document text context
|
||||
question: The prompt/question to ask the LLM
|
||||
|
||||
Returns:
|
||||
str: The LLM's response, or "ERROR" if the call fails
|
||||
"""
|
||||
# Prepare the full prompt with context
|
||||
prompt = f"Document text:\n{context}\n\n{question}"
|
||||
|
||||
try:
|
||||
# Use the unified llm_utils interface (supports both local and EC2 modes)
|
||||
answer = llm_utils.invoke_claude(
|
||||
prompt=prompt,
|
||||
model_id="sonnet_latest", # Uses config.MODEL_ALIASES["sonnet_latest"]
|
||||
filename=filename,
|
||||
max_tokens=8192,
|
||||
)
|
||||
return answer.strip()
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error calling Claude for {filename}: {str(e)}")
|
||||
return "ERROR"
|
||||
|
||||
|
||||
def process_single_file(filename, contract_text, json_folder):
|
||||
"""Process a single file to determine if it's a contract and classify it.
|
||||
|
||||
Args:
|
||||
filename: Name of the file being processed
|
||||
contract_text: Raw text content of the file
|
||||
json_folder: Path to folder for saving individual JSON results
|
||||
|
||||
Returns:
|
||||
tuple: (filename, result_dict)
|
||||
"""
|
||||
try:
|
||||
contract_text = clean_newlines(contract_text)
|
||||
contract_text = clean_law_symbols(contract_text)
|
||||
text_dict = split_text(contract_text)
|
||||
|
||||
# --- search in filename ---
|
||||
yes_keywords_in_filename = [kw for kw, pattern in yes_patterns if pattern.search(filename)]
|
||||
no_keywords_in_filename = [kw for kw, pattern in no_patterns if pattern.search(filename)]
|
||||
|
||||
# --- first 5 pages ---
|
||||
threshold = min(5, len(text_dict))
|
||||
first_five_pages = "\n".join(text_dict[str(page)] for page in range(1, threshold + 1))
|
||||
|
||||
# --- search in body (first 5 pages only) ---
|
||||
yes_keywords_in_body = [kw for kw, pattern in yes_patterns if pattern.search(first_five_pages)]
|
||||
no_keywords_in_body = [kw for kw, pattern in no_patterns if pattern.search(first_five_pages)]
|
||||
|
||||
filename_verdict = len(yes_keywords_in_filename) > 0
|
||||
body_verdict = len(yes_keywords_in_body) > 0
|
||||
final_verdict = filename_verdict or body_verdict
|
||||
|
||||
record.append((
|
||||
filename,
|
||||
len(text_dict),
|
||||
yes_keywords_in_filename,
|
||||
no_keywords_in_filename,
|
||||
yes_keywords_in_body,
|
||||
no_keywords_in_body,
|
||||
filename_verdict,
|
||||
body_verdict,
|
||||
final_verdict,
|
||||
))
|
||||
|
||||
|
||||
columns = [
|
||||
'Contract Name',
|
||||
'total_pages',
|
||||
'yes_keywords_in_filename',
|
||||
'no_keywords_in_filename',
|
||||
'yes_keywords_in_body',
|
||||
'no_keywords_in_body',
|
||||
'filename_verdict',
|
||||
'body_verdict',
|
||||
'final_verdict'
|
||||
]
|
||||
# Loop through n pages
|
||||
is_contract = "NO"
|
||||
contract_type = "N/A"
|
||||
non_contract_type = "N/A"
|
||||
found_on_page = "N/A"
|
||||
|
||||
df = pd.DataFrame(record, columns=columns)
|
||||
for i in range(config.DTC_MAX_PAGES_TO_CHECK):
|
||||
page_key = str(i + 1)
|
||||
|
||||
# final verdict
|
||||
df['final_verdict'] = df['filename_verdict'] | df['body_verdict']
|
||||
df.loc[df['total_pages']<=3,'final_verdict'] = False
|
||||
# Handle edge case: document has fewer pages than config.DTC_MAX_PAGES_TO_CHECK
|
||||
if page_key not in text_dict:
|
||||
break
|
||||
|
||||
# Apply state-specific filtering if state is provided
|
||||
if not string_utils.is_empty(state):
|
||||
exclusion_mask = df.apply(lambda row: doc_classification_funcs.apply_rules_by_state(row, state), axis=1)
|
||||
df.loc[exclusion_mask, 'final_verdict'] = False
|
||||
print(f"Applied state-specific filtering for {state}!")
|
||||
context = text_dict[page_key]
|
||||
|
||||
df.to_excel(OUTPUT_FILE, index=False)
|
||||
# Ask if this page is a contract
|
||||
is_contract_answer = call_llm(
|
||||
filename,
|
||||
context,
|
||||
DOCUMENT_TYPE_CLASSIFICATION_PROMPT,
|
||||
)
|
||||
|
||||
if is_contract_answer.strip().upper() == "YES":
|
||||
# If contract found, get contract type
|
||||
type_of_contract_answer = call_llm(
|
||||
filename, context, CONTRACT_TYPE_PROMPT
|
||||
)
|
||||
is_contract = "YES"
|
||||
contract_type = type_of_contract_answer
|
||||
found_on_page = i + 1
|
||||
break # Stop checking further pages once contract is found
|
||||
|
||||
# If not a contract, determine non-contract document type
|
||||
if is_contract == "NO":
|
||||
# Use the first page for non-contract classification
|
||||
if "1" in text_dict:
|
||||
context = text_dict["1"]
|
||||
non_contract_type_answer = call_llm(
|
||||
filename, context, NON_CONTRACT_TYPE_PROMPT
|
||||
)
|
||||
non_contract_type = non_contract_type_answer
|
||||
found_on_page = 1
|
||||
|
||||
result = {
|
||||
"is_contract": is_contract,
|
||||
"contract_type": contract_type,
|
||||
"non_contract_type": non_contract_type,
|
||||
"found_on_page": found_on_page,
|
||||
}
|
||||
|
||||
# Save result to JSON immediately
|
||||
io_utils.save_result_to_json(filename, result, json_folder)
|
||||
|
||||
# Update progress counter
|
||||
with progress_lock:
|
||||
progress_counter["completed"] += 1
|
||||
completed = progress_counter["completed"]
|
||||
total = progress_counter["total"]
|
||||
if (
|
||||
completed % 5 == 0 or completed == total
|
||||
): # Log every 5 files or at completion
|
||||
logging.info(
|
||||
f"Progress: {completed}/{total} files processed ({completed/total*100:.1f}%)"
|
||||
)
|
||||
|
||||
return filename, result
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing {filename}: {str(e)}")
|
||||
error_result = {
|
||||
"is_contract": "ERROR",
|
||||
"contract_type": "ERROR",
|
||||
"non_contract_type": "ERROR",
|
||||
"found_on_page": "ERROR",
|
||||
}
|
||||
# Save error result to JSON
|
||||
io_utils.save_result_to_json(filename, error_result, json_folder)
|
||||
return filename, error_result
|
||||
|
||||
|
||||
def main(input_dict, run_timestamp):
|
||||
"""Run document type classification on a dictionary of files.
|
||||
|
||||
This function is designed to be called from investment/main.py as a pre-filter.
|
||||
It classifies all files and returns a dict mapping filenames to classification results.
|
||||
|
||||
Args:
|
||||
input_dict: Dictionary mapping filename -> file_text
|
||||
run_timestamp: Timestamp string for organizing output files (format: run_YYYYMMDD_HH-MM_BATCHID)
|
||||
|
||||
Returns:
|
||||
dict: Mapping of filename -> classification result dict with keys:
|
||||
- is_contract: "YES" or "NO"
|
||||
- contract_type: Type if contract, otherwise "N/A"
|
||||
- non_contract_type: Type if non-contract, otherwise "N/A"
|
||||
- found_on_page: Page number where classification was determined
|
||||
|
||||
Side effects:
|
||||
- Saves results to S3 (if config.WRITE_TO_S3) or local filesystem
|
||||
- Saves individual JSON files to config.DTC_JSON_OUTPUT_FOLDER
|
||||
"""
|
||||
logging.info(f"Starting DTC classification for {len(input_dict)} files...")
|
||||
logging.info(f"Max pages to check: {config.DTC_MAX_PAGES_TO_CHECK}")
|
||||
logging.info(f"Results will be saved to: {config.DTC_OUTPUT_FILE}")
|
||||
|
||||
# Initialize progress counter
|
||||
progress_counter["total"] = len(input_dict)
|
||||
progress_counter["completed"] = 0
|
||||
|
||||
# Create output folder
|
||||
os.makedirs(config.DTC_JSON_OUTPUT_FOLDER, exist_ok=True)
|
||||
logging.info(f"JSON output folder: {config.DTC_JSON_OUTPUT_FOLDER}")
|
||||
|
||||
# Record start time
|
||||
start_time = datetime.now()
|
||||
|
||||
answer_dict = {}
|
||||
|
||||
# Process files concurrently
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
|
||||
future_to_file = {}
|
||||
for filename, contract_text in input_dict.items():
|
||||
future = executor.submit(
|
||||
process_single_file,
|
||||
filename,
|
||||
contract_text,
|
||||
config.DTC_JSON_OUTPUT_FOLDER,
|
||||
)
|
||||
future_to_file[future] = filename
|
||||
|
||||
# Collect results as they complete
|
||||
for future in concurrent.futures.as_completed(future_to_file):
|
||||
filename = future_to_file[future]
|
||||
try:
|
||||
result_filename, result_data = future.result()
|
||||
answer_dict[result_filename] = result_data
|
||||
except Exception as e:
|
||||
logging.error(f"Exception occurred for {filename}: {str(e)}")
|
||||
answer_dict[filename] = {
|
||||
"is_contract": "ERROR",
|
||||
"contract_type": "ERROR",
|
||||
"non_contract_type": "ERROR",
|
||||
"found_on_page": "ERROR",
|
||||
}
|
||||
|
||||
# Record end time
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
logging.info(f"DTC classification complete in {duration:.2f} seconds")
|
||||
logging.info(f"Average time per file: {duration/len(input_dict):.2f} seconds")
|
||||
|
||||
# Create DataFrame and save results
|
||||
df = pd.DataFrame([{"FILE_NAME": k, **v} for k, v in answer_dict.items()])
|
||||
|
||||
# Log summary statistics
|
||||
contracts_found = len(df[df["is_contract"] == "YES"])
|
||||
non_contracts = len(df[df["is_contract"] == "NO"])
|
||||
errors = len(df[df["is_contract"] == "ERROR"])
|
||||
|
||||
logging.info(
|
||||
f"DTC Summary: {contracts_found} contracts, {non_contracts} non-contracts, {errors} errors"
|
||||
)
|
||||
|
||||
# Create summary DataFrame with aggregate statistics
|
||||
summary_data = {
|
||||
"total_files": [len(df)],
|
||||
"contracts": [contracts_found],
|
||||
"non_contracts": [non_contracts],
|
||||
"errors": [errors],
|
||||
"contract_rate": [
|
||||
f"{(contracts_found/len(df)*100):.1f}%" if len(df) > 0 else "0%"
|
||||
],
|
||||
"run_timestamp": [run_timestamp],
|
||||
"batch_id": [config.BATCH_ID],
|
||||
}
|
||||
|
||||
# Add contract type breakdown
|
||||
if contracts_found > 0:
|
||||
contract_types = df[df["is_contract"] == "YES"]["contract_type"].value_counts()
|
||||
summary_data["most_common_contract_type"] = [
|
||||
contract_types.index[0] if len(contract_types) > 0 else "N/A"
|
||||
]
|
||||
summary_data["contract_type_breakdown"] = [contract_types.to_dict()]
|
||||
else:
|
||||
summary_data["most_common_contract_type"] = ["N/A"]
|
||||
summary_data["contract_type_breakdown"] = [{}]
|
||||
|
||||
# Add non-contract type breakdown
|
||||
if non_contracts > 0:
|
||||
non_contract_types = df[df["is_contract"] == "NO"][
|
||||
"non_contract_type"
|
||||
].value_counts()
|
||||
summary_data["most_common_non_contract_type"] = [
|
||||
non_contract_types.index[0] if len(non_contract_types) > 0 else "N/A"
|
||||
]
|
||||
summary_data["non_contract_type_breakdown"] = [non_contract_types.to_dict()]
|
||||
else:
|
||||
summary_data["most_common_non_contract_type"] = ["N/A"]
|
||||
summary_data["non_contract_type_breakdown"] = [{}]
|
||||
|
||||
summary_df = pd.DataFrame(summary_data)
|
||||
|
||||
# Save detailed results using same pattern as investment pipeline
|
||||
if config.WRITE_TO_S3:
|
||||
io_utils.write_s3(df, "", run_timestamp, "dtc")
|
||||
logging.info(
|
||||
f"DTC results uploaded to S3: {config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC.csv"
|
||||
)
|
||||
|
||||
# Save summary
|
||||
io_utils.write_s3(summary_df, "", run_timestamp, "dtc_summary")
|
||||
logging.info(
|
||||
f"DTC summary uploaded to S3: {config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC-SUMMARY.csv"
|
||||
)
|
||||
else:
|
||||
io_utils.write_local(df, "", run_timestamp, "dtc")
|
||||
logging.info(
|
||||
f"DTC results saved locally: {config.CONSOLIDATED_OUTPUT_DIRECTORY}/{run_timestamp}/{config.BATCH_ID}-DTC.csv"
|
||||
)
|
||||
|
||||
# Save summary
|
||||
io_utils.write_local(summary_df, "", run_timestamp, "dtc_summary")
|
||||
logging.info(
|
||||
f"DTC summary saved locally: {config.CONSOLIDATED_OUTPUT_DIRECTORY}/{run_timestamp}/{config.BATCH_ID}-DTC-SUMMARY.csv"
|
||||
)
|
||||
|
||||
return answer_dict
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(state=config.STATE)
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
|
||||
# Generate timestamp for this run
|
||||
run_timestamp = f"run_{datetime.now().strftime('%Y%m%d_%H-%M')}_{config.BATCH_ID}"
|
||||
|
||||
# Load input files
|
||||
logging.info("Loading input files...")
|
||||
input_dict = io_utils.read_input()
|
||||
|
||||
# Run classification
|
||||
results = main(input_dict, run_timestamp)
|
||||
|
||||
logging.info(f"Classification complete. Processed {len(results)} files.")
|
||||
|
||||
@@ -17,7 +17,7 @@ import src.utils.logging_utils as logging_utils
|
||||
import src.utils.usage_tracking as usage_tracking
|
||||
from constants.constants import Constants
|
||||
from src import config
|
||||
|
||||
from src.document_classification import main as dtc_main
|
||||
|
||||
def safe_process_file(item, constants, run_timestamp):
|
||||
"""Process a single file with comprehensive error handling.
|
||||
@@ -99,6 +99,19 @@ def main(testing=False, test_params={}):
|
||||
|
||||
logging.info(f"Total Input Files : {len(input_dict)}")
|
||||
|
||||
# Run document type classification if enabled - DTC mode only, then exit
|
||||
if config.PERFORM_DTC:
|
||||
logging.info("=" * 80)
|
||||
logging.info("DOCUMENT TYPE CLASSIFICATION MODE - Running DTC and exiting")
|
||||
logging.info("=" * 80)
|
||||
|
||||
# run DTC classification
|
||||
dtc_main.main(input_dict, run_timestamp)
|
||||
|
||||
logging.info("DTC classification complete. Exiting (investment pipeline not run).")
|
||||
logging.info("=" * 80)
|
||||
return
|
||||
|
||||
# Warm prompt caches before parallel processing
|
||||
# This ensures the cache is registered before multiple threads try to use it
|
||||
logging.info("Warming prompt caches before parallel processing...")
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
[
|
||||
{
|
||||
"field_name": "document_type_classification",
|
||||
"relationship": "one_to_one",
|
||||
"field_type": "classification",
|
||||
"prompt": "Is this document a contract between two parties? Examples include but are not limited to Service Agreements, NDAs, Statements of Work, Amendments, or Addenda.\n Focus only on the written text, not tables or structured data.\n Ignore: tables, forms, schedules, rate sheets, or data fields.\n Answer YES if the text includes:\n Titles like \"Agreement,\" \"Contract,\" \"Addendum,\" etc.\n Phrases such as \"parties agree,\" \"executes this Agreement,\" \"authorized representative\"\n References to \"this Agreement\" or \"the Agreement\"\n Legal or contractual language\n Respond with only: YES or NO",
|
||||
"valid_values": ["YES", "NO"]
|
||||
},
|
||||
{
|
||||
"field_name": "contract_type",
|
||||
"relationship": "one_to_one",
|
||||
"field_type": "classification",
|
||||
"prompt": "Identify the contract type from the text above.\n Classification rules:\n - \"Service Agreement\" includes: Provider Service Agreement, Provider Services Agreement, Professional Services Agreement, Vendor Agreement\n - \"Master Service Agreement\" includes: MSA, Master Agreement, Framework Agreement\n - \"Amendment\" includes: Amendment, Modification, Contract Amendment\n - \"Addendum\" includes: Addendum, Appendix, Exhibit, Attachment (if binding)\n - \"Statement of Work\" includes: SOW, Work Order, Service Order\n - \"Non-Disclosure Agreement\" includes: NDA, Confidentiality Agreement\n - \"Other\" for: any contract not matching above categories\n Output ONLY one of these exact options (no explanation):\n Service Agreement\n Master Service Agreement\n Non-Disclosure Agreement\n Statement of Work\n Amendment\n Addendum\n Other",
|
||||
"valid_values": ["Service Agreement", "Master Service Agreement", "Non-Disclosure Agreement", "Statement of Work", "Amendment", "Addendum", "Other"]
|
||||
},
|
||||
{
|
||||
"field_name": "non_contract_type",
|
||||
"relationship": "one_to_one",
|
||||
"field_type": "classification",
|
||||
"prompt": "Identify the document type from the document above.\n Classification rules:\n - \"Email\" includes: emails, letters, correspondence, faxes, notices\n - \"Memo\" includes: memos, memorandums, internal communications\n - \"Form\" includes: forms, templates, applications, questionnaires, blank documents\n - \"Roster\" includes: rosters, lists, directories, schedules, provider lists\n - \"Other\" for: invoices, bills, reports, policies, claims, or any document not matching above\n Output ONLY one of these exact options (no explanation):\n Email\n Memo\n Form\n Roster\n Other",
|
||||
"valid_values": ["Email", "Memo", "Form", "Roster", "Other"]
|
||||
}
|
||||
]
|
||||
@@ -1,7 +1,9 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from io import StringIO
|
||||
from threading import Lock
|
||||
|
||||
import pandas as pd
|
||||
import src.tracking.tracking_utils as tracking_utils
|
||||
@@ -9,6 +11,9 @@ from botocore.exceptions import ClientError
|
||||
from pyxlsb import open_workbook
|
||||
from src import config
|
||||
|
||||
# Thread-safe lock for JSON file operations
|
||||
json_write_lock = Lock()
|
||||
|
||||
|
||||
def read_local(file_path) -> str | pd.DataFrame | None:
|
||||
"""Reads a local (.txt, .csv, or .xlsx) file and returns its contents.
|
||||
@@ -398,9 +403,27 @@ def write_local(
|
||||
quoting=1,
|
||||
)
|
||||
return None
|
||||
elif output_type == "dtc":
|
||||
output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
df.to_csv(
|
||||
os.path.join(output_dir, f"{config.BATCH_ID}-DTC.csv"),
|
||||
index=False,
|
||||
quoting=1,
|
||||
)
|
||||
return None
|
||||
elif output_type == "dtc_summary":
|
||||
output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
df.to_csv(
|
||||
os.path.join(output_dir, f"{config.BATCH_ID}-DTC-SUMMARY.csv"),
|
||||
index=False,
|
||||
quoting=1,
|
||||
)
|
||||
return None
|
||||
else:
|
||||
logging.error(
|
||||
f"Unknown output_type: {output_type}; output_type must be 'final', 'individual', 'error', 'usage', or 'usage_summary'"
|
||||
f"Unknown output_type: {output_type}. Valid types are: 'final', 'individual', 'error', 'usage', 'usage_summary', 'dtc', 'dtc_summary'"
|
||||
)
|
||||
return None
|
||||
|
||||
@@ -445,6 +468,10 @@ def write_s3(
|
||||
output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE.csv"
|
||||
elif output_type == "usage_summary":
|
||||
output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE-SUMMARY.csv"
|
||||
elif output_type == "dtc":
|
||||
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC.csv"
|
||||
elif output_type == "dtc_summary":
|
||||
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-DTC-SUMMARY.csv"
|
||||
else:
|
||||
# Fallback path for unknown output types
|
||||
output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-unknown-{output_type}.csv"
|
||||
@@ -464,9 +491,60 @@ def write_s3(
|
||||
logging.info(f"Saved usage tracking file: {output_path}")
|
||||
elif output_type == "usage_summary":
|
||||
logging.info(f"Saved usage summary file: {output_path}")
|
||||
elif output_type == "dtc":
|
||||
logging.info(f"Saved DTC output file: {output_path}")
|
||||
elif output_type == "dtc_summary":
|
||||
logging.info(f"Saved DTC summary file: {output_path}")
|
||||
else:
|
||||
logging.info(f"Saved unknown output type file: {output_path}")
|
||||
return None
|
||||
except ClientError as e:
|
||||
logging.error(f"Error uploading to S3: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def save_result_to_json(filename, result, json_folder):
|
||||
"""Save a result dictionary to an individual JSON file in the specified folder.
|
||||
|
||||
This function is thread-safe and can be called concurrently from multiple threads.
|
||||
It's useful for saving intermediate results or individual processing outputs during
|
||||
batch operations.
|
||||
|
||||
Args:
|
||||
filename: Name of the source file that was processed (used to generate JSON filename)
|
||||
result: Dictionary containing the processing results to be saved
|
||||
json_folder: Path to the folder where the JSON file should be saved
|
||||
|
||||
Returns:
|
||||
None; this function is called for its side effect of writing a JSON file to disk.
|
||||
|
||||
Notes:
|
||||
- Creates the output folder if it doesn't exist
|
||||
- Automatically sanitizes the filename to remove invalid characters (/, \\, :)
|
||||
- Adds .json extension if not already present
|
||||
- Thread-safe: uses a lock to prevent concurrent write conflicts
|
||||
- Saves both the original filename and result in the JSON structure
|
||||
|
||||
Example:
|
||||
>>> save_result_to_json(
|
||||
... "contract.pdf",
|
||||
... {"status": "processed", "result": "approved"},
|
||||
... "/path/to/output"
|
||||
... )
|
||||
# Creates: /path/to/output/contract.pdf.json with content:
|
||||
# {"filename": "contract.pdf", "result": {"status": "processed", "result": "approved"}}
|
||||
"""
|
||||
with json_write_lock:
|
||||
# Ensure the folder exists
|
||||
os.makedirs(json_folder, exist_ok=True)
|
||||
|
||||
# Create a safe filename for JSON (replace invalid characters)
|
||||
safe_filename = filename.replace("/", "_").replace("\\", "_").replace(":", "_")
|
||||
if not safe_filename.endswith(".json"):
|
||||
safe_filename = f"{safe_filename}.json"
|
||||
|
||||
json_file_path = os.path.join(json_folder, safe_filename)
|
||||
|
||||
# Save result to individual JSON file
|
||||
with open(json_file_path, "w") as f:
|
||||
json.dump({"filename": filename, "result": result}, f, indent=2)
|
||||
|
||||
@@ -0,0 +1,534 @@
|
||||
"""
|
||||
Unit tests for document_classification/main.py module.
|
||||
|
||||
Tests cover:
|
||||
- call_llm: LLM invocation wrapper
|
||||
- save_result_to_json: JSON file saving
|
||||
- process_single_file: Single document classification
|
||||
- main: Full batch classification with S3/local upload
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from datetime import datetime
|
||||
from unittest.mock import MagicMock, Mock, patch, call
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from src.document_classification import main as dtc_main
|
||||
from src.utils import io_utils
|
||||
|
||||
|
||||
class TestCallLLM(unittest.TestCase):
|
||||
"""Test the call_llm function that wraps llm_utils.invoke_claude."""
|
||||
|
||||
@patch('src.document_classification.main.llm_utils.invoke_claude')
|
||||
def test_call_llm_success(self, mock_invoke):
|
||||
"""Test successful LLM call returns cleaned response."""
|
||||
mock_invoke.return_value = " YES \n"
|
||||
|
||||
result = dtc_main.call_llm(
|
||||
filename="test.txt",
|
||||
context="This is a Service Agreement...",
|
||||
question="Is this a contract?"
|
||||
)
|
||||
|
||||
self.assertEqual(result, "YES")
|
||||
mock_invoke.assert_called_once()
|
||||
call_args = mock_invoke.call_args
|
||||
self.assertIn("This is a Service Agreement", call_args.kwargs['prompt'])
|
||||
self.assertEqual(call_args.kwargs['model_id'], "sonnet_latest")
|
||||
self.assertEqual(call_args.kwargs['filename'], "test.txt")
|
||||
self.assertEqual(call_args.kwargs['max_tokens'], 8192)
|
||||
|
||||
@patch('src.document_classification.main.llm_utils.invoke_claude')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_call_llm_error_handling(self, mock_logging, mock_invoke):
|
||||
"""Test error handling returns ERROR and logs the exception."""
|
||||
mock_invoke.side_effect = Exception("API timeout")
|
||||
|
||||
result = dtc_main.call_llm(
|
||||
filename="test.txt",
|
||||
context="Document text",
|
||||
question="Question"
|
||||
)
|
||||
|
||||
self.assertEqual(result, "ERROR")
|
||||
mock_logging.error.assert_called_once()
|
||||
self.assertIn("test.txt", mock_logging.error.call_args[0][0])
|
||||
|
||||
@patch('src.document_classification.main.llm_utils.invoke_claude')
|
||||
def test_call_llm_prompt_format(self, mock_invoke):
|
||||
"""Test that prompt is correctly formatted with context and question."""
|
||||
mock_invoke.return_value = "Answer"
|
||||
context = "Contract document text here"
|
||||
question = "What type is this?"
|
||||
|
||||
dtc_main.call_llm("test.txt", context, question)
|
||||
|
||||
prompt = mock_invoke.call_args.kwargs['prompt']
|
||||
self.assertIn("Document text:", prompt)
|
||||
self.assertIn(context, prompt)
|
||||
self.assertIn(question, prompt)
|
||||
|
||||
|
||||
class TestSaveResultToJson(unittest.TestCase):
|
||||
"""Test the io_utils.save_result_to_json function for JSON file operations."""
|
||||
|
||||
def setUp(self):
|
||||
"""Create a temporary directory for test files."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directory."""
|
||||
import shutil
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
def test_save_result_creates_folder(self):
|
||||
"""Test that function creates output folder if it doesn't exist."""
|
||||
json_folder = os.path.join(self.temp_dir, "new_folder")
|
||||
self.assertFalse(os.path.exists(json_folder))
|
||||
|
||||
io_utils.save_result_to_json(
|
||||
filename="test.txt",
|
||||
result={"is_contract": "YES"},
|
||||
json_folder=json_folder
|
||||
)
|
||||
|
||||
self.assertTrue(os.path.exists(json_folder))
|
||||
|
||||
def test_save_result_creates_json_file(self):
|
||||
"""Test that JSON file is created with correct content."""
|
||||
result = {
|
||||
"is_contract": "YES",
|
||||
"contract_type": "Service Agreement",
|
||||
"found_on_page": 1
|
||||
}
|
||||
|
||||
io_utils.save_result_to_json(
|
||||
filename="contract123.txt",
|
||||
result=result,
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
json_file = os.path.join(self.temp_dir, "contract123.txt.json")
|
||||
self.assertTrue(os.path.exists(json_file))
|
||||
|
||||
with open(json_file, 'r') as f:
|
||||
saved_data = json.load(f)
|
||||
|
||||
self.assertEqual(saved_data['filename'], "contract123.txt")
|
||||
self.assertEqual(saved_data['result'], result)
|
||||
|
||||
def test_save_result_sanitizes_filename(self):
|
||||
"""Test that invalid filename characters are replaced."""
|
||||
result = {"is_contract": "NO"}
|
||||
|
||||
io_utils.save_result_to_json(
|
||||
filename="path/to/file:with:colons.txt",
|
||||
result=result,
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
# Should replace / : \ with _
|
||||
expected_filename = "path_to_file_with_colons.txt.json"
|
||||
json_file = os.path.join(self.temp_dir, expected_filename)
|
||||
self.assertTrue(os.path.exists(json_file))
|
||||
|
||||
def test_save_result_adds_json_extension(self):
|
||||
"""Test that .json extension is added if not present."""
|
||||
io_utils.save_result_to_json(
|
||||
filename="test.txt",
|
||||
result={"is_contract": "YES"},
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
json_file = os.path.join(self.temp_dir, "test.txt.json")
|
||||
self.assertTrue(os.path.exists(json_file))
|
||||
|
||||
|
||||
class TestProcessSingleFile(unittest.TestCase):
|
||||
"""Test the process_single_file function for document classification."""
|
||||
|
||||
def setUp(self):
|
||||
"""Create temporary directory for JSON output."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directory."""
|
||||
import shutil
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
def test_process_contract_found_first_page(self, mock_clean_nl, mock_clean_law,
|
||||
mock_split, mock_llm):
|
||||
"""Test processing when contract is found on first page."""
|
||||
mock_clean_nl.return_value = "cleaned text"
|
||||
mock_clean_law.return_value = "cleaned text"
|
||||
mock_split.return_value = {
|
||||
"1": "This is a Service Agreement between parties...",
|
||||
"2": "Page 2 content"
|
||||
}
|
||||
|
||||
# First call: is_contract check returns YES
|
||||
# Second call: contract_type returns Service Agreement
|
||||
mock_llm.side_effect = ["YES", "Service Agreement"]
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="contract.txt",
|
||||
contract_text="raw contract text",
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
self.assertEqual(filename, "contract.txt")
|
||||
self.assertEqual(result['is_contract'], "YES")
|
||||
self.assertEqual(result['contract_type'], "Service Agreement")
|
||||
self.assertEqual(result['non_contract_type'], "N/A")
|
||||
self.assertEqual(result['found_on_page'], 1)
|
||||
|
||||
# Should make 2 LLM calls (is_contract + contract_type)
|
||||
self.assertEqual(mock_llm.call_count, 2)
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
def test_process_non_contract(self, mock_clean_nl, mock_clean_law,
|
||||
mock_split, mock_llm):
|
||||
"""Test processing when document is not a contract."""
|
||||
mock_clean_nl.return_value = "cleaned text"
|
||||
mock_clean_law.return_value = "cleaned text"
|
||||
mock_split.return_value = {
|
||||
"1": "This is an email from John...",
|
||||
"2": "Email continues..."
|
||||
}
|
||||
|
||||
# First call: is_contract returns NO
|
||||
# Second call: non_contract_type returns Email
|
||||
mock_llm.side_effect = ["NO", "NO", "Email"]
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="email.txt",
|
||||
contract_text="email text",
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
self.assertEqual(result['is_contract'], "NO")
|
||||
self.assertEqual(result['contract_type'], "N/A")
|
||||
self.assertEqual(result['non_contract_type'], "Email")
|
||||
self.assertEqual(result['found_on_page'], 1)
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
@patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 3)
|
||||
def test_process_contract_found_later_page(self, mock_clean_nl, mock_clean_law,
|
||||
mock_split, mock_llm):
|
||||
"""Test when contract signature is found on page 3."""
|
||||
mock_clean_nl.return_value = "cleaned"
|
||||
mock_clean_law.return_value = "cleaned"
|
||||
mock_split.return_value = {
|
||||
"1": "Cover page",
|
||||
"2": "Table of contents",
|
||||
"3": "SERVICE AGREEMENT between parties",
|
||||
"4": "More content"
|
||||
}
|
||||
|
||||
# Page 1: NO, Page 2: NO, Page 3: YES, then get type
|
||||
mock_llm.side_effect = ["NO", "NO", "YES", "Service Agreement"]
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="contract.txt",
|
||||
contract_text="text",
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
self.assertEqual(result['is_contract'], "YES")
|
||||
self.assertEqual(result['found_on_page'], 3)
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_process_error_handling(self, mock_logging, mock_clean_nl,
|
||||
mock_clean_law, mock_split, mock_llm):
|
||||
"""Test error handling returns error result."""
|
||||
mock_clean_nl.side_effect = Exception("Processing error")
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="bad.txt",
|
||||
contract_text="text",
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
self.assertEqual(result['is_contract'], "ERROR")
|
||||
self.assertEqual(result['contract_type'], "ERROR")
|
||||
self.assertEqual(result['non_contract_type'], "ERROR")
|
||||
self.assertEqual(result['found_on_page'], "ERROR")
|
||||
mock_logging.error.assert_called_once()
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
@patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 2)
|
||||
def test_process_handles_fewer_pages(self, mock_clean_nl, mock_clean_law,
|
||||
mock_split, mock_llm):
|
||||
"""Test when document has fewer pages than MAX_PAGES_TO_CHECK."""
|
||||
mock_clean_nl.return_value = "cleaned"
|
||||
mock_clean_law.return_value = "cleaned"
|
||||
mock_split.return_value = {"1": "Only one page"} # Only 1 page
|
||||
|
||||
mock_llm.side_effect = ["NO", "Memo"]
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="short.txt",
|
||||
contract_text="text",
|
||||
json_folder=self.temp_dir
|
||||
)
|
||||
|
||||
# Should only check page 1, not error on page 2
|
||||
self.assertEqual(result['is_contract'], "NO")
|
||||
|
||||
|
||||
class TestRunDTCClassification(unittest.TestCase):
|
||||
"""Integration tests for main function."""
|
||||
|
||||
def setUp(self):
|
||||
"""Create temporary directory for outputs."""
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up temporary directory."""
|
||||
import shutil
|
||||
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
||||
|
||||
@patch('src.document_classification.main.io_utils.write_s3')
|
||||
@patch('src.document_classification.main.io_utils.write_local')
|
||||
@patch('src.document_classification.main.process_single_file')
|
||||
@patch('src.document_classification.main.config.WRITE_TO_S3', False)
|
||||
@patch('src.document_classification.main.config.BATCH_ID', 'test_batch')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_main_local(self, mock_logging,
|
||||
mock_process, mock_write_local,
|
||||
mock_write_s3):
|
||||
"""Test full DTC run with local file output."""
|
||||
|
||||
# Mock process_single_file to return results
|
||||
mock_process.side_effect = [
|
||||
("file1.txt", {"is_contract": "YES", "contract_type": "Service Agreement",
|
||||
"non_contract_type": "N/A", "found_on_page": 1}),
|
||||
("file2.txt", {"is_contract": "NO", "contract_type": "N/A",
|
||||
"non_contract_type": "Email", "found_on_page": 1})
|
||||
]
|
||||
|
||||
input_dict = {
|
||||
"file1.txt": "Service Agreement text...",
|
||||
"file2.txt": "Email text..."
|
||||
}
|
||||
run_timestamp = "run_20250106_10-30_test"
|
||||
|
||||
results = dtc_main.main(input_dict, run_timestamp)
|
||||
|
||||
# Check results returned correctly
|
||||
self.assertIn("file1.txt", results)
|
||||
self.assertIn("file2.txt", results)
|
||||
self.assertEqual(results["file1.txt"]["is_contract"], "YES")
|
||||
self.assertEqual(results["file2.txt"]["is_contract"], "NO")
|
||||
|
||||
# Check write_local was called twice (once for details, once for summary), not write_s3
|
||||
self.assertEqual(mock_write_local.call_count, 2)
|
||||
mock_write_s3.assert_not_called()
|
||||
|
||||
# Verify write_local called with correct args for main results
|
||||
first_call_args = mock_write_local.call_args_list[0]
|
||||
df_arg = first_call_args[0][0]
|
||||
self.assertIsInstance(df_arg, pd.DataFrame)
|
||||
self.assertEqual(len(df_arg), 2)
|
||||
self.assertEqual(first_call_args[0][2], run_timestamp)
|
||||
self.assertEqual(first_call_args[0][3], "dtc")
|
||||
|
||||
# Verify summary was also saved
|
||||
second_call_args = mock_write_local.call_args_list[1]
|
||||
summary_df = second_call_args[0][0]
|
||||
self.assertIsInstance(summary_df, pd.DataFrame)
|
||||
self.assertEqual(second_call_args[0][3], "dtc_summary")
|
||||
|
||||
@patch('src.document_classification.main.io_utils.write_s3')
|
||||
@patch('src.document_classification.main.io_utils.write_local')
|
||||
@patch('src.document_classification.main.process_single_file')
|
||||
@patch('src.document_classification.main.config.WRITE_TO_S3', True)
|
||||
@patch('src.document_classification.main.config.BATCH_ID', 'test_batch')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_main_s3(self, mock_logging,
|
||||
mock_process, mock_write_local,
|
||||
mock_write_s3):
|
||||
"""Test full DTC run with S3 upload."""
|
||||
|
||||
mock_process.return_value = (
|
||||
"file1.txt",
|
||||
{"is_contract": "YES", "contract_type": "MSA",
|
||||
"non_contract_type": "N/A", "found_on_page": 2}
|
||||
)
|
||||
|
||||
input_dict = {"file1.txt": "Master Service Agreement..."}
|
||||
run_timestamp = "run_20250106_10-30_test"
|
||||
|
||||
results = dtc_main.main(input_dict, run_timestamp)
|
||||
|
||||
# Check S3 upload was called twice (main + summary), not local
|
||||
self.assertEqual(mock_write_s3.call_count, 2)
|
||||
mock_write_local.assert_not_called()
|
||||
|
||||
# Verify write_s3 called with correct args for main results
|
||||
first_call = mock_write_s3.call_args_list[0]
|
||||
self.assertEqual(first_call[0][2], run_timestamp)
|
||||
self.assertEqual(first_call[0][3], "dtc")
|
||||
|
||||
# Verify summary was also saved to S3
|
||||
second_call = mock_write_s3.call_args_list[1]
|
||||
self.assertEqual(second_call[0][2], run_timestamp)
|
||||
self.assertEqual(second_call[0][3], "dtc_summary")
|
||||
|
||||
@patch('src.document_classification.main.io_utils.write_local')
|
||||
@patch('src.document_classification.main.process_single_file')
|
||||
@patch('src.document_classification.main.config.WRITE_TO_S3', False)
|
||||
@patch('src.document_classification.main.config.BATCH_ID', 'test_batch')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_run_dtc_summary_statistics(self, mock_logging,
|
||||
mock_process, mock_write_local):
|
||||
"""Test that summary statistics are logged correctly."""
|
||||
|
||||
# 2 contracts, 1 non-contract, 1 error
|
||||
mock_process.side_effect = [
|
||||
("f1.txt", {"is_contract": "YES", "contract_type": "Service Agreement",
|
||||
"non_contract_type": "N/A", "found_on_page": 1}),
|
||||
("f2.txt", {"is_contract": "YES", "contract_type": "MSA",
|
||||
"non_contract_type": "N/A", "found_on_page": 1}),
|
||||
("f3.txt", {"is_contract": "NO", "contract_type": "N/A",
|
||||
"non_contract_type": "Email", "found_on_page": 1}),
|
||||
("f4.txt", {"is_contract": "ERROR", "contract_type": "ERROR",
|
||||
"non_contract_type": "ERROR", "found_on_page": "ERROR"})
|
||||
]
|
||||
|
||||
input_dict = {f"f{i}.txt": "text" for i in range(1, 5)}
|
||||
|
||||
dtc_main.main(input_dict, "run_test")
|
||||
|
||||
# Check that summary was logged
|
||||
log_calls = [str(call) for call in mock_logging.info.call_args_list]
|
||||
summary_log = [c for c in log_calls if "DTC Summary" in c]
|
||||
self.assertTrue(len(summary_log) > 0)
|
||||
|
||||
# Verify counts in summary
|
||||
summary_str = str(summary_log[0])
|
||||
self.assertIn("2 contracts", summary_str)
|
||||
self.assertIn("1 non-contracts", summary_str)
|
||||
self.assertIn("1 errors", summary_str)
|
||||
|
||||
|
||||
@patch('src.document_classification.main.io_utils.write_local')
|
||||
@patch('src.document_classification.main.process_single_file')
|
||||
@patch('src.document_classification.main.config.WRITE_TO_S3', False)
|
||||
@patch('src.document_classification.main.config.BATCH_ID', 'test_batch')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_run_dtc_no_contracts_summary(self, mock_logging,
|
||||
mock_process, mock_write_local):
|
||||
"""Test summary when there are no contracts (covers else block)."""
|
||||
# All non-contracts
|
||||
mock_process.side_effect = [
|
||||
("f1.txt", {"is_contract": "NO", "contract_type": "N/A",
|
||||
"non_contract_type": "Email", "found_on_page": 1}),
|
||||
("f2.txt", {"is_contract": "NO", "contract_type": "N/A",
|
||||
"non_contract_type": "Form", "found_on_page": 1}),
|
||||
]
|
||||
|
||||
input_dict = {f"f{i}.txt": "text" for i in range(1, 3)}
|
||||
dtc_main.main(input_dict, "run_test")
|
||||
|
||||
# Verify write_local was called twice
|
||||
self.assertEqual(mock_write_local.call_count, 2)
|
||||
|
||||
# Check summary DataFrame
|
||||
summary_df = mock_write_local.call_args_list[1][0][0]
|
||||
self.assertEqual(summary_df['contracts'].iloc[0], 0)
|
||||
self.assertEqual(summary_df['non_contracts'].iloc[0], 2)
|
||||
self.assertEqual(summary_df['most_common_contract_type'].iloc[0], "N/A")
|
||||
|
||||
@patch('src.document_classification.main.io_utils.write_local')
|
||||
@patch('src.document_classification.main.process_single_file')
|
||||
@patch('src.document_classification.main.config.WRITE_TO_S3', False)
|
||||
@patch('src.document_classification.main.config.BATCH_ID', 'test_batch')
|
||||
@patch('src.document_classification.main.logging')
|
||||
def test_run_dtc_future_exception(self, mock_logging,
|
||||
mock_process, mock_write_local):
|
||||
"""Test exception handling in futures loop (covers lines 244-246)."""
|
||||
# Simulate an exception being raised by a future
|
||||
def raise_exception(filename, text, folder):
|
||||
if filename == "bad.txt":
|
||||
raise ValueError("Simulated processing error")
|
||||
return (filename, {"is_contract": "YES", "contract_type": "MSA",
|
||||
"non_contract_type": "N/A", "found_on_page": 1})
|
||||
|
||||
mock_process.side_effect = raise_exception
|
||||
|
||||
input_dict = {
|
||||
"good.txt": "text",
|
||||
"bad.txt": "text"
|
||||
}
|
||||
|
||||
results = dtc_main.main(input_dict, "run_test")
|
||||
|
||||
# Should handle exception and still return results
|
||||
self.assertIn("bad.txt", results)
|
||||
self.assertEqual(results["bad.txt"]["is_contract"], "ERROR")
|
||||
|
||||
# Verify error was logged
|
||||
error_logs = [call for call in mock_logging.error.call_args_list
|
||||
if "Exception occurred" in str(call)]
|
||||
self.assertTrue(len(error_logs) > 0)
|
||||
|
||||
@patch('src.document_classification.main.call_llm')
|
||||
@patch('src.document_classification.main.split_text')
|
||||
@patch('src.document_classification.main.clean_law_symbols')
|
||||
@patch('src.document_classification.main.clean_newlines')
|
||||
@patch('src.document_classification.main.config.DTC_MAX_PAGES_TO_CHECK', 10)
|
||||
def test_process_progress_logging_every_5(self, mock_clean_nl, mock_clean_law,
|
||||
mock_split, mock_llm):
|
||||
"""Test that progress is logged every 5 files (covers line 169)."""
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
try:
|
||||
mock_clean_nl.return_value = "cleaned"
|
||||
mock_clean_law.return_value = "cleaned"
|
||||
mock_split.return_value = {"1": "Contract text"}
|
||||
mock_llm.side_effect = ["YES", "Service Agreement"]
|
||||
|
||||
# Set up progress counter to trigger line 169
|
||||
from src.document_classification.main import progress_counter, progress_lock
|
||||
with progress_lock:
|
||||
progress_counter["total"] = 10
|
||||
progress_counter["completed"] = 4 # Next will be 5
|
||||
|
||||
filename, result = dtc_main.process_single_file(
|
||||
filename="file5.txt",
|
||||
contract_text="text",
|
||||
json_folder=temp_dir
|
||||
)
|
||||
|
||||
# Verify file was processed
|
||||
self.assertEqual(filename, "file5.txt")
|
||||
self.assertEqual(result["is_contract"], "YES")
|
||||
finally:
|
||||
import shutil
|
||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user