diff --git a/fieldExtraction/src/investment/main.py b/fieldExtraction/src/investment/main.py index 4e37ed1..85c6c5a 100644 --- a/fieldExtraction/src/investment/main.py +++ b/fieldExtraction/src/investment/main.py @@ -1,5 +1,4 @@ import concurrent.futures -import gc import logging import os import random @@ -10,10 +9,7 @@ import pandas as pd random.seed(42) -import traceback - from constants.constants import Constants -from sentence_transformers import SentenceTransformer import src.investment.file_processing as file_processing import src.utils.io_utils as io_utils @@ -21,13 +17,11 @@ from src import config def safe_process_file(item, constants, run_timestamp): - """Process a single file with error handling and resource cleanup. + """Process a single file with comprehensive error handling. This function wraps the main file processing logic to ensure that: - Individual file failures don't crash the entire batch - - Resources are properly cleaned up after each file - Detailed error information is captured for debugging - - Garbage collection is forced to help prevent file descriptor leaks Args: item: Tuple of (filename, contract_text) representing the file to process @@ -41,16 +35,13 @@ def safe_process_file(item, constants, run_timestamp): Note: Always returns a DataFrame to maintain consistent output format for pd.concat. - Forces garbage collection after processing to help with resource cleanup, - particularly important for preventing file descriptor leaks in concurrent processing. + With FAISS migration, no special resource cleanup is needed. """ file_id = ( item[0] if isinstance(item, (list, tuple)) and len(item) > 0 else item ) # unpack file_id from item (passed in as a tuple below) try: result = file_processing.process_file(item, constants, run_timestamp) - # Force garbage collection after each file - gc.collect() return result except ( Exception @@ -65,8 +56,6 @@ def safe_process_file(item, constants, run_timestamp): logging.error(f"Full traceback:\n{full_traceback}") - # Force garbage collection on error too - gc.collect() return pd.DataFrame( [ { diff --git a/fieldExtraction/src/investment/smart_chunking_funcs.py b/fieldExtraction/src/investment/smart_chunking_funcs.py index 4ded0f6..020b1f4 100644 --- a/fieldExtraction/src/investment/smart_chunking_funcs.py +++ b/fieldExtraction/src/investment/smart_chunking_funcs.py @@ -4,59 +4,35 @@ Smart Chunking Functions for Contract Processing This module provides intelligent text chunking and semantic search capabilities for contract processing. It includes: -- Vectorstore caching to prevent "too many open files" errors in concurrent processing +- Lightweight in-memory vectorstore operations using FAISS - Multiple chunking strategies (hierarchical, OR, AND, regex) - Semantic and keyword-based chunk retrieval -- Thread-safe operations for multi-threaded contract processing -- Monitoring and cleanup to prevent file descriptor leaks ("too many open files" errors) +- Optimized for ephemeral, per-contract vectorstore usage Key Components: -- field_context(): Main function for semantic chunk retrieval with caching +- 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 module-level vectorstore caching to share embeddings across threads -- Reference counting ensures proper cleanup of temporary vectorstores -- Designed to handle concurrent processing without file descriptor leaks -- Emergency cleanup via atexit handler prevents resource leaks on process termination - -File Descriptor Management: -- Monitors open file descriptors to prevent "too many open files" errors -- Implements reference counting for shared vectorstores -- Forces Chroma client connection cleanup via _client.reset() -- Includes retry logic for directory cleanup operations -- Logs detailed statistics every 10 files for monitoring - -Troubleshooting: -- Check logs for "High number of open files" warnings -- Monitor cache hit/miss ratios for performance -- Watch for cleanup failure counts in cache statistics -- Emergency cleanup logs appear during process shutdown +- 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 atexit -import gc -import hashlib import json -import logging import os import re -import shutil -import threading -import time -import weakref from collections import defaultdict from itertools import combinations from typing import Any, Dict, List, Optional, Set, Tuple, Union -import chromadb import nltk -import psutil from constants.constants import Constants from langchain_aws.embeddings.bedrock import BedrockEmbeddings -from langchain_chroma import Chroma from langchain_community.retrievers import BM25Retriever +from langchain_community.vectorstores import FAISS from langchain_text_splitters import RecursiveCharacterTextSplitter from nltk.tokenize import word_tokenize @@ -70,72 +46,6 @@ nltk.data.path.append( ) ) -logging.basicConfig(level=logging.INFO) -logger = logging.getLogger(__name__) - -# Enable Chroma reset functionality -try: - # Enable reset functionality to allow proper cleanup - chromadb.config.Settings(allow_reset=True) - logger.info("Chroma reset functionality enabled") -except Exception as e: - logger.warning(f"Could not enable Chroma reset: {e}") - - -# =================================================================== -# VECTORSTORE CACHING -# =================================================================== -# Thread-safe cache to prevent "too many open files" errors -_vectorstore_lock = threading.Lock() -_vectorstore_cache: Dict[str, Dict[str, Any]] = {} - -# Track active vectorstores for cleanup -_active_vectorstores = weakref.WeakSet() - - -def _cleanup_all_vectorstores(): - """Emergency cleanup function for atexit""" - logger.info("Emergency cleanup starting...") - with _vectorstore_lock: - cleanup_count = 0 - for cache_key in list(_vectorstore_cache.keys()): - try: - vectorstore_data = _vectorstore_cache[cache_key] - vectorstore = vectorstore_data["vectorstore"] - - # Force close the vectorstore - if hasattr(vectorstore, "_client") and vectorstore._client: - vectorstore._client.reset() - if hasattr(vectorstore, "delete_collection"): - vectorstore.delete_collection() - - # Clean up directory - chroma_filename = vectorstore_data.get("chroma_filename") - if os.path.exists(chroma_filename): - shutil.rmtree(chroma_filename, ignore_errors=True) - - cleanup_count += 1 - except Exception as e: - logger.warning(f"Error during emergency cleanup of {cache_key}: {e}") - - _vectorstore_cache.clear() - logger.info( - f"Emergency cleanup completed. Cleaned up {cleanup_count} vectorstores." - ) - - -# Register emergency cleanup -atexit.register(_cleanup_all_vectorstores) - -# Monitoring counters -_cache_stats = { - "cache_hits": 0, - "cache_misses": 0, - "vectorstores_created": 0, - "vectorstores_cleaned": 0, - "cleanup_failures": 0, - "open_files_peak": 0, -} # Setup bedrock embeddings_model = BedrockEmbeddings( @@ -156,72 +66,6 @@ TOP_K = 8 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 get_open_file_count() -> int: - """Get current number of open file descriptors for monitoring. - - Uses Linux /proc filesystem when available, falls back to psutil. - Returns -1 if neither method works (shouldn't happen in normal operation). - - Returns: - int: Number of open file descriptors, or -1 if unable to determine - """ - - try: # Linux-specific: count files in /proc/PID/fd directory - return len(os.listdir(f"/proc/{os.getpid()}/fd")) - except: - # Fallback if /proc not available (e.g., on Windows or Mac) - try: - return len(psutil.Process().open_files()) - except: - return -1 - - -def log_cache_stats(): - """Log comprehensive cache statistics for monitoring and debugging. - - Logs include: - - Cache hit/miss ratios (higher hit ratio = better performance) - - Vectorstore creation/cleanup counts (should be roughly equal) - - Current and peak open file descriptor counts - - Active cache size (number of cached vectorstores) - - Cleanup failure count (should remain at 0) - """ - open_files = get_open_file_count() - _cache_stats["open_files_peak"] = max(_cache_stats["open_files_peak"], open_files) - - logger.info( - f"Cache Stats - Hits: {_cache_stats['cache_hits']}, " - f"Misses: {_cache_stats['cache_misses']}, " - f"Created: {_cache_stats['vectorstores_created']}, " - f"Cleaned: {_cache_stats['vectorstores_cleaned']}, " - f"Failures: {_cache_stats['cleanup_failures']}, " - f"Open Files: {open_files}, " - f"Peak Files: {_cache_stats['open_files_peak']}, " - f"Active Cache: {len(_vectorstore_cache)}" - ) - - -def monitor_file_descriptors(filename: str, operation: str) -> int: - """Monitor file descriptor usage and warn if approaching limits. - - Args: - filename: Name of file being processed (for logging context) - operation: Operation being performed (e.g., "start", "vectorstore_create") - - Returns: - int: Current number of open file descriptors - - Note: - Logs warning if file descriptor count exceeds 800 (threshold for 1024 limit) - """ - open_files = get_open_file_count() - if open_files > 800: # Warning threshold - logger.warning( - f"High number of open files ({open_files}) detected during {operation} on {filename}" - ) - return open_files - - def parse_chunk(chunk: str) -> tuple[int, str]: """Extract chunk number and content from a chunk string @@ -253,18 +97,6 @@ def find_overlap(prev_chunk, next_chunk, min_overlap=3): return 0 # No overlap found -def get_vectorstore_cache_key(contract_text: str) -> str: - """Generate a cache key for the vector store based on the contract text. - - Args: - contract_text (str): The text of the contract to generate a key for. - - Returns: - str: The generated cache key. - """ - return hashlib.md5(contract_text.encode()).hexdigest()[:16] - - 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 @@ -310,48 +142,11 @@ def stitch_chunks(chunks: list[str], overlap_size: int = CHUNK_OVERLAP) -> str: def field_context( one_to_one_fields: FieldSet, constants: Constants, contract_text: str, filename: str ) -> Tuple[FieldSet, List[str]]: - """Thread-safe field context extraction with vectorstore caching. - - This function creates chunks from whole contract text and returns relevant chunks - for each field based on semantic search and keyword search. - - This function is the core of the file descriptor leak prevention system. - It implements: - - Thread-safe vectorstore caching with reference counting - - Comprehensive cleanup with retries and forced garbage collection - - File descriptor monitoring throughout the process - - Emergency cleanup registration for process termination - - Args: - one_to_one_fields: Fields to process for semantic chunking - constants: Configuration constants - contract_text: Full contract text to chunk and embed - filename: Filename for logging context - - Returns: - Tuple of (updated fieldset with retrieved text, list of all chunks) - - Raises: - Various exceptions from Chroma/embedding operations - - Note: - This function is designed to be called concurrently from multiple threads. - The vectorstore cache prevents duplicate embedding work and the reference - counting ensures cleanup only happens when all threads are done. - """ - - # === MONITORING SETUP === - # Track processing time and file descriptor usage for this operation - start_time = time.time() - initial_files = monitor_file_descriptors(filename, "start_field_context") + """Lightweight field context extraction using FAISS (in-memory only).""" # Convert contract_text to lowercase for case-insensitive processing contract_text = contract_text.lower() - # Generate unique cache key for this contract - cache_key = get_vectorstore_cache_key(contract_text) - chroma_filename = f"chroma_cache_{cache_key}" - text_splitter = RecursiveCharacterTextSplitter( chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP ) @@ -363,172 +158,43 @@ def field_context( for idx, split in enumerate(splits) ] - vectorstore = None - try: - # === THREAD-SAFE CACHE ACCESS === - # Multiple threads may try to create / access the same vectorstore - # Lock ensures only one thread creates it, others wait and reuse - with _vectorstore_lock: - if cache_key not in _vectorstore_cache: - # === VECTORSTORE CREATION === - # This is the expensive operation we would like to cache - # Creates Chroma database with persistent storage for performance - logger.info(f"Creating new vectorstore for {filename} (cache miss)") - _cache_stats["cache_misses"] += 1 - _cache_stats["vectorstores_created"] += 1 + # Create lightweight FAISS vectorstore (in-memory only) + vectorstore = FAISS.from_texts(texts=splits, embedding=embeddings_model) - # Create vectorstore with persistence (for performance) - vectorstore = Chroma.from_texts( - texts=splits, - embedding=embeddings_model, - persist_directory=chroma_filename, - ) + vectorstore_retriever = vectorstore.as_retriever(search_kwargs={"k": TOP_K}) + keyword_retriever = BM25Retriever.from_texts( + splits, k=TOP_K, preprocess_func=word_tokenize + ) - _vectorstore_cache[cache_key] = { - "vectorstore": vectorstore, - "splits": splits, - "chroma_filename": chroma_filename, - "ref_count": 0, - "created_at": time.time(), - } + 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) - # === EMERGENCY CLEANUP REGISTRATION === - # Add to weak set so emergency cleanup can find it if process crashes - _active_vectorstores.add(vectorstore) # Track for emergency cleanup + # Convert field keywords to lowercase for case-insensitive processing + field_keywords = [keyword.lower() for keyword in field.keywords] - after_create_files = monitor_file_descriptors( - filename, "vectorstore_create" - ) - logger.info( - f"Vectorstore created. Files {initial_files} -> {after_create_files}" - ) - else: - # === CACHE HIT === - # Reusing existing vectorstore saves significant time and resources - logger.info(f"Using cached vectorstore for {filename} (cache hit)") - _cache_stats["cache_hits"] += 1 + # 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)) + ) - # === REFERENCE COUNTING === - # Track how many threads are using this vectorstore - # Only cleanup when count reaches 0 - _vectorstore_cache[cache_key]["ref_count"] += 1 - vectorstore_data = _vectorstore_cache[cache_key] - - # Use the cached/created vectorstore - vectorstore = vectorstore_data["vectorstore"] - splits = vectorstore_data["splits"] - - 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)) - ) - - # Check if no context or full_context, if so, switch attribute - if len(field.retrieved_text) == 0 or len(field.retrieved_text) == len( - splits - ): - field.field_type = "full_context" - - finally: - # === CLEANUP WITH REFERENCE COUNTING === - # This section is critical for preventing file descriptor leaks - with _vectorstore_lock: - if cache_key in _vectorstore_cache: - _vectorstore_cache[cache_key]["ref_count"] -= 1 - - # === ACTUAL CLEANUP === - # Only cleanup when no other threads are using this vectorstore - if _vectorstore_cache[cache_key]["ref_count"] == 0: - logger.info(f"Cleaning up vectorstore for {filename}") - try: - vectorstore_to_cleanup = _vectorstore_cache[cache_key][ - "vectorstore" - ] - chroma_filename_to_cleanup = _vectorstore_cache[cache_key][ - "chroma_filename" - ] - - # === FORCE CONNECTION CLEANUP === - # Directory deletion can fail due to timing issues - # Retry with small delays to handle race conditions - if ( - hasattr(vectorstore_to_cleanup, "_client") - and vectorstore_to_cleanup._client - ): - vectorstore_to_cleanup._client.reset() - - # Delete collection first - if hasattr(vectorstore_to_cleanup, "delete_collection"): - vectorstore_to_cleanup.delete_collection() - - # === RETRY LOGIC === - # Directory deletion can fail due to timing issues - # Retry with small delays to handle race conditions - # Clean up directory with retries - for attempt in range(3): - try: - if os.path.exists(chroma_filename_to_cleanup): - shutil.rmtree( - chroma_filename_to_cleanup, ignore_errors=True - ) - break # Success - except OSError as e: - if attempt == 2: # Last attempt - logger.warning( - f"Failed to cleanup {chroma_filename_to_cleanup} after 3 attempts: {e}" - ) - _cache_stats["cleanup_failures"] += 1 - else: - time.sleep(0.1) # Brief delay before retry - - # Force garbage collection - gc.collect() - _cache_stats["vectorstores_cleaned"] += 1 - - except Exception as e: - logger.warning(f"Error cleaning up {cache_key}: {e}") - _cache_stats["cleanup_failures"] += 1 - finally: - # Always remove from cache - del _vectorstore_cache[cache_key] - - end_time = time.time() - final_files = monitor_file_descriptors(filename, "completion") - processing_time = end_time - start_time - logger.info( - f"Completed {filename} in {processing_time:.2f}s. " - f"Files: {initial_files} -> {final_files}" - ) - - # Log stats every 10 files - if (_cache_stats["cache_hits"] + _cache_stats["cache_misses"]) % 10 == 0: - log_cache_stats() + # 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