jacob-mathison/doczyai-pipelines
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Activity

Merged in feature/multithreading (pull request #828)

Browse Source 
Feature/multithreading

* fixes

* Merge main into feature/multithreading

Resolved conflicts:
- Kept timing instrumentation in file_processing.py
- Kept new 3-step Exhibit-based approach for one-to-n processing
- Maintained parallelization improvements (20 workers)

Changes include:
- Timing utils integration for performance monitoring
- Increased max_workers from 5 to 20 across all components
- Parallelized code_breakout and grouper_breakout
- Fixed tin_npi_funcs function call parameters

* Fix max_workers error for empty documents

- Add check to skip parallel processing when no pages exist
- Use min(len(all_page_tasks), 20) to prevent max_workers=0
- Handles edge case of documents with no exhibits or pages

* Fix max_workers=0 errors in one_to_n_funcs

- Add checks before all ThreadPoolExecutor creations
- Prevents errors when processing empty lists:
  - carveout_and_special_case
  - breakout
  - special_case_breakout
  - filter_services_without_reimbursements
  - run_lob_relationship
- Ensures executor only created when there are items to process

* Reduce code processing parallelism to prevent API throttling

Lower max_workers from 20 to 10 for code_breakout and grouper_breakout
to prevent overwhelming Bedrock API with concurrent requests

* fixed

* Merged main into feature/multithreading

* Move documentation into folder

* Fix logging statements

* Merge branch 'main' into feature/multithreading

* Refactor for clarity

* Update previous exhibit passing logic

* properly simplify exhibits

* Merged main into feature/multithreading

* update conditional for None

* exhibit multithreading changed

* exhibit multithreading changed

* Merge remote-tracking branch 'origin/main' into feature/multithreading

* synced with main

* Parallelize dynamic assignment, refactor HSC field worker, add timing/exhibit unit tests, tidy imports/ignore helpers

* analyze_regression.py edited online with Bitbucket
* count_pages.py edited online with Bitbucket
* compare_regressed_with_baseline.py edited online with Bitbucket
* simple_testbed_compare.py edited online with Bitbucket
* run_testbed_metrics_regressed.py edited online with Bitbucket

Approved-by: Katon Minhas
This commit is contained in:
Faizan Mohiuddin
2026-01-09 20:18:01 +00:00
committed by Katon Minhas
parent 69f23653a1
commit 8c9060e425
17 changed files with 2633 additions and 479 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4 -1
View File
@@ -18,7 +18,7 @@ dist/
# Package files
*.jar
*.png
# Maven
target/
dist/
@@ -103,6 +103,9 @@ new/
# Embeddings
fieldExtraction/embeddings/
# Output reports
fieldExtraction/output_reports/
# PRD Documents (local only, living documents)
*.prd
*prd.md
fieldExtraction/documentation/EXTRACTION_FLOW_VISUAL.md
+650
View File
@@ -0,0 +1,650 @@
# Field Extraction Visual Flow Guide
## Complete System Flow Diagram
```
┌─────────────────────────────────────────────────────────────────────┐
│ START: main.py │
│ │
│ 1. Load configuration (batch_id, fields, max_workers) │
│ 2. Initialize batch tracker │
│ 3. Read input contracts (from local folder or S3) │
│ 4. Filter contracts with "reimbursement" for B processing │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ Split Processing Paths │
└─────────────────────────────────────────────────────────────────────┘
↓ ↓
╔═══════════════════════╗ ╔═══════════════════════╗
║ AC PROCESSING ║ ║ B PROCESSING ║
║ (Contract Fields) ║ ║ (Reimbursement) ║
╚═══════════════════════╝ ╚═══════════════════════╝
↓ ↓
[See AC Flow Diagram] [See B Flow Diagram]
↓ ↓
╔═══════════════════════╗ ╔═══════════════════════╗
║ AC Output CSV ║ ║ B Output CSV ║
╚═══════════════════════╝ ╚═══════════════════════╝
↓ ↓
└───────────┬───────────────┘
↓
┌───────────────────────────────┐
│ Consolidation & QA/QC │
│ - Merge AC + B data │
│ - Quality checks │
│ - Generate reports │
└───────────────────────────────┘
↓
┌───────────────────────────────┐
│ Final Outputs │
│ - ABC combined CSV │
│ - QA/QC Excel report │
└───────────────────────────────┘
```
---
## AC Processing Detailed Flow
```
┌─────────────────────────────────────────────────────────────────────┐
│ AC PROCESSING START │
│ (file_processing.py: run_ac_prompts) │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 1: PREPROCESSING (preprocess.py) │
│ │
│ Input: Raw contract text │
│ │
│ Actions: │
│ ├─ Clean text (remove symbols, fix newlines) │
│ ├─ Split into pages │
│ └─ Create smart chunks based on keywords │
│ │
│ Output: text_dict (pages), ac_chunks (focused text per field) │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 2: REGEX EXTRACTION │
│ │
│ Extract using patterns: │
│ ├─ TIN (Tax ID): XX-XXXXXXX format │
│ ├─ NPI (Provider ID): 10-digit number │
│ └─ Other IRS names from specific pages │
│ │
│ Output: Partial ac_answers_dict │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 3: SMART CHUNKED PROMPTS (smart_chunking_funcs.py) │
│ │
│ For each field group with keyword matches: │
│ │
│ ┌─────────────────────────────────────────────────┐ │
│ │ Example: CONTRACT_EFFECTIVE_DT │ │
│ ├─────────────────────────────────────────────────┤ │
│ │ 1. Keywords found: "effective date", "execution"│ │
│ │ 2. Chunk contains: Pages 1-3 │ │
│ │ 3. Build prompt with focused question │ │
│ │ 4. Send to Claude AI │ │
│ │ 5. Parse response: "2024-01-15" │ │
│ │ 6. Store in ac_answers_dict │ │
│ └─────────────────────────────────────────────────┘ │
│ │
│ Repeat for ~40 field groups │
│ │
│ Output: ac_answers_dict with most fields filled │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 4: FULL CONTEXT PROMPTS │
│ │
│ For fields with no keyword matches: │
│ ├─ Use entire contract as context │
│ ├─ More expensive but comprehensive │
│ └─ Ensures rare fields aren't missed │
│ │
│ Output: Remaining fields filled │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 5: CONDITIONAL PROMPTS │
│ │
│ Context-based derivations: │
│ ├─ Use existing answers to inform new questions │
│ └─ Example: If effective date is N/A, check signature date │
│ │
│ Output: Final ac_answers_dict │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 6: POSTPROCESSING (postprocess.py: ac_postprocess) │
│ │
│ Cleanup and standardization: │
│ ├─ Convert dates to YYYY-MM-DD format │
│ ├─ Clean TIN/NPI formatting │
│ ├─ Derive indicator fields (Y/N from text) │
│ ├─ Add metadata (filename, page count, parent code) │
│ ├─ Remove null/invalid values │
│ ├─ Rename columns to user-friendly names │
│ └─ Reorder columns │
│ │
│ Output: Clean DataFrame │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 7: WRITE OUTPUT │
│ │
│ Save to: {batch_id}-AC.csv │
│ Format: One row per contract, ~45 columns │
└─────────────────────────────────────────────────────────────────────┘
```
---
## B Processing Detailed Flow
```
┌─────────────────────────────────────────────────────────────────────┐
│ B PROCESSING START │
│ (file_processing.py: run_b_prompts) │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 1: PREPROCESSING (preprocess.py) │
│ │
│ Input: Raw contract text │
│ │
│ Actions: │
│ ├─ Clean text (symbols, newlines) │
│ ├─ Split into pages │
│ ├─ Identify exhibit pages (headers like "Exhibit", "Attachment") │
│ ├─ Chunk consecutive pages together │
│ └─ Align and format tables │
│ │
│ Output: text_dict with chunked pages and aligned tables │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 2: BOTTOM-UP PRIMARY (one_to_n_funcs.py) │
│ │
│ Process each page/chunk with reimbursement info: │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ For Page/Chunk: │ │
│ │ │ │
│ │ Prompt (prompts.py: BOTTOM_UP_PRIMARY): │ │
│ │ "Extract EVERY reimbursement term on this page" │ │
│ │ │ │
│ │ Input Page Example: │ │
│ │ ┌────────────────────────────────────────┐ │ │
│ │ │ PHYSICIAN SERVICES │ │ │
│ │ │ Office Visits: 110% Medicare │ │ │
│ │ │ Surgery: 120% Medicare │ │ │
│ │ │ Lab Tests: $25.00 per test │ │ │
│ │ └────────────────────────────────────────┘ │ │
│ │ │ │
│ │ AI Extraction: │ │
│ │ [ │ │
│ │ { │ │
│ │ "FULL_SERVICE": "Office Visits", │ │
│ │ "FULL_METHODOLOGY": "110% Medicare", │ │
│ │ "PROV_TYPE": "Professional" │ │
│ │ }, │ │
│ │ { │ │
│ │ "FULL_SERVICE": "Surgery", │ │
│ │ "FULL_METHODOLOGY": "120% Medicare", │ │
│ │ "PROV_TYPE": "Professional" │ │
│ │ }, │ │
│ │ { │ │
│ │ "FULL_SERVICE": "Lab Tests", │ │
│ │ "FULL_METHODOLOGY": "$25.00 per test", │ │
│ │ "PROV_TYPE": "Ancillary" │ │
│ │ } │ │
│ │ ] │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
│ Output: List of dictionaries (one per service found) │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 3: BOTTOM-UP SECONDARY (one_to_n_funcs.py) │
│ │
│ For each extracted service, break down methodology: │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ Methodology Breakout │ │
│ │ │ │
│ │ Prompt (prompts.py: BOTTOM_UP_METHODOLOGY_BREAKOUT): │ │
│ │ "Parse this methodology into structured fields" │ │
│ │ │ │
│ │ Input: "Lesser of billed charges or 110% Medicare" │ │
│ │ │ │
│ │ AI Breakdown: │ │
│ │ { │ │
│ │ "LESSER": "Y", │ │
│ │ "RATE_STANDARD": "110% Medicare", │ │
│ │ "RATE_SHORT": "1.1", │ │
│ │ "FLAT_FEE_STANDARD": "N/A", │ │
│ │ "LESSER_RATE": "100% of BC", │ │
│ │ "NOT_TO_EXCEED": "N/A" │ │
│ │ } │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ Escalator Detection │ │
│ │ │ │
│ │ Prompt (prompts.py: BOTTOM_UP_ESCALATOR): │ │
│ │ "Does this rate increase over time?" │ │
│ │ │ │
│ │ Output: │ │
│ │ { │ │
│ │ "RATE_ESCALATOR_IND": "Y", │ │
│ │ "RATE_ESCALATOR_DT": "2025-01-01" │ │
│ │ } │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
│ Output: Enhanced list with structured rate information │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 4: TOP-DOWN EXTRACTION (one_to_n_funcs.py) │
│ │
│ Extract page-level context that applies to all services: │
│ │
│ Prompt (prompts.py: TOP_DOWN_PRIMARY): │
│ ├─ EXHIBIT: Full exhibit/attachment name │
│ ├─ CONTRACT_LOB: Line of Business (Medicare, Medicaid, etc.) │
│ ├─ DEFAULT_TERM: What happens if no rate specified │
│ ├─ CDM_IND: Chargemaster neutralization? │
│ └─ ADD_ON_REIMBURSEMENT_LANGUAGE: Extra payment clauses │
│ │
│ Example Output: │
│ { │
│ "EXHIBIT": "Attachment C: Medicare Advantage", │
│ "CONTRACT_LOB": "Medicare Advantage", │
│ "DEFAULT_TERM": "100% of Medicare Fee Schedule", │
│ "CDM_IND": "N" │
│ } │
│ │
│ This context is merged with ALL services on that page │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 5: CONDITIONAL PROMPTS (conditional_funcs.py) │
│ │
│ Fill in additional service-level details: │
│ │
│ ├─ IP/OP: Inpatient or Outpatient? │
│ │ Prompt: "Is this service IP or OP?" │
│ │ Uses keywords like "inpatient", "outpatient" │
│ │ │
│ ├─ PROV_TYPE_2: More specific provider category │
│ │ Prompt: "What specific provider type?" │
│ │ Options: PCP, Specialist, ER, Clinic, etc. │
│ │ │
│ └─ LOB_CHECK: Assign single LOB to each service │
│ Prompt: "Which LOB does this service fall under?" │
│ Narrows from page-level to service-level │
│ │
│ Output: Fully enriched service list │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 6: POSTPROCESSING (postprocess.py: b_postprocess) │
│ │
│ Cleanup and standardization: │
│ ├─ Add metadata (filename, parent code, page count) │
│ ├─ Flag single code, multiple rates (SCMR) │
│ ├─ Filter out add-on services │
│ ├─ Clean methodology text │
│ ├─ Standardize provider types │
│ ├─ Clean lesser of rates │
│ ├─ Standardize Line of Business values │
│ ├─ Fix Y/N fields │
│ ├─ Rename columns │
│ └─ Reorder columns │
│ │
│ Output: Clean DataFrame │
└─────────────────────────────────────────────────────────────────────┘
↓
┌─────────────────────────────────────────────────────────────────────┐
│ STEP 7: WRITE OUTPUT │
│ │
│ Save to: {batch_id}-B.csv │
│ Format: Multiple rows per contract (one per service), ~35 columns │
└─────────────────────────────────────────────────────────────────────┘
```
---
## Prompt Flow for Key Extractions
### Contract Effective Date Extraction
```
┌────────────────────────────────────────────────────┐
│ CONTRACT_EFFECTIVE_DT EXTRACTION │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Smart Chunking │
│ Keywords: "effective date", "execution date" │
│ Result: Pages 1-3 contain keywords │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Build Prompt (prompts.py: get_effective_date_prompt) │
│ │
│ Context: [Page 1-3 text] │
│ │
│ Instruction: │
│ "Extract ONLY the effective date of the contract │
│ Follow these rules: │
│ 1. Look for 'Effective Date:' label │
│ 2. Don't use signature date │
│ 3. Return in YYYY-MM-DD format or N/A │
│ 4. Enclose answer in |pipes|" │
│ │
│ Example Text: │
│ "This Agreement is effective as of January 15, │
│ 2024 ('Effective Date')..." │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Send to Claude AI │
│ Model: Claude 3.5 Sonnet │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ AI Response │
│ "...reviewing the contract text... │
│ Answer: |2024-01-15|" │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Parse Response │
│ - Extract text between |pipes| │
│ - Result: "2024-01-15" │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Validate & Convert │
│ - Check format is YYYY-MM-DD │
│ - If invalid, set to N/A │
└────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────┐
│ Store in Output │
│ CONTRACT_EFFECTIVE_DT = "2024-01-15" │
└────────────────────────────────────────────────────┘
```
### Reimbursement Methodology Extraction
```
┌────────────────────────────────────────────────────────────────┐
│ REIMBURSEMENT METHODOLOGY EXTRACTION │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ Page Text (Example) │
│ ──────────────────────────────────────────────────────────── │
│ COMPENSATION SCHEDULE - PROFESSIONAL SERVICES │
│ │
│ CCA shall reimburse Provider for Covered Services rendered │
│ to Covered Persons the lesser of: │
│ (i) Allowable Charges; or │
│ (ii) The rates specified below: │
│ │
│ Service Category | Contracted Rate │
│ ──────────────────────────────────────────── │
│ Office Visits | 110% of Medicare Fee Schedule │
│ Preventive Care | 100% of Medicare Fee Schedule │
│ Specialist Consults | 125% of Medicare Fee Schedule │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ STAGE 1: Bottom-Up Primary │
│ Prompt: BOTTOM_UP_PRIMARY │
│ │
│ Instruction: "Extract EVERY reimbursement term" │
│ │
│ Key Rules: │
│ - Include complete "lesser of" statements │
│ - Capture table headers for context │
│ - One entry per service │
│ - Include provider type │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ AI Response (JSON) │
│ [ │
│ { │
│ "FULL_SERVICE": "Office Visits", │
│ "SUBHEADER": "N/A", │
│ "FULL_METHODOLOGY": "Lesser of: (i) Allowable Charges; │
│ or (ii) The rates specified below: Service Category | │
│ Contracted Rate\nOffice Visits | 110% of Medicare │
│ Fee Schedule", │
│ "PROV_TYPE": "Professional" │
│ }, │
│ { │
│ "FULL_SERVICE": "Preventive Care", │
│ "SUBHEADER": "N/A", │
│ "FULL_METHODOLOGY": "Lesser of: (i) Allowable Charges; │
│ or (ii) The rates specified below: Service Category | │
│ Contracted Rate\nPreventive Care | 100% of Medicare │
│ Fee Schedule", │
│ "PROV_TYPE": "Professional" │
│ }, │
│ { │
│ "FULL_SERVICE": "Specialist Consults", │
│ "SUBHEADER": "N/A", │
│ "FULL_METHODOLOGY": "Lesser of: (i) Allowable Charges; │
│ or (ii) The rates specified below: Service Category | │
│ Contracted Rate\nSpecialist Consults | 125% of │
│ Medicare Fee Schedule", │
│ "PROV_TYPE": "Professional" │
│ } │
│ ] │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ STAGE 2: Bottom-Up Secondary (Methodology Breakout) │
│ │
│ For each service, parse methodology: │
│ │
│ Input (Office Visits): │
│ "Lesser of: (i) Allowable Charges; or (ii) 110% of Medicare │
│ Fee Schedule" │
│ │
│ Prompt: BOTTOM_UP_METHODOLOGY_BREAKOUT │
│ "Break this methodology into structured fields" │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ AI Response (Office Visits) │
│ { │
│ "LESSER": "Y", │
│ "RATE_STANDARD": "110% Medicare", │
│ "RATE_SHORT": "1.1", │
│ "FLAT_FEE_STANDARD": "N/A", │
│ "LESSER_RATE": "100% of AC", // AC = Allowable Charges │
│ "NOT_TO_EXCEED": "N/A" │
│ } │
│ │
│ Repeat for Preventive Care → RATE_STANDARD: "100% Medicare" │
│ Repeat for Specialist → RATE_STANDARD: "125% Medicare" │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ STAGE 3: Top-Down Context │
│ │
│ Extract page-level info: │
│ { │
│ "EXHIBIT": "COMPENSATION SCHEDULE - PROFESSIONAL SERVICES", │
│ "CONTRACT_LOB": "Commercial", │
│ "DEFAULT_TERM": "N/A", │
│ "CDM_IND": "N" │
│ } │
│ │
│ Merge with each service │
└────────────────────────────────────────────────────────────────┘
↓
┌────────────────────────────────────────────────────────────────┐
│ FINAL OUTPUT (Office Visits row) │
│ ──────────────────────────────────────────────────────────── │
│ Contract Name: "contract_123.pdf" │
│ Exhibit: "COMPENSATION SCHEDULE - PROFESSIONAL SERVICES" │
│ Line of Business: "Commercial" │
│ Full Service: "Office Visits" │
│ Provider Type: "Professional" │
│ Reimbursement Methodology: "Lesser of: (i) Allowable │
│ Charges; or (ii) 110% of Medicare Fee Schedule" │
│ Rate Standard: "110% Medicare" │
│ Rate Short: "1.1" │
│ Flat Fee: "N/A" │
│ Lesser Of: "Y" │
│ Lesser Of Rate: "100% of AC" │
│ Not to Exceed: "N/A" │
│ ... (additional columns) │
└────────────────────────────────────────────────────────────────┘
```
---
## Error Handling Flow
```
┌────────────────────────────────────────────┐
│ File Processing Attempt │
└────────────────────────────────────────────┘
↓
┌──────────────────────┐
│ Success? │
└──────────────────────┘
↓ Yes ↓ No
┌─────────┐ ┌─────────────────────┐
│ Log │ │ Capture Error │
│ Success │ │ - Error message │
│ │ │ - Stack trace │
│ Update │ │ - Timestamp │
│ Tracking│ └─────────────────────┘
│ │ ↓
│ Continue│ ┌─────────────────────┐
│ to Next │ │ Log to Tracking │
└─────────┘ │ Status: Failed │
│ Error: [details] │
└─────────────────────┘
↓
┌─────────────────────┐
│ Mark for Self-Repair│
└─────────────────────┘
↓
┌────────────────────────────────────┐
│ After All Files Processed │
└────────────────────────────────────┘
↓
┌────────────────────────────────────┐
│ Self-Repair Mode │
│ - Identify failed files │
│ - Retry processing │
│ - Consolidate outputs │
└────────────────────────────────────┘
```
---
## Data Flow Summary
```
Raw Contract Text
↓
┌─────────────┐
│ Preprocessing│
│ - Clean │
│ - Split │
│ - Chunk │
└─────────────┘
↓
┌─────────────┐ ┌──────────────┐
│ AC Path │ │ B Path │
│ │ │ │
│ Smart Chunk │ │ Bottom-Up │
│ Prompts │ │ Primary │
│ ↓ │ │ ↓ │
│ Full Context│ │ Bottom-Up │
│ Prompts │ │ Secondary │
│ ↓ │ │ ↓ │
│ Conditional │ │ Top-Down │
│ Prompts │ │ ↓ │
│ ↓ │ │ Conditional │
│ Postprocess │ │ ↓ │
│ │ │ Postprocess │
└─────────────┘ └──────────────┘
↓ ↓
┌─────────────┐ ┌──────────────┐
│ AC CSV │ │ B CSV │
│ (1 row) │ │ (many rows) │
└─────────────┘ └──────────────┘
↓ ↓
└──────────┬──────────┘
↓
┌───────────────┐
│ Consolidation │
│ - Merge AC+B │
│ - QA/QC │
└───────────────┘
↓
┌───────────────┐
│ Final Outputs │
│ - ABC CSV │
│ - QA Report │
└───────────────┘
```
---
## Parallel Processing Visualization
```
Main Process
│
┌────────────────┼────────────────┐
│ │ │
Worker 1 Worker 2 Worker 3
│ │ │
Contract A Contract D Contract G
↓ ↓ ↓
[AC + B] [AC + B] [AC + B]
Processing Processing Processing
↓ ↓ ↓
Success Success Failed
│ │ │
└────────────────┼────────────────┘
│
Track Results
│
┌────────────────┼────────────────┐
│ │ │
Worker 1 Worker 2 Worker 3
│ │ │
Contract B Contract E Contract H
↓ ↓ ↓
Processing Processing Processing
│ │ │
└────────────────┼────────────────┘
│
↓
Upload Progress (every 10 files)
│
↓
Continue...
```
---
This visual guide complements the detailed text guide by showing the exact flow of data and processing through the system in an easy-to-understand diagram format.
fieldExtraction/documentation/FIELD_EXTRACTION_GUIDE.md
+632
View File
@@ -0,0 +1,632 @@
# Field Extraction System - Complete Guide
## What This System Does
This system reads healthcare provider contracts (like agreements between insurance companies and doctors or hospitals) and automatically extracts important information from them. Think of it as a smart assistant that reads through hundreds of pages of legal contracts and pulls out key details like payment rates, contract dates, provider names, and reimbursement terms.
---
## Table of Contents
1. [Overview](#overview)
2. [How to Run the System](#how-to-run-the-system)
3. [Complete Processing Flow](#complete-processing-flow)
4. [Field Types Explained](#field-types-explained)
5. [The Extraction Process Step-by-Step](#the-extraction-process-step-by-step)
6. [Understanding the Prompts](#understanding-the-prompts)
7. [Output Files](#output-files)
---
## Overview
### What Gets Extracted?
The system extracts two main categories of information:
1. **AC Fields** (Administrative/Contract Fields): Basic contract information like:
- Contract title and effective date
- Payer name (insurance company)
- Provider names and identification numbers
- Termination clauses
- Legal terms and conditions
2. **B Fields** (Benefit/Reimbursement Fields): Payment-related information like:
- What services are covered
- How much providers get paid
- Payment methodologies (percentages of Medicare/Medicaid rates, flat fees)
- Which line of business the rates apply to (Medicare, Medicaid, Commercial)
---
## How to Run the System
### Starting the Process
The main entry point is **[main.py](fieldExtraction/src/main.py)**. You can run it with different configurations:
```bash
python main.py batch_id=my_batch fields=abc max_workers=5
```
**Key Parameters:**
- `batch_id`: A unique name for this extraction run (e.g., "batch_january_2025")
- `fields`: What to extract - "ac" for contract fields, "b" for reimbursement, "abc" for both
- `max_workers`: How many contracts to process at the same time (parallel processing)
- `read_mode`: Where to read files from ("local" or "s3")
- `write_to_s3`: Whether to save results to cloud storage
---
## Complete Processing Flow
### The Big Picture
```
Input Contracts (PDF → Text)
↓
main.py starts
↓
┌────────────────────────────┐
│ Split into two paths │
└────────────────────────────┘
↓ ↓
AC Processing B Processing
(Contract Info) (Reimbursement)
↓ ↓
Preprocessing Preprocessing
↓ ↓
AI Extraction AI Extraction
↓ ↓
Postprocessing Postprocessing
↓ ↓
AC Output CSV B Output CSV
↓ ↓
┌────────────────────────────┐
│ Combine & Quality Check │
└────────────────────────────┘
↓
Final Reports
```
### Step-by-Step Flow
#### 1. **Initial Setup** ([main.py:220-244](fieldExtraction/src/main.py#L220-L244))
When you run the system:
- Creates a batch tracker to monitor progress
- Reads all contract text files from the input folder (or S3 bucket)
- Checks which files contain "reimbursement" information
- Filters out already-processed files if requested
#### 2. **Parallel Processing** ([main.py:296-334](fieldExtraction/src/main.py#L296-L334))
The system processes multiple contracts at once:
- Creates separate processing tasks for AC and B fields
- Runs them in parallel using multiple "workers" (like having multiple assistants)
- Tracks progress and handles errors for each contract
- Updates progress statistics in real-time
#### 3. **Contract Processing Paths**
**For Each Contract:**
- If processing AC fields → calls `process_ac()`
- If processing B fields → calls `process_b()`
- Each path follows its own specialized extraction process
---
## Field Types Explained
### AC Fields (Contract Information)
These fields capture the administrative details of the contract. Examples:
- **CONTRACT_TITLE**: The official name of the agreement
- **PAYER_NAME**: The insurance company (e.g., "Centene", "Superior Health Plan")
- **PROV_GROUP_NAME**: The doctor's office or hospital name
- **CONTRACT_EFFECTIVE_DT**: When the contract starts
- **TERMINATION_UPON_NOTICE**: How many days notice to cancel
- **TERM_CLAUSE**: The full legal text about contract duration
**Why these matter**: These fields help understand who is involved in the contract, when it's valid, and basic legal terms.
### B Fields (Reimbursement Information)
These fields capture payment details. Examples:
- **FULL_SERVICE**: What medical service is being paid for (e.g., "Office Visits", "Surgery", "Lab Tests")
- **FULL_METHODOLOGY**: The complete description of how payment is calculated
- **RATE_STANDARD**: The payment rate (e.g., "110% Medicare" or "85% Medicaid")
- **FLAT_FEE_STANDARD**: Fixed dollar amounts (e.g., "$50.00 Per Visit")
- **LESSER_RATE**: If there's a "lesser of" clause (e.g., "100% of Billed Charges")
- **PROV_TYPE**: Whether it's for doctors (Professional), hospitals (Facility), or other services (Ancillary)
**Why these matter**: These tell us exactly how much a provider gets paid for each type of service.
---
## The Extraction Process Step-by-Step
### AC Processing Flow
**Entry Point**: [file_processing.py:15-100](fieldExtraction/src/file_processing.py#L15-L100)
#### Step 1: Preprocessing ([preprocess.py:53-56](fieldExtraction/src/preprocess.py#L53-L56))
**What happens:**
1. Clean up the text (remove weird symbols, fix line breaks)
2. Split the contract into individual pages
3. Create "smart chunks" - group pages by topic using keywords
**Smart Chunking** ([smart_chunking_funcs.py](fieldExtraction/src/smart_chunking_funcs.py)):
- Looks for specific keywords related to each field
- Example: For "CONTRACT_EFFECTIVE_DT", it searches for pages containing "effective date", "execution date", etc.
- Only sends relevant pages to AI, saving time and cost
#### Step 2: Extract Basic Information
**Regex-based extraction** ([file_processing.py:36-40](fieldExtraction/src/file_processing.py#L36-L40)):
- Some fields are found using patterns (like phone numbers or tax IDs)
- TIN (Tax ID): Looks for 9-digit numbers with hyphen (XX-XXXXXXX)
- NPI (Provider ID): Looks for 10-digit numbers
#### Step 3: Smart Chunked Prompts ([smart_chunking_funcs.py:33-150](fieldExtraction/src/smart_chunking_funcs.py#L33-L150))
**For each field or group of fields:**
1. Take the relevant chunk of text (found by keywords)
2. Build a specific question or set of questions
3. Send to AI (Claude) with the question and text
4. Parse the AI's response
5. Store the answer
**Example for Contract Effective Date**:
```
Context: [Pages 1-3 of the contract]
Question: "What is the contract effective date?"
AI Response: "2024-01-15"
System converts to: "2024-01-15" (standardized format)
```
#### Step 4: Full Context Prompts
**For fields not found in chunks:**
- Uses the entire contract as context
- More expensive but ensures nothing is missed
- Used for rare fields or when keyword search fails
#### Step 5: Postprocessing ([postprocess.py:55-114](fieldExtraction/src/postprocess.py#L55-L114))
**Clean up and validate:**
- Convert dates to standard format (YYYY-MM-DD)
- Clean up Tax IDs and NPIs
- Derive indicator fields (Yes/No flags)
- Add metadata (filename, page count)
- Remove invalid values
- Organize columns in standard order
### B Processing Flow
**Entry Point**: [file_processing.py:103-153](fieldExtraction/src/file_processing.py#L103-L153)
#### Step 1: Preprocessing ([preprocess.py:23-50](fieldExtraction/src/preprocess.py#L23-L50))
**What happens:**
1. Clean the text
2. Find pages with "Exhibit" or "Attachment" headers (where reimbursement info lives)
3. Chunk consecutive pages together (keeps related compensation schedules together)
4. Detect and format tables (many rates are in table format)
**Table Alignment**:
- Contracts often have messy tables
- System realigns columns and rows
- Preserves the structure for accurate extraction
#### Step 2: Bottom-Up Extraction ([one_to_n_funcs.py](fieldExtraction/src/one_to_n_funcs.py))
**Primary Bottom-Up** ([prompts.py:9-196](fieldExtraction/src/prompts.py#L9-L196)):
This is the core reimbursement extraction. For each page:
**Prompt asks AI to find**:
- Every service mentioned (e.g., "Office Visits", "Emergency Room")
- Complete payment methodology
- Provider type (Professional/Facility/Ancillary)
- Any subheaders or categories
**Example extraction**:
```
Page text: "Physician Services: 110% of Medicare Fee Schedule"
AI extracts:
{
"FULL_SERVICE": "Physician Services",
"FULL_METHODOLOGY": "110% of Medicare Fee Schedule",
"PROV_TYPE": "Professional"
}
```
**Critical Rules in the Prompt**:
- Extract EVERY rate on the page (don't skip anything)
- Include "lesser of" statements (e.g., "lesser of billed charges or 100% Medicare")
- Capture complete table context (headers + values)
- Split services with multiple rates into separate entries
**Secondary Bottom-Up** ([prompts.py:198-458](fieldExtraction/src/prompts.py#L198-L458)):
Breaks down each methodology into structured fields:
**Methodology Breakout Prompt**:
```
Input: "Lesser of billed charges or 110% Medicare Fee Schedule"
AI extracts:
{
"LESSER": "Y",
"RATE_STANDARD": "110% Medicare",
"RATE_SHORT": "1.1",
"LESSER_RATE": "100% of BC" (Billed Charges)
}
```
**Escalator Detection**:
- Checks if rates increase over time
- Finds effective dates for rate changes
#### Step 3: Top-Down Extraction ([prompts.py:482-554](fieldExtraction/src/prompts.py#L482-L554))
**Extracts page-level information**:
- Exhibit names (e.g., "Attachment C: Medicare Advantage")
- Line of Business (Medicare, Medicaid, Commercial, Dual)
- Default terms (what happens if no rate is specified)
- Chargemaster protections
- Add-on reimbursement language
**Why both Bottom-Up and Top-Down?**
- Bottom-up finds specific rates and services
- Top-down finds context that applies to all services on a page
- Combined, they give complete information
#### Step 4: Conditional Prompts ([conditional_funcs.py](fieldExtraction/src/conditional_funcs.py))
**Fills in additional details**:
- Inpatient vs Outpatient (IP/OP)
- More specific provider types
- Line of Business assignment for each service
#### Step 5: Postprocessing ([postprocess.py:5-52](fieldExtraction/src/postprocess.py#L5-L52))
**Final cleanup**:
- Filter out invalid services
- Clean methodology text
- Standardize Line of Business values
- Add metadata fields
- Rename columns to user-friendly names
- Reorder columns
---
## Understanding the Prompts
### What Are Prompts?
Prompts are detailed instructions given to the AI (Claude). They're like giving very specific directions to a research assistant.
### Anatomy of a Prompt
Every prompt has:
1. **Context**: The contract text to analyze
2. **Task**: What to extract
3. **Instructions**: How to extract it
4. **Examples**: Show the AI what good answers look like
5. **Constraints**: What NOT to do
6. **Output format**: How to return the answer (JSON, text, etc.)
### Example: Bottom-Up Primary Prompt
**Location**: [prompts.py:9-196](fieldExtraction/src/prompts.py#L9-L196)
```
Structure:
├── Give page text
├── Explain the task ("extract EVERY reimbursement term")
├── Define each field to extract
│ ├── FULL_SERVICE: What service is being reimbursed?
│ ├── FULL_METHODOLOGY: Complete payment description
│ └── PROV_TYPE: Professional, Facility, or Ancillary
├── Provide detailed rules
│ ├── How to handle tables
│ ├── How to capture "lesser of" statements
│ └── When to split into multiple entries
├── Give 3-5 examples of correct extractions
└── Request JSON output
```
**Key Instruction Examples**:
- "Extract EVERY percentage rate, dollar amount, or lesser of statement"
- "For tables, include: [Lesser of statements] + [Table header] + [Service row]"
- "Split entries when encountering multiple percentage rates for different services"
### Example: Methodology Breakout Prompt
**Location**: [prompts.py:198-324](fieldExtraction/src/prompts.py#L198-L324)
**Purpose**: Parse a methodology string into structured fields
```
Input: "Lesser of billed charges or 110% of Medicare Fee Schedule"
Prompt asks:
1. Is there a "lesser of"? → LESSER: Y or N
2. What's the percentage rate? → RATE_STANDARD: "110% Medicare"
3. Convert to decimal → RATE_SHORT: "1.1"
4. What's the lesser of rate? → LESSER_RATE: "100% of BC"
5. Any flat fees? → FLAT_FEE_STANDARD: N/A
6. Any "not to exceed"? → NOT_TO_EXCEED: N/A
Output:
{
"LESSER": "Y",
"RATE_STANDARD": "110% Medicare",
"RATE_SHORT": "1.1",
"FLAT_FEE_STANDARD": "N/A",
"LESSER_RATE": "100% of BC",
"NOT_TO_EXCEED": "N/A"
}
```
**Why this structure?**
- Makes rates searchable and comparable
- Standardizes different ways of expressing the same thing
- Enables analysis (e.g., "find all services paid above 100% Medicare")
---
## Preprocessing Details
### Text Cleaning
**Functions in [preprocessing_funcs.py](fieldExtraction/src/preprocessing_funcs.py)**:
1. **clean_newlines**: Fixes line break issues
2. **clean_law_symbols**: Removes § symbols and other legal markers
3. **split_text**: Divides contract into page-by-page dictionary
### Smart Chunking for AC Fields
**How it works** ([preprocessing_funcs.py](fieldExtraction/src/preprocessing_funcs.py)):
1. **Define keyword groups** ([keywords.py](fieldExtraction/src/keywords.py)):
```python
"CONTRACT_EFFECTIVE_DT": {
"keywords": ["effective date", "execution date", "date of execution"],
"case_sensitive": False
}
```
2. **Search each page** for keywords
3. **Extract matching pages** (and nearby pages for context)
4. **Build focused chunk** with only relevant text
5. **Result**: Instead of sending 50 pages, send only 3 relevant pages
**Benefits**:
- Faster processing
- Lower cost (less text to analyze)
- More accurate (AI focuses on relevant content)
### Chunking for B Fields
**Different approach** ([preprocess.py:24-50](fieldExtraction/src/preprocess.py#L24-L50)):
1. **Find exhibit pages**: Pages starting with "Exhibit", "Attachment", "Amendment"
2. **Chunk consecutive pages**: Keep pages 10-15 together if they're part of the same exhibit
3. **Align tables**: Detect tables and reformat them for consistency
**Why chunk this way?**
- Compensation schedules span multiple pages
- Need full context of each exhibit
- Tables need to stay together for accurate extraction
---
## Output Files
### AC Output
**File**: `{batch_id}-AC.csv`
**Contains one row per contract with columns like**:
- Contract Name
- Contract Title
- Payer Name
- Provider Group Name
- Contract Effective Date
- Termination Terms
- Legal Clauses
- (40+ fields total)
### B Output
**File**: `{batch_id}-B.csv`
**Contains multiple rows per contract (one per service/rate) with columns like**:
- Contract Name
- Exhibit Name
- Line of Business
- Full Service Description
- Provider Type
- Reimbursement Methodology
- Rate Standard
- Rate Short (decimal)
- Flat Fee
- Lesser Of Indicator
- (30+ fields total)
### Combined ABC Output
**File**: `{batch_id}-ABC.csv`
**Merges AC and B data**:
- Each B row includes AC contract information
- Enables analysis like "find all services in Medicare contracts with rates above 110%"
### QA/QC Report
**File**: Excel workbook with multiple sheets:
- **Blanks Analysis**: Fields with missing values
- **Date Formatting Check**: Invalid dates
- **State Validation**: Health plan state issues
- **TIN/NPI Validation**: Provider ID issues
- **Y/N Field Check**: Incorrect Yes/No values
---
## Error Handling and Recovery
### Built-in Safeguards
1. **File-level tracking** ([tracking.py](fieldExtraction/src/tracking.py)):
- Records which files are processed
- Skips already-completed files on re-run
- Logs errors with details
2. **Self-repair mode** ([main.py:170-218](fieldExtraction/src/main.py#L170-L218)):
- Automatically retries failed files
- Consolidates output after repairs
3. **Progress monitoring**:
- Real-time progress percentages
- Completion tracking for AC and B separately
- Upload every 10 files for incremental results
### What Happens When Things Go Wrong
**If a file fails**:
1. Error is logged with details
2. Processing continues with other files
3. Failed files marked in tracking
4. Self-repair attempts to reprocess
**If AI returns bad JSON**:
1. System tries to fix common JSON errors
2. Falls back to asking AI to fix it ([prompts.py:727-735](fieldExtraction/src/prompts.py#L727-L735))
3. If still broken, logs error and continues
---
## Key Design Decisions
### Why Two Separate Paths (AC and B)?
- **Different information types**: Contract metadata vs reimbursement details
- **Different extraction patterns**: One-to-one (AC) vs one-to-many (B)
- **Different prompting strategies**: Single-pass vs multi-stage
### Why Multiple Prompt Stages?
**For B fields**:
1. **Bottom-up primary**: Find everything on the page
2. **Bottom-up secondary**: Structure each finding
3. **Top-down**: Add page-level context
4. **Conditional**: Fill in gaps
**Reasoning**:
- Breaking it into stages improves accuracy
- Each stage has a focused task
- Easier to debug and improve each stage independently
### Why Use AI Instead of Rules?
**Contracts are highly variable**:
- Different formats (tables, paragraphs, bullet points)
- Different terminology (same concept, different words)
- Different structures (some have exhibits, some don't)
- Legal language that changes meaning with small variations
**AI advantages**:
- Understands context and meaning
- Adapts to variations
- Can follow complex instructions
- Generalizes to new contract types
**Where rules ARE used**:
- Date formatting (after extraction)
- TIN/NPI patterns (regex)
- Data cleanup (standardization)
---
## Performance and Scalability
### Processing Speed
**Typical performance**:
- AC fields: ~30-60 seconds per contract
- B fields: ~2-5 minutes per contract (depends on reimbursement complexity)
- With 5 workers: Can process 100 contracts in ~1-2 hours
### Parallel Processing
**How it works** ([main.py:296-334](fieldExtraction/src/main.py#L296-L334)):
```python
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
# Submit all contracts for processing
futures = [executor.submit(process_ac, contract) for contract in contracts]
# Collect results as they complete
for future in concurrent.futures.as_completed(futures):
result = future.result()
```
**Benefits**:
- Process multiple contracts simultaneously
- Better utilization of resources
- Significant time savings on large batches
### Cost Management
**Strategies**:
- Smart chunking reduces tokens sent to AI
- Caching avoids re-processing
- Progress tracking prevents duplicate work
- Cost logging tracks spend per file ([costlog_funcs.py](fieldExtraction/src/costlog_funcs.py))
---
## Troubleshooting Common Issues
### "No reimbursement information found"
**Cause**: Contract doesn't contain payment terms (might be administrative-only)
**Solution**: Normal behavior - file is skipped for B processing, AC still processed
### "AC output file not created"
**Cause**: Error during AC processing or postprocessing
**Solution**: Check error logs, review tracking file for details
### Missing fields in output
**Cause**: Information not found in contract OR extraction failed
**Solution**: Check raw contract for presence of information, review prompt effectiveness
### Rate extraction errors
**Cause**: Complex table structures, unusual formatting
**Solution**: Review table alignment step, check if manual cleanup needed
---
## Summary
This field extraction system automates the tedious work of reading healthcare contracts and pulling out critical information. By combining smart text processing, focused AI prompting, and thorough quality checks, it can process hundreds of contracts efficiently and accurately.
**The key to its success**:
1. **Intelligent preprocessing**: Only send relevant text to AI
2. **Structured prompting**: Clear instructions with examples
3. **Multi-stage extraction**: Break complex tasks into manageable steps
4. **Robust postprocessing**: Clean and validate results
5. **Error handling**: Gracefully handle failures and retry
Whether you're looking to understand how the system works, modify it for new fields, or troubleshoot issues, this guide provides the foundation you need.
fieldExtraction/documentation/PERFORMANCE_ANALYSIS_GUIDE.md
+200
View File
@@ -0,0 +1,200 @@
# Performance Analysis Guide
## Overview
This guide explains the timing instrumentation added to the pipeline and how to interpret the results to identify bottlenecks and parallelization opportunities.
## Instrumented Components
### 1. Main Pipeline (`main.py`)
- **`read_constants`**: Loading configuration and constants
- **`read_input`**: Reading input files from S3 or local storage
- **`warm_prompt_caches`**: Pre-warming Claude prompt caches
- **`parallel_file_processing`**: Main parallel processing of all files (CRITICAL METRIC)
- **`concat_results`**: Concatenating all file results into final dataframes
- **`qc_qa_validation`**: Quality control and QA validation
- **`write_outputs`**: Writing final outputs to S3/local storage
- **`parent_child_mapping`**: Parent-child relationship mapping (if enabled)
### 2. Per-File Processing (`file_processing.py`)
Each file goes through these stages (tracked with `context=filename`):
- **`{filename}.preprocess`**: Text cleaning, splitting, header/footer detection
- **`{filename}.one_to_n_exhibit_chunking`**: Chunking document into exhibits
- **`{filename}.one_to_n_extraction`**: Extracting one-to-many relationships (parallel exhibits)
- **`{filename}.one_to_n_generate_reimb_ids`**: Generating reimbursement IDs
- **`{filename}.one_to_one_extraction`**: Extracting one-to-one fields (see below)
- **`{filename}.merge_one_to_one_into_one_to_n`**: Merging results
- **`{filename}.code_processing`**: Code breakout and grouper processing
- **`{filename}.postprocess`**: Final postprocessing and validation
- **`{filename}.write_individual`**: Writing individual file results
### 3. One-to-One Processing (`file_processing.py` -> `run_one_to_one_prompts`)
- **`{filename}.one_to_one.provider_info`**: Extracting TIN/NPI provider information
- **`{filename}.one_to_one.hybrid_smart_chunking`**: RAG-based field extraction (CRITICAL)
- **`{filename}.one_to_one.full_context`**: Full-context field extraction (fallback)
### 4. Hybrid Smart Chunking (RAG) (`hybrid_smart_chunking_funcs.py`)
This is often the bottleneck. Detailed timing:
- **`{filename}.hsc.setup`**: Setting up embeddings, reranker, and text splitter
- **`{filename}.hsc.preprocessing`**: Document preprocessing and chunking
- **`{filename}.hsc.vectorstore_creation`**: Creating FAISS vectorstore + hybrid retriever
- **`{filename}.hsc.parallel_field_processing`**: Processing all fields in parallel (max 5 workers)
- **Per-field timing** (inside parallel processing):
- **`{filename}.hsc.field.{field_name}.retrieval`**: BM25 + semantic retrieval + reranking
- **`{filename}.hsc.field.{field_name}.llm_call`**: Claude API call for field extraction
## How to Analyze Results
### Step 1: Run Your Pipeline
```bash
poetry run python -m src.investment.main run_mode=local read_mode=local input_dir=test-multi write_to_s3=False state=MS-IL-NY-CA-TX-WI-AR-KY-TN-ME-OH-OR-FL-AZ-IN-NJ-MI max_workers=50 client=Centene-Parkland-Molina-Moda-Healthnet-Trillium
```
### Step 2: Review Timing Summary
At the end of execution, you'll see:
1. **Total pipeline duration**
2. **Hierarchical timing summary** organized by component
Look for the output sections:
```
================================================================================
PIPELINE COMPLETE - Total time: XXX.XXs
================================================================================
HIERARCHICAL TIMING SUMMARY
================================================================================
```
### Step 3: Identify Bottlenecks
#### Key Metrics to Check:
1. **`parallel_file_processing` total time**
- This should be the largest component
- Compare to total pipeline time to see overhead
- **If this is much shorter than pipeline time**: Look for overhead in setup, I/O, or postprocessing
2. **Per-file timings** (look at individual file contexts)
- **Find the slowest files**: Look for files with longest total time
- **Identify slow components within files**: Check which stage takes longest
- Is it `one_to_one.hybrid_smart_chunking`?
- Is it `one_to_n_extraction`?
- Is it `postprocess`?
3. **Hybrid Smart Chunking (HSC) breakdown**
- **`hsc.vectorstore_creation`**: Should be relatively quick (<5s per file)
- **`hsc.parallel_field_processing`**: Often the bottleneck
- Check if fields are processing in parallel (should see ~5 fields at a time)
- Look at per-field retrieval vs llm_call times
- **`hsc.field.{field_name}.llm_call`**: LLM latency (network + Claude processing)
- High variance? Could be prompt cache issues or rate limiting
- Consistently slow? Could be context size or model load
4. **One-to-N parallel processing**
- Check if exhibits and pages are processing in parallel
- Look for imbalanced workload (some exhibits much slower than others)
### Step 4: Optimization Opportunities
#### Scenario 1: HSC Vectorstore Creation is Slow
- **Problem**: Document preprocessing/chunking is slow
- **Solution**:
- Consider caching vectorstores per document
- Reduce chunk_size or chunk_overlap in HSC_CONFIG
- Optimize preprocessing pipeline
#### Scenario 2: HSC LLM Calls are Slow
- **Problem**: Too many sequential LLM calls or high latency
- **Current**: max_workers=5 for field processing
- **Solution**:
- Increase `max_workers` in `hybrid_smart_chunking_funcs.py:304` (currently 5)
- Check Bedrock connection pool (now set to 60, should be sufficient)
- Verify prompt caching is working (check cache hits in logs)
#### Scenario 3: HSC Retrieval/Reranking is Slow
- **Problem**: RAG retrieval taking too long per field
- **Solution**:
- Reduce `ENSEMBLE_TOP_K` or `RERANKER_TOP_K` in RAG config
- Consider using faster reranker model
- Optimize BM25 index size
#### Scenario 4: Parallel File Processing Not Saturated
- **Problem**: Not all `max_workers` are being used
- **Check**:
- Connection pool exhaustion (should be fixed with `max_pool_connections=60`)
- Memory constraints (monitor memory usage)
- I/O bottlenecks (local disk or S3 bandwidth)
- **Solution**:
- Increase `max_workers` carefully (currently 50)
- Profile memory usage to find leaks
#### Scenario 5: One-to-N Exhibits Not Parallel
- **Current**: max_workers=5 for exhibit processing, max_workers=5 for page processing
- **Solution**:
- Increase these limits in `file_processing.py:244` (exhibits) and `:192` (pages)
- Balance against total file-level parallelism to avoid oversubscription
#### Scenario 6: Postprocessing is Slow
- **Problem**: Code processing or postprocessing taking too long
- **Solution**:
- Profile specific postprocessing functions
- Look for opportunities to batch operations
- Consider caching expensive lookups
## Current Parallelization Strategy
### Three Levels of Parallelism:
1. **File Level** (main.py)
- `max_workers=50` (from command line arg)
- 50 files processed simultaneously
2. **Exhibit/Page Level** (within each file)
- Exhibits: `max_workers=5` per file
- Pages: `max_workers=5` per exhibit
- This means potentially 50 files × 5 exhibits × 5 pages = 1,250 threads (too many!)
- **Currently limited by nested executors**
3. **Field Level** (within HSC per file)
- `max_workers=5` for field extraction
- 5 fields extracted in parallel per file
### Total Theoretical Concurrency:
- **Bedrock API calls**: Up to 60 concurrent (connection pool limit)
- **Thread pools**: Nested executors create complex concurrency
- **Bottleneck**: Likely the 60-connection Bedrock pool with 50 file workers
## Recommendations After First Run
1. **Run with timing enabled** and review the hierarchical summary
2. **Identify the top 3 slowest components** by total time
3. **Check if parallelism is effective**:
- Are all max_workers being utilized?
- Is there thread contention or blocking?
4. **Look for outliers**: Files that take much longer than others
5. **Profile memory usage**: Ensure no memory leaks with many threads
6. **Monitor Bedrock throttling**: Check for rate limit errors in logs
## Next Steps
Based on timing results, we can:
1. **Increase parallelism** where bottlenecks are found
2. **Add caching** for repeated operations (vectorstores, embeddings)
3. **Optimize prompts** to reduce token count and latency
4. **Batch operations** where sequential processing is required
5. **Add more granular timing** to specific slow functions
## Timing Log Format
Timing logs appear as:
```
INFO - ⏱️ read_constants: 0.50s
INFO - ⏱️ parallel_file_processing: 450.23s
DEBUG - ⏱️ file1.pdf.preprocess: 2.34s
DEBUG - ⏱️ file1.pdf.one_to_one.hybrid_smart_chunking: 45.67s
DEBUG - ⏱️ file1.pdf.hsc.setup: 0.12s
DEBUG - ⏱️ file1.pdf.hsc.vectorstore_creation: 3.45s
DEBUG - ⏱️ file1.pdf.hsc.parallel_field_processing: 40.12s
DEBUG - ⏱️ file1.pdf.hsc.field.EFFECTIVE_DT.retrieval: 0.89s
DEBUG - ⏱️ file1.pdf.hsc.field.EFFECTIVE_DT.llm_call: 2.34s
```
Use these to drill down into specific bottlenecks!
fieldExtraction/src/codes/code_funcs.py
+47 -21
View File
@@ -1,6 +1,7 @@
import ast
import json
import logging
import concurrent.futures
import os
import re
@@ -720,15 +721,20 @@ def code_breakout(merged_results: pd.DataFrame, constants: Constants):
Returns:
pd.DataFrame: DataFrame with updated code fields.
"""
import concurrent.futures
# Fill empty AARETE_DERIVED_CLAIM_TYPE_CD with mode
answer_dicts = fill_claim_type(merged_results.to_dict(orient="records"))
answer_dicts_with_code = []
for answer_dict in answer_dicts:
# Parallelize code extraction
def process_single_code(answer_dict):
code_answer_dict = extract_codes_from_service(answer_dict, constants)
answer_dict.update(code_answer_dict)
answer_dicts_with_code.append(answer_dict)
return answer_dict
max_workers = min(len(answer_dicts), 10)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
answer_dicts_with_code = list(executor.map(process_single_code, answer_dicts))
# Fill in GROUPER_CD_DESC based on GROUPER_TYPE and GROUPER_CD
answer_dicts_with_code = [
@@ -754,32 +760,52 @@ def grouper_breakout(results_with_code: pd.DataFrame):
Returns:
pd.DataFrame: DataFrame with updated grouper fields.
"""
# try grouper fields again if they are empty and grouper_cd is not empty
answer_dicts_with_code = results_with_code.to_dict(orient="records")
final_answer_dicts = []
GROUPER_QUESTIONS = FieldSet(
file_path=config.FIELD_JSON_PATH, field_type="grouper_breakout"
).print_prompt_dict()
# Separate dicts that need grouper breakout from those that don't
needs_breakout = []
no_breakout_needed = []
for answer_dict in answer_dicts_with_code:
if string_utils.is_empty(
answer_dict.get("GROUPER_TYPE")
) and not string_utils.is_empty(answer_dict.get("GROUPER_CD")):
# run groper breakout prompt
grouper_breakout_prompt = prompt_templates.GROUPER_BREAKOUT(
answer_dict.get("SERVICE_TERM", ""),
answer_dict.get("REIMB_TERM", ""),
GROUPER_QUESTIONS,
)
claude_answer_raw = llm_utils.invoke_claude(
grouper_breakout_prompt, "sonnet_latest", ""
)
try:
grouper_answer = string_utils.universal_json_load(claude_answer_raw)
except Exception as e:
grouper_answer = {"GROUPER_TYPE": e}
needs_breakout.append(answer_dict)
else:
no_breakout_needed.append(answer_dict)
answer_dict.update(grouper_answer)
final_answer_dicts.append(answer_dict)
logging.debug(f"Grouper breakout: {len(needs_breakout)} rows need LLM processing, {len(no_breakout_needed)} rows skip")
return pd.DataFrame(final_answer_dicts)
# Parallelize LLM calls for rows that need breakout
def process_grouper_breakout(answer_dict):
grouper_breakout_prompt = prompt_templates.GROUPER_BREAKOUT(
answer_dict.get("SERVICE_TERM", ""),
answer_dict.get("REIMB_TERM", ""),
GROUPER_QUESTIONS,
)
claude_answer_raw = llm_utils.invoke_claude(
grouper_breakout_prompt, "sonnet_latest", ""
)
try:
grouper_answer = string_utils.universal_json_load(claude_answer_raw)
except Exception as e:
grouper_answer = {"GROUPER_TYPE": e}
answer_dict.update(grouper_answer)
return answer_dict
if needs_breakout:
max_workers = min(len(needs_breakout), 10)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
processed_breakout = list(executor.map(process_grouper_breakout, needs_breakout))
else:
processed_breakout = []
# Combine all results
final_answer_dicts = no_breakout_needed + processed_breakout
return pd.DataFrame(final_answer_dicts)
fieldExtraction/src/config.py
+1 -1
View File
@@ -193,7 +193,7 @@ if RUN_MODE == "local":
AWS_SESSION_TOKEN = get_value_from_env("AWS_SESSION_TOKEN")
# Bedrock
config = Config(read_timeout=2000)
config = Config(read_timeout=2000, max_pool_connections=60)
BEDROCK_RUNTIME = boto3.client(
service_name="bedrock-runtime",
region_name="us-east-1",
fieldExtraction/src/investment/dynamic_funcs.py
+34 -19
View File
@@ -1,12 +1,13 @@
import logging
import concurrent.futures
import src.investment.prompt_calls as prompt_calls
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
from src.investment import prompt_calls, preprocessing_funcs
from src.utils import string_utils
from constants.constants import Constants
from src import config
from src.prompts import prompt_templates
from src.prompts.fieldset import Field, FieldSet
from src.investment.exhibit_funcs import Exhibit
def dynamic_primary(
@@ -140,24 +141,38 @@ def dynamic(
return exhibit_level_answers, dynamic_reimbursement_fields
def dynamic_assignment(reimbursement_primary_answers: list[dict[str, str]], dynamic_reimbursement_fields: FieldSet, exhibit_text_simplified: str, page_num: str, constants: Constants, filename: str) -> list[dict[str, str]]:
# Process each answer and field individually
answer_dicts = []
def dynamic_assignment(reimbursement_primary_answers: list[dict[str, str]],
dynamic_reimbursement_fields: FieldSet,
text_dict: dict[str, str],
exhibit: Exhibit,
constants: Constants,
filename: str
) -> list[dict[str, str]]:
"""
Enrich reimbursement rows with dynamic fields in parallel.
"""
if not reimbursement_primary_answers or not dynamic_reimbursement_fields.fields:
return reimbursement_primary_answers
# TODO: Cache here
# instruction, prompt = prompt_templates.DYNAMIC_ASSIGNMENT
# TODO: Multithread this loop
for answer_dict in reimbursement_primary_answers:
service_term = answer_dict['SERVICE_TERM']
reimb_term = answer_dict['REIMB_TERM']
def _assign_single(answer_dict: dict[str, str]) -> dict[str, str]:
# copy to avoid mutating caller-provided dicts when running in threads
updated = answer_dict.copy()
page_num = updated["REIMB_PAGE"]
exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(
text_dict, exhibit.exhibit_page_nums, page_num
)
service_term = updated['SERVICE_TERM']
reimb_term = updated['REIMB_TERM']
for dynamic_field in dynamic_reimbursement_fields.fields:
dynamic_field_answer = prompt_calls.prompt_dynamic_assignment(service_term, reimb_term, dynamic_field, exhibit_text_simplified, page_num, constants, filename)
answer_dict.update(dynamic_field_answer)
answer_dicts.append(answer_dict)
return answer_dicts
dynamic_field_answer = prompt_calls.prompt_dynamic_assignment(
service_term, reimb_term, dynamic_field, exhibit_text_simplified, page_num, constants, filename
)
updated.update(dynamic_field_answer)
return updated
max_workers = min(len(reimbursement_primary_answers), 8)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
return list(executor.map(_assign_single, reimbursement_primary_answers))
def add_one_to_one_field(one_to_one_fields, field_to_add, answer_dicts, constants: Constants):
fieldExtraction/src/investment/exhibit_funcs.py
+143
View File
@@ -0,0 +1,143 @@
"""
Exhibit class for tracking exhibit state during one-to-n processing.
This module provides the Exhibit class which maintains all state and field values
for a single exhibit during field extraction, including a linked list structure
for propagating dynamic fields from exhibit n-1 to exhibit n.
"""
from typing import Optional
from src.prompts.fieldset import FieldSet
class Exhibit:
"""
Tracks all state and field values for a single exhibit during processing.
Maintains a linked list structure via prev_exhibit to enable field propagation
from exhibit n-1 to exhibit n during dynamic assignment.
Attributes:
exhibit_page (str): The page number where this exhibit starts
exhibit_page_nums (list[str]): All page numbers that belong to this exhibit
exhibit_header (str): The header text of the exhibit
exhibit_text (str): The full text content of the exhibit
num_pages (int): Number of pages in this exhibit
prev_exhibit (Optional[Exhibit]): Pointer to the previous exhibit (n-1)
dynamic_primary_fields (FieldSet): Dynamic primary fields from exhibit_level
dynamic_reimbursement_fields (FieldSet): Dynamic reimbursement fields from exhibit_level
exhibit_level_answers (dict): Exhibit-level field answers
reimbursement_rows (list[dict]): Reimbursement rows from page processing
special_case_rows (list[dict]): Special case rows from page processing
has_reimbursements (bool): Whether this exhibit has any reimbursements
final_rows (list[dict]): Final combined and cleaned rows
"""
def __init__(self, exhibit_page: str, exhibit_page_nums: list[str],
exhibit_header: str, exhibit_text: str, prev_exhibit: Optional['Exhibit'] = None):
"""
Initialize a new Exhibit instance.
Args:
exhibit_page: The starting page number for this exhibit
exhibit_page_nums: List of all page numbers in this exhibit
exhibit_header: The header text of the exhibit
exhibit_text: The full text content of the exhibit
prev_exhibit: Optional pointer to the previous exhibit for field propagation
"""
# Basic exhibit metadata
self.exhibit_page = exhibit_page
self.exhibit_page_nums = exhibit_page_nums
self.exhibit_header = exhibit_header
self.exhibit_text = exhibit_text
self.num_pages = len(exhibit_page_nums)
# Linked list for exhibit chain
self.prev_exhibit = prev_exhibit
# Field tracking (populated in Step 2)
self.dynamic_primary_fields = FieldSet()
self.dynamic_reimbursement_fields = FieldSet()
self.exhibit_level_answers = {}
# Reimbursement data (populated in Step 1)
self.reimbursement_rows = []
self.special_case_rows = []
self.has_reimbursements = False
# Final combined rows (populated in Step 3)
self.final_rows = []
def add_reimbursement_rows(self, reimbursement_rows: list[dict], special_case_rows: list[dict]):
"""
Add reimbursement rows from page processing.
This is called during Step 1 (parallel page processing) to accumulate
reimbursement and special case rows from all pages in this exhibit.
Args:
reimbursement_rows: List of reimbursement dictionaries from a page
special_case_rows: List of special case dictionaries from a page
"""
self.reimbursement_rows.extend(reimbursement_rows)
self.special_case_rows.extend(special_case_rows)
self.has_reimbursements = len(reimbursement_rows) > 0 or len(special_case_rows) > 0
def set_exhibit_level_data(self, exhibit_level_answers: dict, dynamic_reimbursement_fields: FieldSet):
"""
Set exhibit-level field data from exhibit_level processing.
This is called during Step 2 (exhibit-level processing) to store the
exhibit-level answers and dynamic fields that will be used in Step 3.
Args:
exhibit_level_answers: Dictionary of exhibit-level field answers
dynamic_reimbursement_fields: FieldSet of dynamic reimbursement fields
"""
self.exhibit_level_answers = exhibit_level_answers
self.dynamic_reimbursement_fields = dynamic_reimbursement_fields
# Extract dynamic primary fields from exhibit level answers
# These are the fields that can be propagated to next exhibit
self.dynamic_primary_fields = dynamic_reimbursement_fields
def get_previous_exhibit_dynamic_fields(self) -> FieldSet:
"""
Get dynamic primary fields from previous exhibit (n-1) if available.
This enables field propagation across exhibits, allowing exhibit n to
inherit dynamic fields from exhibit n-1 when appropriate.
Returns:
FieldSet: The previous exhibit's dynamic primary fields, or an empty
FieldSet if there is no previous exhibit
"""
if self.prev_exhibit is not None:
return self.prev_exhibit.dynamic_primary_fields
return FieldSet()
def __repr__(self):
"""String representation for debugging."""
return (f"Exhibit(page={self.exhibit_page}, pages={len(self.exhibit_page_nums)}, "
f"has_reimb={self.has_reimbursements}, reimb_count={len(self.reimbursement_rows)})")
def get_exhibit_list(text_dict: dict[str, str], exhibit_chunk_mapping: dict[str, str], all_exhibit_headers: dict[str, str]) -> list[Exhibit]:
exhibits = []
prev_exhibit = None
for exhibit_page, exhibit_page_nums in exhibit_chunk_mapping.items():
exhibit_text = "\n".join([text_dict[page_num] for page_num in exhibit_page_nums])
exhibit_header = all_exhibit_headers.get(exhibit_page, "")
exhibit = Exhibit(
exhibit_page=exhibit_page,
exhibit_page_nums=exhibit_page_nums,
exhibit_header=exhibit_header,
exhibit_text=exhibit_text,
prev_exhibit=prev_exhibit
)
exhibits.append(exhibit)
prev_exhibit = exhibit
return exhibits
fieldExtraction/src/investment/file_processing.py
+214 -179
View File
@@ -1,12 +1,14 @@
import concurrent.futures
import logging
import pandas as pd
import src.codes.code_funcs as code_funcs
from src.investment import preprocessing_funcs, aarete_derived, dynamic_funcs, one_to_n_funcs, one_to_one_funcs, postprocess, postprocessing_funcs, preprocess, hybrid_smart_chunking_funcs, row_funcs, tin_npi_funcs
from src.utils import io_utils, logging_utils, string_utils
from src.investment import preprocessing_funcs, aarete_derived, dynamic_funcs, one_to_n_funcs, one_to_one_funcs, postprocess, postprocessing_funcs, preprocess, hybrid_smart_chunking_funcs, row_funcs, tin_npi_funcs, exhibit_funcs
from src.investment.exhibit_funcs import Exhibit
from src.utils import io_utils, logging_utils, string_utils, timing_utils
from constants.constants import Constants
from src import config
from src.prompts.fieldset import FieldSet, Field
from src.prompts.fieldset import FieldSet
from src.utils.string_utils import datetime_str
@@ -27,32 +29,37 @@ def process_file(file_object, constants: Constants, run_timestamp):
final_results = pd.DataFrame([{"FILE_NAME": filename}])
################## PREPROCESS ##################
contract_text = preprocess.clean_text(contract_text)
text_dict, top_sheet_dict = preprocess.split_text(contract_text)
text_dict, header, footer = preprocess.find_headers_and_footers(text_dict)
with timing_utils.timed_block("preprocess", context=filename):
contract_text = preprocess.clean_text(contract_text)
text_dict, top_sheet_dict = preprocess.split_text(contract_text)
text_dict, header, footer = preprocess.find_headers_and_footers(text_dict)
# ONE TO N PROCESSING
one_to_n_results = pd.DataFrame() # Initialize empty DataFrame for one_to_n_results
if process_one_to_n:
exhibit_chunk_mapping, all_exhibit_headers = preprocess.one_to_n_exhibit_chunking(
text_dict, constants.EXHIBIT_HEADER_MARKERS, filename
)
logging.info(f"{datetime_str()} Preprocessing Complete - {filename}")
one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}]) # Initialize here
if string_utils.contains_reimbursement(contract_text):
one_to_n_results, dynamic_one_to_one_fields, first_reimbursement_page = run_one_to_n_prompts(
text_dict, exhibit_chunk_mapping, all_exhibit_headers, constants, filename
with timing_utils.timed_block("one_to_n_exhibit_chunking", context=filename):
exhibit_chunk_mapping, all_exhibit_headers = preprocess.one_to_n_exhibit_chunking(
text_dict, constants.EXHIBIT_HEADER_MARKERS, filename
)
logging.info(f"{datetime_str()} Preprocessing Complete - {filename}")
one_to_n_results = pd.DataFrame([{"FILE_NAME": filename}]) # Initialize here
if string_utils.contains_reimbursement(contract_text):
logging.info(f"{datetime_str()} Starting One-to-N extraction for {len(exhibit_chunk_mapping)} exhibits - {filename}")
with timing_utils.timed_block("one_to_n_extraction", context=filename):
one_to_n_results, dynamic_one_to_one_fields, first_reimbursement_page = run_one_to_n_prompts(
text_dict, exhibit_chunk_mapping, all_exhibit_headers, constants, filename
)
if not one_to_n_results.empty:
one_to_n_results["FILE_NAME"] = filename
one_to_n_results = postprocessing_funcs.generate_reimb_ids(one_to_n_results)
with timing_utils.timed_block("one_to_n_generate_reimb_ids", context=filename):
one_to_n_results = postprocessing_funcs.generate_reimb_ids(one_to_n_results)
logging.info(f"{datetime_str()} One to N Complete - {filename}")
else:
first_reimbursement_page = '1'
logging.info(f"{datetime_str()} No Reimbursement Found, Skipping - {filename}")
final_results = one_to_n_results # Set as final results
else:
first_reimbursement_page = '1'
@@ -60,46 +67,51 @@ def process_file(file_object, constants: Constants, run_timestamp):
# ONE TO ONE PROCESSING
if process_one_to_one:
one_to_one_results = run_one_to_one_prompts(
filename,
contract_text,
text_dict,
top_sheet_dict,
dynamic_one_to_one_fields,
first_reimbursement_page,
constants,
)
with timing_utils.timed_block("one_to_one_extraction", context=filename):
one_to_one_results = run_one_to_one_prompts(
filename,
contract_text,
text_dict,
top_sheet_dict,
dynamic_one_to_one_fields,
first_reimbursement_page,
constants,
)
one_to_one_results["FILE_NAME"] = filename
logging.info(f"{datetime_str()} One to One Complete - {filename}")
# Decide how to handle one_to_one results
if not one_to_n_results.empty:
# BOTH processed - merge into one_to_n
final_results = row_funcs.merge_one_to_one_into_one_to_n(
one_to_n_results, one_to_one_results, constants
)
else:
# ONLY one_to_one processed - convert dict to DataFrame
final_results = pd.DataFrame([one_to_one_results])
with timing_utils.timed_block("merge_one_to_one_into_one_to_n", context=filename):
if not one_to_n_results.empty:
# BOTH processed - merge into one_to_n
final_results = row_funcs.merge_one_to_one_into_one_to_n(
one_to_n_results, one_to_one_results, constants
)
else:
# ONLY one_to_one processed - convert dict to DataFrame
final_results = pd.DataFrame([one_to_one_results])
else:
logging.info(f"{datetime_str()} Fields not configured for One to One, Skipping - {filename}")
# APPLY CODES IF ONE_TO_N WAS PROCESSED
if not one_to_n_results.empty:
results_with_code = code_funcs.code_breakout(final_results, constants)
final_results = code_funcs.grouper_breakout(results_with_code)
with timing_utils.timed_block("code_processing", context=filename):
results_with_code = code_funcs.code_breakout(final_results, constants)
final_results = code_funcs.grouper_breakout(results_with_code)
logging.info(f"{datetime_str()} Codes Complete - {filename}")
# POSTPROCESS
final_df = postprocess.postprocess(final_results, constants)
with timing_utils.timed_block("postprocess", context=filename):
final_df = postprocess.postprocess(final_results, constants)
logging.info(f"{datetime_str()} Postprocessing Complete - {filename}")
################## WRITE INDIVIDUAL ##################
if config.WRITE_TO_S3:
io_utils.write_s3(final_df, filename, run_timestamp, "individual")
else:
io_utils.write_local(final_df, filename, "", "individual")
with timing_utils.timed_block("write_individual", context=filename):
if config.WRITE_TO_S3:
io_utils.write_s3(final_df, filename, run_timestamp, "individual")
else:
io_utils.write_local(final_df, filename, "", "individual")
logging.info(f"{datetime_str()} Writing Complete - {filename}")
@@ -121,22 +133,25 @@ def run_one_to_one_prompts(
relationship="one_to_one", file_path=config.FIELD_JSON_PATH
).combine(dynamic_one_to_one_fields)
################## RUN PROVIDER INFO ##################
with timing_utils.timed_block("one_to_one.provider_info", context=filename):
one_to_one_results = tin_npi_funcs.run_provider_info_fields(
contract_text, text_dict, filename, payer_name=""
)
################## 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
)
with timing_utils.timed_block("one_to_one.hybrid_smart_chunking", context=filename):
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 PROVIDER INFO ##################
payer_name = hybrid_smart_chunked_answers_dict.get('PAYER_NAME')
one_to_one_results = tin_npi_funcs.run_provider_info_fields(
contract_text, text_dict, filename, payer_name
)
################## RUN FULL CONTEXT PROMPTS ##################
full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
one_to_one_fields, contract_text, text_dict, first_reimbursement_page, constants, filename
)
with timing_utils.timed_block("one_to_one.full_context", context=filename):
full_context_answers_dict = one_to_one_funcs.run_full_context_fields(
one_to_one_fields, contract_text, text_dict, first_reimbursement_page, constants, filename
)
################## COMBINE ANSWERS ################
# Prefer HSC answers over full_context when HSC value is not N/A
for key, value in full_context_answers_dict.items():
@@ -168,111 +183,12 @@ def run_one_to_one_prompts(
return one_to_one_results[0]
def run_one_to_n_prompts(text_dict: dict[str, str],
exhibit_chunk_mapping: dict[str, list[str]],
all_exhibit_headers: dict[str, str],
constants: Constants,
filename: str):
one_to_n_results = []
first_reimbursement_page = '1'
previous_exhibit = None
#################################### PROCESS EACH EXHIBIT SEPARATELY ####################################
for exhibit_page, exhibit_page_nums in exhibit_chunk_mapping.items():
exhibit_text_original = "\n".join([text_dict[page_num] for page_num in exhibit_page_nums])
exhibit_header = all_exhibit_headers.get(exhibit_page, "")
############################### Get Exhibit Level ###############################
exhibit_level_answers, dynamic_reimbursement_fields = (
one_to_n_funcs.exhibit_level(
exhibit_text_original,
exhibit_header,
exhibit_page,
constants,
filename,
)
) # dict, FieldSet
logging.debug(f"Exhibit Level Answers for {filename}, Page {exhibit_page}: {exhibit_level_answers}")
# Store original values before inheritance (for tracking in previous_exhibit)
original_exhibit_page_nums = exhibit_page_nums.copy()
original_exhibit_text = exhibit_text_original
original_dynamic_reimbursement_fields = dynamic_reimbursement_fields
############################### Check and Combine Exhibit Inheritance ###############################
# IF dynamic_reimbursement_fields is empty FieldSet, and previous exhibit's dynamic_reimbursement_fields is NOT empty,
# and the previous exhibit has NO reimbursement rows, then replace dynamic_reimbursement_fields with the previous
# exhibit's dynamic_reimbursement_field values AND set exhibit_page_nums = exhibit_page_nums for exhibit N-1 + exhibit_page_nums for exhibit N
should_inherit, exhibit_page_nums, dynamic_reimbursement_fields = (
one_to_n_funcs.check_and_combine_exhibit_inheritance(
previous_exhibit,
exhibit_page_nums,
exhibit_text_original,
dynamic_reimbursement_fields,
)
)
all_exhibit_reimbursements, all_exhibit_special_case = [], []
#################################### PAGE-BY-PAGE REIMBURSEMENT PRIMARY ####################################
for page_num in exhibit_page_nums:
reimbursement_level_answers, special_case_answers, first_reimbursement_page = process_page_one_to_n(
text_dict,
exhibit_page_nums,
exhibit_page,
page_num,
first_reimbursement_page,
dynamic_reimbursement_fields,
filename,
constants
)
all_exhibit_reimbursements += reimbursement_level_answers
all_exhibit_special_case += special_case_answers
################################ Combine Answers ###############################
all_exhibit_rows = row_funcs.combine_one_to_n(
exhibit_text_original,
all_exhibit_reimbursements,
all_exhibit_special_case,
exhibit_level_answers,
filename
) # returns list of dicts
################################ Mapping and Cleaning ###############################
all_exhibit_rows = one_to_n_funcs.one_to_n_cleaning(all_exhibit_rows, exhibit_text_original, constants, filename)
################################ Add to Total ###############################
one_to_n_results += all_exhibit_rows
################################ Store Current Exhibit for Next Iteration ###############################
# Store original (non-combined) values for next exhibit to potentially inherit from
# Determine if current exhibit had reimbursements
had_reimbursements = len(all_exhibit_reimbursements) > 0 or len(all_exhibit_special_case) > 0
previous_exhibit = {
'exhibit_page': exhibit_page,
'exhibit_page_nums': original_exhibit_page_nums,
'exhibit_text_original': original_exhibit_text,
'dynamic_reimbursement_fields': original_dynamic_reimbursement_fields,
'exhibit_level_answers': exhibit_level_answers.copy(),
'had_reimbursements': had_reimbursements
}
################## Add N/A Dynamic or Exhibit to One-to-One ##################
dynamic_one_to_one_fields = dynamic_funcs.get_dynamic_one_to_one_fields(
one_to_n_results, constants
)
################## CONVERT TO DF ##################
one_to_n_df = pd.DataFrame(one_to_n_results)
return one_to_n_df, dynamic_one_to_one_fields, first_reimbursement_page
def process_page_one_to_n(text_dict: dict[str, str], exhibit_page_nums: list[str], exhibit_page: str, page_num: str, first_reimbursement_page: str, dynamic_reimbursement_fields: FieldSet, filename: str, constants: Constants):
def process_page_reimbursements(exhibit: Exhibit, page_num: str, text_dict: dict[str, str],
exhibit_page_nums: list[str], constants: Constants, filename: str):
"""
STEP 1 HELPER: Extract reimbursements from a single page (without exhibit-level fields).
Runs: reimbursement_level → carveout_and_special_case → lesser_of_distribution → breakout
"""
page_text = text_dict[page_num]
exhibit_text_simplified = preprocessing_funcs.simplify_exhibit(text_dict, exhibit_page_nums, page_num)
@@ -282,34 +198,153 @@ def process_page_one_to_n(text_dict: dict[str, str], exhibit_page_nums: list[str
constants,
filename
)
if not reimbursement_level_answers:
return [], [], first_reimbursement_page
# Track the first page with reimbursements
if first_reimbursement_page == '1':
first_reimbursement_page = exhibit_page
logging.debug(f"Reimbursement Primary Answers for {filename}, Page {page_num}: {reimbursement_level_answers}")
return [], []
################################ Carveouts and Special Case ################################
reimbursement_level_answers, special_case_answers = one_to_n_funcs.carveout_and_special_case(
reimbursement_level_answers, constants, filename
) # Breaks the reimbursement answers into reimbursments (regular lines) and special cases (distributed across exhibit)
############################### Dynamic Assignment ###############################
reimbursement_level_answers = dynamic_funcs.dynamic_assignment(reimbursement_level_answers, dynamic_reimbursement_fields, exhibit_text_simplified, page_num, constants, filename)
)
############################### Lesser of Distribution ###############################
reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution(reimbursement_level_answers, exhibit_text_simplified, page_num, constants, filename)
# Note: We run lesser_of without dynamic fields in Step 1
reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution(
reimbursement_level_answers, exhibit_text_simplified, page_num, constants, filename
)
################################ Get Breakouts ###############################
reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout(
reimbursement_level_answers,
special_case_answers,
filename,
constants,
) # list[dict[str, str]], list[dict[str, str]]
)
################################ Assign REIMB_PAGE ###############################
for answer_dict in reimbursement_level_answers:
answer_dict["REIMB_PAGE"] = page_num
return reimbursement_level_answers, special_case_answers
def run_one_to_n_prompts(text_dict: dict[str, str],
exhibit_chunk_mapping: dict[str, list[str]],
all_exhibit_headers: dict[str, str],
constants: Constants,
filename: str):
"""
3-STEP APPROACH (per exhibit):
Step 1: Extract reimbursements for exhibit pages (parallel within exhibit)
Step 2: Run exhibit_level when exhibit has reimbursements
Step 3: Run dynamic_assignment and combine results
"""
total_exhibits = len(exhibit_chunk_mapping)
logging.debug(f"{datetime_str()} Processing {total_exhibits} exhibits with NEW 3-step approach - {filename}")
# Create Exhibit objects in order (maintains exhibit chain via prev_exhibit)
exhibits = exhibit_funcs.get_exhibit_list(
text_dict,
exhibit_chunk_mapping,
all_exhibit_headers
) # <- Returns list[Exhibit]
return reimbursement_level_answers, special_case_answers, first_reimbursement_page
# Process exhibits serially; within each exhibit, process pages in parallel
total_pages = sum(len(exhibit.exhibit_page_nums) for exhibit in exhibits)
completed_pages = 0
one_to_n_results = []
first_reimbursement_page = '1'
for exhibit in exhibits:
exhibit_pages = exhibit.exhibit_page_nums
if exhibit_pages:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(exhibit_pages), 20)) as executor:
page_futures = {
executor.submit(
process_page_reimbursements,
exhibit, page_num, text_dict, exhibit_pages, constants, filename
): page_num
for page_num in exhibit_pages
}
for future in concurrent.futures.as_completed(page_futures):
try:
page_num = page_futures[future]
reimbursement_rows, special_case_rows = future.result()
exhibit.add_reimbursement_rows(reimbursement_rows, special_case_rows)
completed_pages += 1
if completed_pages % 20 == 0 or completed_pages == total_pages:
logging.debug(f"{datetime_str()} Page progress: {completed_pages}/{total_pages} - {filename}")
except Exception as e:
logging.error(f"Error processing page: {str(e)}")
if not exhibit.has_reimbursements:
continue
# STEP 2: exhibit_level for this exhibit
exhibit_level_answers, dynamic_reimbursement_fields = one_to_n_funcs.exhibit_level(
exhibit.exhibit_text,
exhibit.exhibit_header,
exhibit.exhibit_page,
constants,
filename,
)
exhibit.set_exhibit_level_data(exhibit_level_answers, dynamic_reimbursement_fields)
logging.debug(f"Exhibit Level Answers for {filename}, Page {exhibit.exhibit_page}: {exhibit_level_answers}")
# STEP 3: dynamic assignment & combine for this exhibit
# Get dynamic fields from previous exhibit if needed (for future use)
if exhibit.dynamic_reimbursement_fields:
dynamic_fields = exhibit.dynamic_reimbursement_fields
elif exhibit.prev_exhibit and not exhibit.prev_exhibit.has_reimbursements and exhibit.get_previous_exhibit_dynamic_fields().fields:
dynamic_fields = exhibit.get_previous_exhibit_dynamic_fields()
else:
dynamic_fields = None
# Run dynamic assignment for ALL reimbursement rows in this exhibit
# Note: dynamic_assignment expects a list of rows and returns a list of rows
reimbursement_rows_with_dynamic = exhibit.reimbursement_rows
if exhibit.reimbursement_rows and dynamic_fields is not None:
reimbursement_rows_with_dynamic = dynamic_funcs.dynamic_assignment(
exhibit.reimbursement_rows,
dynamic_fields,
text_dict,
exhibit,
constants,
filename
)
# Combine reimbursement rows with exhibit-level answers
combined_rows = row_funcs.combine_one_to_n(
exhibit.exhibit_text,
reimbursement_rows_with_dynamic,
exhibit.special_case_rows,
exhibit.exhibit_level_answers,
filename
)
# Run cleaning
combined_rows = one_to_n_funcs.one_to_n_cleaning(
combined_rows, exhibit.exhibit_text, constants, filename
)
exhibit.final_rows = combined_rows
one_to_n_results.extend(combined_rows)
# Track first reimbursement page
if first_reimbursement_page == '1' and exhibit.has_reimbursements:
first_reimbursement_page = exhibit.exhibit_page
logging.debug(f"{datetime_str()} STEP 3 COMPLETE: Combined {len(one_to_n_results)} total rows - {filename}")
################## Add N/A Dynamic or Exhibit to One-to-One ##################
dynamic_one_to_one_fields = dynamic_funcs.get_dynamic_one_to_one_fields(
one_to_n_results, constants
)
################## CONVERT TO DF ##################
one_to_n_df = pd.DataFrame(one_to_n_results)
return one_to_n_df, dynamic_one_to_one_fields, first_reimbursement_page
fieldExtraction/src/investment/hybrid_smart_chunking_funcs.py
+112 -96
View File
@@ -5,6 +5,7 @@ This module provides RAG (Retrieval-Augmented Generation) based field extraction
for contract processing using hybrid retrieval (BM25 + semantic search) with reranking.
"""
import concurrent.futures
import gc
import logging
import re
@@ -15,6 +16,7 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
import src.utils.timing_utils as timing_utils
import src.investment.preprocessing_funcs as preprocessing_funcs
import src.investment.hybrid_smart_chunking_preprocessing as hybrid_smart_chunking_preprocessing
from constants.constants import Constants
@@ -35,6 +37,84 @@ from src.utils.rag_utils import (
)
def process_single_field(field, retriever, reranker, rag_config, text_dict, constants, filename):
"""Process a single field in parallel for HSC."""
try:
retrieval_query = field.retrieval_question
if not retrieval_query:
return (field.field_name, "N/A", field)
# Retrieve + rerank using retrieval question
with timing_utils.timed_block(f"hsc.field.{field.field_name}.retrieval", context=filename):
retrieved = retriever.invoke(retrieval_query)[: rag_config.ENSEMBLE_TOP_K]
reranked = rerank_documents(
retrieval_query,
retrieved,
reranker,
rag_config.RERANKER_TOP_K,
rag_config.RERANKER,
)
# Build context and use actual LLM prompt for extraction
context = format_chunks_for_llm(text_dict, reranked)
if len(context) == 0:
# Fallback: mark field for full_context processing
field.field_type = "full_context"
logging.warning(
f"No context retrieved for {field.field_name}, marking as full_context"
)
return (field.field_name, "N/A", field)
llm_prompt = field.get_prompt_dict(constants) # Returns dict of {field_name : prompt}
prompt = prompt_templates.ONE_TO_ONE_SINGLE_FIELD_TEMPLATE(context, llm_prompt)
logging.debug(f"Field: {field.field_name}")
logging.debug(f"Retrieval question: {retrieval_query}")
logging.debug(f"LLM prompt: {llm_prompt}")
# LLM call
with timing_utils.timed_block(f"hsc.field.{field.field_name}.llm_call", context=filename):
raw = llm_utils.invoke_claude(
prompt, "sonnet_latest", filename, max_tokens=8192
)
try:
parsed = string_utils.universal_json_load(raw) # returns dict
val = parsed.get(field.field_name)
if isinstance(val, list):
parsed[field.field_name] = "|".join(map(str, val))
elif isinstance(val, str):
parsed[field.field_name] = val
field_value = parsed[field.field_name]
if field_value == "N/A":
field.field_type = "full_context"
if field.field_name == "PAYER_NAME":
field.prompt = field.prompt + prompt_templates.FULL_CONTEXT_PAYER_NAME_ADDITIONAL_INSTRUCTION()
logging.warning(
f"Obtained N/A for {field.field_name}, marking as full_context"
)
return (field.field_name, field_value, field)
except Exception as e:
logging.warning(
f"Failed to parse JSON response for field {field.field_name}: "
f"{type(e).__name__}: {str(e)}. Setting field value to N/A."
)
return (field.field_name, "N/A", field)
except Exception as e:
# Mark field for full_context retry and preserve error info
field.field_type = "full_context"
logging.error(
f"Error on {field.field_name}: {type(e).__name__}: {str(e)}, marking as full_context"
)
return (field.field_name, "N/A", field)
def clip_context_by_tokens(context: str, max_tokens: int = 150000) -> str:
"""
Clip context text if it exceeds max_tokens.
@@ -184,112 +264,48 @@ def run_hybrid_smart_chunked_fields(
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", " ", ""],
)
with timing_utils.timed_block("hsc.setup", context=filename):
reranker = get_reranker_client(rag_config.RERANKER)
embeddings = get_embeddings(
rag_config.EMBEDDING_MODEL_ID, rag_config.BEDROCK_REGION
)
splitter = RecursiveCharacterTextSplitter(
chunk_size=rag_config.CHUNK_SIZE,
chunk_overlap=rag_config.CHUNK_OVERLAP,
add_start_index=True,
strip_whitespace=True,
separators=["\n\n", "\n", " ", ""],
)
sections, is_preserved = preprocessing_pipeline(contract_text)
docs = prepare_documents_for_chunking(sections, filename)
chunks = splitter.split_documents(docs)
del docs
gc.collect()
with timing_utils.timed_block("hsc.preprocessing", context=filename):
sections, is_preserved = preprocessing_pipeline(contract_text)
docs = prepare_documents_for_chunking(sections, filename)
chunks = splitter.split_documents(docs)
del docs
gc.collect()
if not chunks:
return {}
# Build retriever
vectorstore = create_vectorstore(chunks, embeddings)
retriever = create_hybrid_retriever(chunks, vectorstore, rag_config)
with timing_utils.timed_block("hsc.vectorstore_creation", context=filename):
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
with timing_utils.timed_block("hsc.parallel_field_processing", context=filename):
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(fields_to_process), 20)) as executor:
futures = [
executor.submit(process_single_field, field, retriever, reranker, rag_config, text_dict, constants, filename)
for field in fields_to_process
]
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_dict(
constants
) # Returns dict of {field_name : prompt}
prompt = prompt_templates.ONE_TO_ONE_SINGLE_FIELD_TEMPLATE(
context, llm_prompt
)
logging.debug(f"Field: {field.field_name}")
logging.debug(f"Retrieval question: {retrieval_query}")
logging.debug(f"LLM prompt: {llm_prompt}")
# LLM call
raw = llm_utils.invoke_claude(
prompt, "sonnet_latest", filename, max_tokens=8192
)
logging.debug(f"Hybrid Smart Chunking RAG Raw response for {field.field_name}: {raw}")
try:
parsed = string_utils.universal_json_load(
raw
) # returns dict
val = parsed.get(field.field_name)
if isinstance(val, list):
parsed[field.field_name] = "|".join(map(str, val))
elif isinstance(val, str):
parsed[field.field_name] = val
answers_dict[field.field_name] = parsed[field.field_name]
if answers_dict[field.field_name] == "N/A":
field.field_type = "full_context"
if field.field_name == "PAYER_NAME":
field.prompt = field.prompt + prompt_templates.FULL_CONTEXT_PAYER_NAME_ADDITIONAL_INSTRUCTION()
logging.warning(
f"Obtained N/A for {field.field_name}, marking as full_context"
)
except Exception as e:
logging.warning(
f"Failed to parse JSON response for field {field.field_name}: "
f"{type(e).__name__}: {str(e)}. Setting field value to N/A."
)
answers_dict[field.field_name] = "N/A"
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"
for future in concurrent.futures.as_completed(futures):
try:
field_name, field_value, field = future.result()
answers_dict[field_name] = field_value
except Exception as e:
logging.error(f"Error processing field: {str(e)}")
del retriever, vectorstore, chunks
gc.collect()
fieldExtraction/src/investment/main.py
+118 -86
View File
@@ -2,6 +2,7 @@ import concurrent.futures
import logging
import os
import random
import time
import traceback
from datetime import datetime
@@ -15,6 +16,7 @@ import src.utils.io_utils as io_utils
import src.utils.llm_utils as llm_utils
import src.utils.logging_utils as logging_utils
import src.utils.usage_tracking as usage_tracking
import src.utils.timing_utils as timing_utils
from constants.constants import Constants
from src import config
from src.parent_child import main as parent_child_main
@@ -84,20 +86,25 @@ def main(testing=False, test_params={}):
# Windows paths cannot contain ':'; use '-' in time component
run_timestamp = datetime.now().strftime(f"run_%Y%m%d_%H-%M_{config.BATCH_ID}")
# Track overall pipeline time
pipeline_start = time.time()
# Read Constants
constants = Constants()
with timing_utils.timed_block("read_constants", log_level="INFO"):
constants = Constants()
# Read and process input
if not testing:
input_dict = io_utils.read_input()
max_workers = config.MAX_WORKERS
else:
input_dict = io_utils.read_input(test_params["local_input_dir"])
max_workers = test_params["max_workers"]
if "input_files" in test_params:
input_dict = {
k: v for k, v in input_dict.items() if k in test_params["input_files"]
}
with timing_utils.timed_block("read_input", log_level="INFO"):
if not testing:
input_dict = io_utils.read_input()
max_workers = config.MAX_WORKERS
else:
input_dict = io_utils.read_input(test_params["local_input_dir"])
max_workers = test_params["max_workers"]
if "input_files" in test_params:
input_dict = {
k: v for k, v in input_dict.items() if k in test_params["input_files"]
}
logging.info(f"Total Input Files : {len(input_dict)}")
@@ -117,82 +124,98 @@ def main(testing=False, test_params={}):
# 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...")
try:
cacheable_instructions = prompt_templates.get_cacheable_instructions(constants)
for cache_key, instruction in cacheable_instructions.items():
llm_utils.warm_prompt_cache(
instruction=instruction,
cache_key=cache_key,
model_id="sonnet_latest"
)
logging.info(f"Successfully warmed {len(cacheable_instructions)} prompt caches")
except Exception as e:
logging.warning(f"Error warming caches (will proceed anyway): {str(e)}")
with timing_utils.timed_block("warm_prompt_caches", log_level="INFO"):
try:
cacheable_instructions = prompt_templates.get_cacheable_instructions(constants)
for cache_key, instruction in cacheable_instructions.items():
llm_utils.warm_prompt_cache(
instruction=instruction,
cache_key=cache_key,
model_id="sonnet_latest"
)
logging.info(f"Successfully warmed {len(cacheable_instructions)} prompt caches")
except Exception as e:
logging.warning(f"Error warming caches (will proceed anyway): {str(e)}")
# Process files concurrently - each thread gets its own per-file logging context
successful_results = []
error_results = []
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
# Each submitted task will set its own logging context in file_processing.process_file()
# We use an executor to process files concurrently; use submit() instead of map()
# so that we can catch exceptions and continue processing other files
futures = [
executor.submit(safe_process_file, item, constants, run_timestamp)
for item in input_dict.items()
]
# collect results as they complete
for future in concurrent.futures.as_completed(futures):
try:
results = future.result()
if "error" in results.columns:
error_results.append(results)
else:
successful_results.append(results)
except (
Exception
) as e: # When there's an issue with the future itself (not the processing inside the future)
logging.error(f"Error processing future: {str(e)}")
# Coerce this to the format expected by pd.concat (so a single-row dataframe)
error_results.append(pd.DataFrame([{"error": str(e)}]))
logging.info(f"Starting parallel file processing with {max_workers} workers...")
with timing_utils.timed_block("parallel_file_processing", log_level="INFO"):
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
# Each submitted task will set its own logging context in file_processing.process_file()
# We use an executor to process files concurrently; use submit() instead of map()
# so that we can catch exceptions and continue processing other files
futures = [
executor.submit(safe_process_file, item, constants, run_timestamp)
for item in input_dict.items()
]
# collect results as they complete
completed_count = 0
for future in concurrent.futures.as_completed(futures):
try:
results = future.result()
if "error" in results.columns:
error_results.append(results)
else:
successful_results.append(results)
completed_count += 1
if completed_count % 5 == 0: # Log progress every 5 files
logging.info(f"Processed {completed_count}/{len(futures)} files...")
except (
Exception
) as e: # When there's an issue with the future itself (not the processing inside the future)
logging.error(f"Error processing future: {str(e)}")
# Coerce this to the format expected by pd.concat (so a single-row dataframe)
error_results.append(pd.DataFrame([{"error": str(e)}]))
# only concat if we have results
if len(successful_results) > 0:
FINAL_RESULT_DF = pd.concat(successful_results)
else:
FINAL_RESULT_DF = pd.DataFrame()
with timing_utils.timed_block("concat_results", log_level="INFO"):
if len(successful_results) > 0:
FINAL_RESULT_DF = pd.concat(successful_results, ignore_index=True)
else:
FINAL_RESULT_DF = pd.DataFrame()
if len(error_results) > 0:
ERROR_RESULT_DF = pd.concat(error_results)
else:
ERROR_RESULT_DF = pd.DataFrame()
if len(error_results) > 0:
ERROR_RESULT_DF = pd.concat(error_results, ignore_index=True)
else:
ERROR_RESULT_DF = pd.DataFrame()
# Run QC/QA validation if enabled (creates separate validated copy)
if config.ENABLE_QC_QA and not FINAL_RESULT_DF.empty:
logging.info("=" * 80)
logging.info("Running QC/QA Validation Pipeline...")
logging.info("=" * 80)
try:
# Create a copy for validation - preserve original FINAL_RESULT_DF
validated_df = run_qc_qa_pipeline(FINAL_RESULT_DF.copy(), stats_df=None)
logging.info("QC/QA validation completed successfully")
with timing_utils.timed_block("qc_qa_validation", log_level="INFO"):
try:
# Ensure unique index before QC/QA (safety check)
if FINAL_RESULT_DF.index.has_duplicates:
logging.warning("Detected duplicate indices in FINAL_RESULT_DF, resetting index...")
FINAL_RESULT_DF = FINAL_RESULT_DF.reset_index(drop=True)
# Generate validation statistics
logging.info("Generating QC/QA statistics...")
qc_qa_stats_df = generate_statistics(validated_df)
logging.info("QC/QA statistics generated successfully")
# Create a copy for validation - preserve original FINAL_RESULT_DF
validated_df = run_qc_qa_pipeline(FINAL_RESULT_DF.copy(), stats_df=None)
logging.info("QC/QA validation completed successfully")
# Save QC/QA outputs (local + S3) - separate from main results
save_qc_qa_outputs(
validated_df=validated_df,
stats_df=qc_qa_stats_df,
run_timestamp=run_timestamp,
write_to_s3=config.WRITE_TO_S3
)
except Exception as e:
logging.error(f"QC/QA validation failed: {str(e)}")
logging.error(f"Full traceback:\n{traceback.format_exc()}")
logging.warning("Continuing with original unvalidated results...")
# Generate validation statistics
logging.info("Generating QC/QA statistics...")
qc_qa_stats_df = generate_statistics(validated_df)
logging.info("QC/QA statistics generated successfully")
# Save QC/QA outputs (local + S3) - separate from main results
save_qc_qa_outputs(
validated_df=validated_df,
stats_df=qc_qa_stats_df,
run_timestamp=run_timestamp,
write_to_s3=config.WRITE_TO_S3
)
except Exception as e:
logging.error(f"QC/QA validation failed: {str(e)}")
logging.error(f"Full traceback:\n{traceback.format_exc()}")
logging.warning("Continuing with original unvalidated results...")
logging.info("=" * 80)
if not testing:
@@ -201,29 +224,38 @@ def main(testing=False, test_params={}):
config.BATCH_ID, run_timestamp
)
if config.WRITE_TO_S3:
doczy_output_for_pc = io_utils.write_s3(FINAL_RESULT_DF, "", run_timestamp, "final")
if not ERROR_RESULT_DF.empty:
doczy_output_for_pc = io_utils.write_s3(ERROR_RESULT_DF, "", run_timestamp, "error")
with timing_utils.timed_block("write_outputs", log_level="INFO"):
if config.WRITE_TO_S3:
doczy_output_for_pc = io_utils.write_s3(FINAL_RESULT_DF, "", run_timestamp, "final")
if not ERROR_RESULT_DF.empty:
doczy_output_for_pc = io_utils.write_s3(ERROR_RESULT_DF, "", run_timestamp, "error")
# Export usage and cost tracking data to S3 only if both dataframes have data
if not per_file_usage_df.empty and not batch_summary_df.empty:
logging.info("Exporting usage and cost tracking data to S3...")
io_utils.write_s3(per_file_usage_df, "", run_timestamp, "usage")
io_utils.write_s3(batch_summary_df, "", run_timestamp, "usage_summary")
else:
doczy_output_for_pc = io_utils.write_local(FINAL_RESULT_DF, "", run_timestamp, "final")
if not ERROR_RESULT_DF.empty:
doczy_output_for_pc = io_utils.write_local(ERROR_RESULT_DF, "", run_timestamp, "error")
# Export usage and cost tracking data to S3 only if both dataframes have data
if not per_file_usage_df.empty and not batch_summary_df.empty:
logging.info("Exporting usage and cost tracking data to S3...")
io_utils.write_s3(per_file_usage_df, "", run_timestamp, "usage")
io_utils.write_s3(batch_summary_df, "", run_timestamp, "usage_summary")
else:
doczy_output_for_pc = io_utils.write_local(FINAL_RESULT_DF, "", run_timestamp, "final")
if not ERROR_RESULT_DF.empty:
doczy_output_for_pc = io_utils.write_local(ERROR_RESULT_DF, "", run_timestamp, "error")
if config.PERFORM_PARENT_CHILD_MAPPING:
parent_child_main.main(doczy_output_for_pc)
with timing_utils.timed_block("parent_child_mapping", log_level="INFO"):
parent_child_main.main(doczy_output_for_pc)
# Export usage and cost tracking data locally only if both dataframes have data
if not per_file_usage_df.empty and not batch_summary_df.empty:
logging.info("Exporting usage and cost tracking data locally...")
io_utils.write_local(per_file_usage_df, "", run_timestamp, "usage")
io_utils.write_local(batch_summary_df, "", run_timestamp, "usage_summary")
# Print timing summary
pipeline_duration = time.time() - pipeline_start
logging.info("=" * 80)
logging.info(f"PIPELINE COMPLETE - Total time: {pipeline_duration:.2f}s")
logging.info("=" * 80)
timing_utils.print_hierarchical_timing_summary()
else:
return FINAL_RESULT_DF, ERROR_RESULT_DF
fieldExtraction/src/investment/one_to_n_funcs.py
+180 -74
View File
@@ -1,3 +1,4 @@
import concurrent.futures
import logging
from src.investment import dynamic_funcs, prompt_calls, aarete_derived, postprocessing_funcs
@@ -167,6 +168,48 @@ def breakout(
return reimbursement_level_answers, special_case_answers
def process_single_carveout(answer_dict, carveout_definitions, special_case_definitions, filename):
"""Process a single carveout/special case in parallel"""
service_term, reimb_term = answer_dict.get("SERVICE_TERM"), answer_dict.get("REIMB_TERM")
reimbursement_categorization = prompt_calls.prompt_carveout_check(
service_term, reimb_term, carveout_definitions, special_case_definitions, filename
)
if reimbursement_categorization in carveout_definitions.keys():
answer_dict["CARVEOUT_CD"] = reimbursement_categorization
# Standard N/A OR Base response gets CARVEOUT_IND=N
if reimbursement_categorization in ["N/A", "BASE_COVERED_SERVICES", "DEFAULT_TERM"]:
answer_dict["CARVEOUT_IND"] = "N"
else:
answer_dict["CARVEOUT_IND"] = "Y"
# Determine DEFAULT_IND based on specific carveout_answer values
answer_dict["DEFAULT_IND"] = (
"Y" if reimbursement_categorization in ["DEFAULT_TERM", "UNLISTED_CODE"] else "N"
)
# Carveout breakouts
if reimbursement_categorization == "TRIGGER_CAP":
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
prompt_templates.TRIGGER_CAP_BREAKOUT, reimb_term, filename
)
answer_dict.update(breakout_answer_dict)
if reimbursement_categorization == "ADDITION":
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
prompt_templates.ADDITION_BREAKOUT, reimb_term, filename
)
answer_dict.update(breakout_answer_dict)
return ("reimbursement", answer_dict)
elif reimbursement_categorization in special_case_definitions.keys():
answer_dict[reimbursement_categorization] = reimb_term
answer_dict.pop("REIMB_TERM", None)
answer_dict.pop("SERVICE_TERM", None)
return ("special_case", answer_dict)
else:
return (None, answer_dict)
def carveout_and_special_case(
reimbursement_breakout_answers: list[dict], constants: Constants, filename: str
) -> tuple[list[dict], list[dict]]:
@@ -185,43 +228,23 @@ def carveout_and_special_case(
reimbursement_level_answers, special_case_answers = [], []
for answer_dict in reimbursement_breakout_answers:
service_term, reimb_term = answer_dict.get("SERVICE_TERM"), answer_dict.get("REIMB_TERM")
reimbursement_categorization = prompt_calls.prompt_carveout_check(
service_term, reimb_term, carveout_definitions, special_case_definitions, filename
)
if reimbursement_categorization in carveout_definitions.keys():
answer_dict["CARVEOUT_CD"] = reimbursement_categorization
# Standard N/A OR Base response gets CARVEOUT_IND=N
if reimbursement_categorization in ["N/A", "BASE_COVERED_SERVICES", "DEFAULT_TERM"]:
answer_dict["CARVEOUT_IND"] = "N"
else:
answer_dict["CARVEOUT_IND"] = "Y"
# Determine DEFAULT_IND based on specific carveout_answer values
answer_dict["DEFAULT_IND"] = (
"Y" if reimbursement_categorization in ["DEFAULT_TERM", "UNLISTED_CODE"] else "N"
)
# Parallelize carveout processing
if reimbursement_breakout_answers:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(reimbursement_breakout_answers), 5)) as executor:
futures = [
executor.submit(process_single_carveout, answer_dict.copy(), carveout_definitions, special_case_definitions, filename)
for answer_dict in reimbursement_breakout_answers
]
# Carveout breakouts
if reimbursement_categorization == "TRIGGER_CAP":
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
prompt_templates.TRIGGER_CAP_BREAKOUT, reimb_term, filename
)
answer_dict.update(breakout_answer_dict)
if reimbursement_categorization == "ADDITION":
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
prompt_templates.ADDITION_BREAKOUT, reimb_term, filename
)
answer_dict.update(breakout_answer_dict)
reimbursement_level_answers.append(answer_dict)
elif reimbursement_categorization in special_case_definitions.keys():
answer_dict[reimbursement_categorization] = reimb_term
answer_dict.pop("REIMB_TERM", None)
answer_dict.pop("SERVICE_TERM", None)
special_case_answers.append(answer_dict)
for future in concurrent.futures.as_completed(futures):
try:
result_type, result_dict = future.result()
if result_type == "reimbursement":
reimbursement_level_answers.append(result_dict)
elif result_type == "special_case":
special_case_answers.append(result_dict)
except Exception as e:
logging.error(f"Error in carveout processing: {str(e)}")
return reimbursement_level_answers, special_case_answers
@@ -245,10 +268,21 @@ def methodology_breakout(
"""
final_answers = []
for answer_dict in reimbursement_primary_answers:
final_answers += methodology_breakout_single_row(
answer_dict, constants, filename
)
# Parallelize methodology breakout processing
if reimbursement_primary_answers:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(reimbursement_primary_answers), 5)) as executor:
futures = [
executor.submit(methodology_breakout_single_row, answer_dict, constants, filename)
for answer_dict in reimbursement_primary_answers
]
for future in concurrent.futures.as_completed(futures):
try:
final_answers += future.result()
except Exception as e:
logging.error(f"Error in methodology breakout: {str(e)}")
return final_answers
@@ -398,6 +432,19 @@ def methodology_breakout_secondary(
}
def process_single_special_case_breakout(answer_dict, special_case_fields, filename):
"""Process a single special case breakout in parallel"""
term_field_name = [field_name for field_name in answer_dict.keys() if "_TERM" in field_name and field_name not in ["SERVICE_TERM","REIMB_TERM",]][0]
term_answer = answer_dict[term_field_name]
term_field = special_case_fields.get_field(term_field_name)
breakout_template = term_field.get_breakout_template()
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
breakout_template, term_answer, filename
)
answer_dict.update(breakout_answer_dict)
return answer_dict
def special_case_breakout(
special_case_answers: list[dict],
filename: str,
@@ -423,18 +470,36 @@ def special_case_breakout(
dict[str, list]: A dictionary containing the processed breakout answers for each
special case field.
"""
if not special_case_answers:
return special_case_answers
special_case_fields = FieldSet(config.FIELD_JSON_PATH, field_type="special_case_check")
for answer_dict in special_case_answers:
term_field_name = [field_name for field_name in answer_dict.keys() if "_TERM" in field_name and field_name not in ["SERVICE_TERM","REIMB_TERM",]][0]
term_answer = answer_dict[term_field_name]
term_field = special_case_fields.get_field(term_field_name)
breakout_template = term_field.get_breakout_template()
breakout_answer_dict = prompt_calls.prompt_special_case_breakout(
breakout_template, term_answer, filename
)
answer_dict.update(breakout_answer_dict)
return special_case_answers
# Parallelize special case breakout processing
updated_answers = []
if special_case_answers:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(special_case_answers), 5)) as executor:
futures = [
executor.submit(process_single_special_case_breakout, answer_dict.copy(), special_case_fields, filename)
for answer_dict in special_case_answers
]
for future in concurrent.futures.as_completed(futures):
try:
updated_answers.append(future.result())
except Exception as e:
logging.error(f"Error in special case breakout: {str(e)}")
return updated_answers
def validate_single_reimbursement(answer_dict, filename):
"""Validate a single reimbursement in parallel"""
if prompt_calls.validate_reimbursements_for_llm(answer_dict, filename):
return answer_dict
else:
logging.debug(f"LLM filtered out: {answer_dict}")
return None
def filter_services_without_reimbursements(
@@ -467,11 +532,29 @@ def filter_services_without_reimbursements(
valid_reimbursements = []
for answer_dict in reimbursement_primary_answers:
if prompt_calls.validate_reimbursements_for_llm(answer_dict, filename):
valid_reimbursements.append(answer_dict)
else:
logging.debug(f"LLM filtered out: {answer_dict}")
# Parallelize validation
if reimbursement_primary_answers:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(reimbursement_primary_answers), 5)) as executor:
# Create futures with index tracking
future_to_index = {
executor.submit(validate_single_reimbursement, answer_dict, filename): i
for i, answer_dict in enumerate(reimbursement_primary_answers)
}
# Collect results with their original indices
results_with_indices = []
for future in concurrent.futures.as_completed(future_to_index):
try:
result = future.result()
if result is not None:
idx = future_to_index[future]
results_with_indices.append((idx, result))
except Exception as e:
logging.error(f"Error in reimbursement validation: {str(e)}")
# Sort by original index to preserve order
results_with_indices.sort(key=lambda x: x[0])
valid_reimbursements = [result for _, result in results_with_indices]
filtered_count = len(reimbursement_primary_answers) - len(valid_reimbursements)
if filtered_count > 0:
@@ -583,6 +666,29 @@ def split_reimb_dates(one_to_n_results: list, filename: str) -> list:
return one_to_n_results
def process_single_lob_relationship(answer_dict, exhibit_text, filename):
"""Process LOB relationship for a single answer dict in parallel"""
if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_LOB")):
return answer_dict
# Process for program
if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
answer_dict["LOB_PROGRAM_RELATIONSHIP"] = (
prompt_calls.prompt_lob_relationship(
answer_dict, "AARETE_DERIVED_PROGRAM", exhibit_text, filename
)
)
# Process for Product
if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PRODUCT")):
answer_dict["LOB_PRODUCT_RELATIONSHIP"] = (
prompt_calls.prompt_lob_relationship(
answer_dict, "AARETE_DERIVED_PRODUCT", exhibit_text, filename
)
)
return answer_dict
def get_lob_relationship(answer_dicts, exhibit_text, filename):
"""
Populates LOB_PROGRAM_RELATIONSHIP and LOB_PRODUCT_RELATIONSHIP fields in answer_dicts
@@ -596,25 +702,25 @@ def get_lob_relationship(answer_dicts, exhibit_text, filename):
Returns:
list[dict]: The updated list of dictionaries with LOB relationships populated.
"""
for answer_dict in answer_dicts:
if string_utils.is_empty(answer_dict.get("AARETE_DERIVED_LOB")):
continue
# Process for program
if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PROGRAM")):
answer_dict["LOB_PROGRAM_RELATIONSHIP"] = (
prompt_calls.prompt_lob_relationship(
answer_dict, "AARETE_DERIVED_PROGRAM", exhibit_text, filename
)
)
# Process for Product
if not string_utils.is_empty(answer_dict.get("AARETE_DERIVED_PRODUCT")):
answer_dict["LOB_PRODUCT_RELATIONSHIP"] = (
prompt_calls.prompt_lob_relationship(
answer_dict, "AARETE_DERIVED_PRODUCT", exhibit_text, filename
)
)
if not answer_dicts:
return answer_dicts
return answer_dicts
# Parallelize LOB relationship processing
updated_dicts = []
if answer_dicts:
with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(answer_dicts), 5)) as executor:
futures = [
executor.submit(process_single_lob_relationship, answer_dict.copy(), exhibit_text, filename)
for answer_dict in answer_dicts
]
for future in concurrent.futures.as_completed(futures):
try:
updated_dicts.append(future.result())
except Exception as e:
logging.error(f"Error in LOB relationship processing: {str(e)}")
return updated_dicts
def one_to_n_cleaning(all_exhibit_rows: list[dict], exhibit_text: str, constants: Constants, filename: str):
fieldExtraction/src/parent_child/main.py
+1 -1
View File
@@ -39,7 +39,7 @@ def main(doczy_output_for_pc):
)
return mapped_df.shape[0]
except Exception as e:
logging.info(f"Parent child mapping failed due to {e}")
logging.error(f"Parent child mapping failed due to {e}")
if __name__ == "__main__":
fieldExtraction/src/testbed/testbed_metrics.py
+1 -1
View File
@@ -10,7 +10,7 @@ warnings.filterwarnings("ignore")
########################## READ AND PREPROCESS ##########################
# Read files
testbed = pd.read_excel("Doczy-Testbed.xlsx") # This is the testbed
results = pd.read_csv("Doczy-Results.csv") # This is the doczy output
results = pd.read_excel("Doczy-Results.xlsx") # This is the doczy output
########################## Preprocess ##########################
testbed = testbed_utils.testbed_preprocess(testbed)
fieldExtraction/src/utils/timing_utils.py
+232
View File
@@ -0,0 +1,232 @@
"""
Timing utilities for performance monitoring and bottleneck identification.
Provides decorators and context managers for tracking execution time of functions and code blocks.
Includes thread-safe statistics collection for concurrent operations.
"""
import functools
import logging
import time
from contextlib import contextmanager
from dataclasses import dataclass, field
from threading import Lock
from typing import Dict, List, Optional, Callable, Any
from collections import defaultdict
@dataclass
class TimingStats:
"""Statistics for a timed operation."""
name: str
count: int = 0
total_time: float = 0.0
min_time: float = float('inf')
max_time: float = 0.0
times: List[float] = field(default_factory=list)
def add(self, duration: float):
"""Add a timing measurement."""
self.count += 1
self.total_time += duration
self.min_time = min(self.min_time, duration)
self.max_time = max(self.max_time, duration)
self.times.append(duration)
@property
def avg_time(self) -> float:
"""Calculate average time."""
return self.total_time / self.count if self.count > 0 else 0.0
def __str__(self) -> str:
"""Format timing statistics."""
if self.count == 0:
return f"{self.name}: No data"
return (
f"{self.name}: "
f"count={self.count}, "
f"total={self.total_time:.2f}s, "
f"avg={self.avg_time:.2f}s, "
f"min={self.min_time:.2f}s, "
f"max={self.max_time:.2f}s"
)
class TimingTracker:
"""Thread-safe timing statistics tracker."""
def __init__(self):
self._stats: Dict[str, TimingStats] = {}
self._lock = Lock()
self._enabled = True
def record(self, name: str, duration: float, context: Optional[str] = None):
"""Record a timing measurement."""
if not self._enabled:
return
key = f"{context}.{name}" if context else name
with self._lock:
if key not in self._stats:
self._stats[key] = TimingStats(name=key)
self._stats[key].add(duration)
def get_stats(self, name: Optional[str] = None) -> Dict[str, TimingStats]:
"""Get timing statistics."""
with self._lock:
if name:
return {k: v for k, v in self._stats.items() if name in k}
return dict(self._stats)
def print_summary(self, min_time: float = 0.0):
"""Print timing summary sorted by total time."""
with self._lock:
if not self._stats:
logging.info("No timing data collected")
return
# Filter and sort by total time
filtered_stats = [s for s in self._stats.values() if s.total_time >= min_time]
sorted_stats = sorted(filtered_stats, key=lambda s: s.total_time, reverse=True)
logging.info("=" * 80)
logging.info("TIMING SUMMARY (sorted by total time)")
logging.info("=" * 80)
for stat in sorted_stats:
logging.info(str(stat))
# Calculate total tracked time
total_tracked = sum(s.total_time for s in self._stats.values())
logging.info("=" * 80)
logging.info(f"Total tracked time: {total_tracked:.2f}s")
logging.info("=" * 80)
def print_hierarchical_summary(self):
"""Print timing summary organized hierarchically by component."""
with self._lock:
if not self._stats:
logging.info("No timing data collected")
return
# Group stats by top-level component
components = defaultdict(list)
for name, stat in self._stats.items():
component = name.split('.')[0]
components[component].append(stat)
logging.info("=" * 80)
logging.info("HIERARCHICAL TIMING SUMMARY")
logging.info("=" * 80)
# Sort components by total time
component_totals = {
comp: sum(s.total_time for s in stats)
for comp, stats in components.items()
}
sorted_components = sorted(component_totals.items(), key=lambda x: x[1], reverse=True)
for component, total_time in sorted_components:
logging.info(f"\n{component}: {total_time:.2f}s total")
logging.info("-" * 60)
# Sort stats within component by total time
sorted_stats = sorted(components[component], key=lambda s: s.total_time, reverse=True)
for stat in sorted_stats:
indent = " " * (stat.name.count('.'))
logging.info(f"{indent}{stat}")
logging.info("=" * 80)
def reset(self):
"""Reset all timing statistics."""
with self._lock:
self._stats.clear()
def enable(self):
"""Enable timing tracking."""
self._enabled = True
def disable(self):
"""Disable timing tracking."""
self._enabled = False
# Global timing tracker instance
_global_tracker = TimingTracker()
def get_tracker() -> TimingTracker:
"""Get the global timing tracker."""
return _global_tracker
@contextmanager
def timed_block(name: str, context: Optional[str] = None, log_level: str = "DEBUG"):
"""
Context manager for timing a block of code.
Usage:
with timed_block("database_query", context="user_123"):
# code to time
result = db.query(...)
"""
start_time = time.time()
try:
yield
finally:
duration = time.time() - start_time
_global_tracker.record(name, duration, context)
# Also log the timing
log_func = getattr(logging, log_level.lower(), logging.debug)
full_name = f"{context}.{name}" if context else name
log_func(f"⏱️ {full_name}: {duration:.2f}s")
def timed(name: Optional[str] = None, context: Optional[str] = None, log_level: str = "DEBUG"):
"""
Decorator for timing function execution.
Usage:
@timed("process_document")
def process_document(doc):
# function code
pass
"""
def decorator(func: Callable) -> Callable:
func_name = name or func.__name__
@functools.wraps(func)
def wrapper(*args, **kwargs) -> Any:
start_time = time.time()
try:
result = func(*args, **kwargs)
return result
finally:
duration = time.time() - start_time
_global_tracker.record(func_name, duration, context)
# Also log the timing
log_func = getattr(logging, log_level.lower(), logging.debug)
full_name = f"{context}.{func_name}" if context else func_name
log_func(f"⏱️ {full_name}: {duration:.2f}s")
return wrapper
return decorator
def print_timing_summary(min_time: float = 0.0):
"""Print timing summary (convenience function)."""
_global_tracker.print_summary(min_time)
def print_hierarchical_timing_summary():
"""Print hierarchical timing summary (convenience function)."""
_global_tracker.print_hierarchical_summary()
def reset_timing_stats():
"""Reset timing statistics (convenience function)."""
_global_tracker.reset()
fieldExtraction/tests/test_exhibit_funcs.py
+29
View File
@@ -0,0 +1,29 @@
from src.investment.exhibit_funcs import Exhibit, get_exhibit_list
from src.prompts.fieldset import FieldSet
def test_add_reimbursement_rows_sets_flag():
exhibit = Exhibit(exhibit_page="1", exhibit_page_nums=["1"], exhibit_header="H", exhibit_text="text")
exhibit.add_reimbursement_rows([{"a": 1}], [])
assert exhibit.has_reimbursements is True
assert len(exhibit.reimbursement_rows) == 1
def test_set_exhibit_level_and_previous_dynamic_fields():
prev = Exhibit(exhibit_page="1", exhibit_page_nums=["1"], exhibit_header="H1", exhibit_text="text1")
dyn_fields = FieldSet()
prev.dynamic_primary_fields = dyn_fields
current = Exhibit(exhibit_page="2", exhibit_page_nums=["2"], exhibit_header="H2", exhibit_text="text2", prev_exhibit=prev)
current.set_exhibit_level_data({"X": "Y"}, FieldSet())
assert current.exhibit_level_answers["X"] == "Y"
assert current.get_previous_exhibit_dynamic_fields() is dyn_fields
def test_get_exhibit_list_links_previous():
text_dict = {"1": "A", "2": "B"}
mapping = {"1": ["1"], "2": ["2"]}
headers = {"1": "H1", "2": "H2"}
exhibits = get_exhibit_list(text_dict, mapping, headers)
assert len(exhibits) == 2
assert exhibits[1].prev_exhibit is exhibits[0]
fieldExtraction/tests/test_timing_utils.py
+35
View File
@@ -0,0 +1,35 @@
import time
from src.utils import timing_utils
def test_timing_stats_add_and_avg():
stats = timing_utils.TimingStats(name="unit")
stats.add(0.1)
stats.add(0.3)
assert stats.count == 2
assert stats.min_time == 0.1
assert stats.max_time == 0.3
assert abs(stats.avg_time - 0.2) < 1e-6
def test_timing_tracker_record_and_reset():
tracker = timing_utils.TimingTracker()
tracker.record("step", 0.05, context="fileA")
tracker.record("step", 0.10, context="fileA")
stats = tracker.get_stats()
assert "fileA.step" in stats
assert stats["fileA.step"].count == 2
tracker.reset()
assert tracker.get_stats() == {}
def test_timed_block_records_duration():
tracker = timing_utils.get_tracker()
tracker.reset()
with timing_utils.timed_block("sleep_test", context="unit"):
time.sleep(0.01)
stats = tracker.get_stats()
key = "unit.sleep_test"
assert key in stats
assert stats[key].count == 1