diff --git a/src/config.py b/src/config.py index e71e652..11563b2 100644 --- a/src/config.py +++ b/src/config.py @@ -479,3 +479,13 @@ ENABLE_QC_QA = get_arg_value("enable_qc_qa", "True") == "True" # Leading state-based brand names (e.g., "Oklahoma Complete Health") # are preserved. STATE_FLAG = True + +############## INSTRUMENTATION SETTINGS ############## +# On by default; opt out with DOCZY_INSTRUMENTATION=0 (or false/no/off). +INSTRUMENTATION_ENABLED = os.environ.get("DOCZY_INSTRUMENTATION", "").lower() not in ( + "0", + "false", + "no", + "off", +) +INSTRUMENTATION_CSV_PATH = os.environ.get("DOCZY_INSTRUMENTATION_CSV", "") diff --git a/src/pipelines/runner.py b/src/pipelines/runner.py index 38184c4..8bf5c59 100644 --- a/src/pipelines/runner.py +++ b/src/pipelines/runner.py @@ -28,6 +28,7 @@ random.seed(42) warnings.filterwarnings("ignore", message=".*protected namespace.*") import src.prompts.prompt_templates as prompt_templates +import src.utils.instrumentation as instrumentation import src.utils.io_utils as io_utils import src.utils.llm_utils as llm_utils import src.utils.logging_utils as logging_utils @@ -167,6 +168,8 @@ def main(client: str = "saas", 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}") + instrumentation.configure_for_run(config.BATCH_ID, run_timestamp) + # Track overall pipeline time pipeline_start = time.time() @@ -427,6 +430,9 @@ def main(client: str = "saas", testing=False, test_params={}): io_utils.write_s3( batch_summary_df, "", run_timestamp, "usage_summary" ) + if instrumentation.is_enabled(): + instrumentation.close() + io_utils.upload_instrumentation_csv(run_timestamp) else: doczy_output_for_pc = io_utils.write_local( FINAL_RESULT_DF_CC, "", run_timestamp, "cc_results_full" @@ -509,6 +515,9 @@ def main(client: str = "saas", testing=False, test_params={}): io_utils.write_local(per_file_usage_df, "", run_timestamp, "usage") io_utils.write_local(batch_summary_df, "", run_timestamp, "usage_summary") + if instrumentation.is_enabled(): + instrumentation.close() + # Print timing summary pipeline_duration = time.time() - pipeline_start logging.info("=" * 80) diff --git a/src/pipelines/saas/file_processing.py b/src/pipelines/saas/file_processing.py index 5996bc5..79c0afa 100644 --- a/src/pipelines/saas/file_processing.py +++ b/src/pipelines/saas/file_processing.py @@ -28,7 +28,14 @@ from src.pipelines.vendors.shared.extraction import ( run_full_context_fields as vendor_run_full_context_fields, ) from src.pipelines.shared.extraction.exhibit_funcs import Exhibit, ExhibitChunk -from src.utils import io_utils, logging_utils, string_utils, timing_utils +from src.utils import ( + instrumentation, + instrumentation_context, + io_utils, + logging_utils, + string_utils, + timing_utils, +) from src.constants.constants import Constants from src import config from src.prompts.fieldset import FieldSet @@ -46,6 +53,7 @@ def process_file(file_object, constants: Constants, run_timestamp): # This includes logging from all called functions (preprocess, one_to_n_funcs, etc.) logging_utils.set_current_file(filename) logging.debug(f"{datetime_str()} Processing {filename}...") + instrumentation.log("file_start", filename=filename) # Set default values dynamic_one_to_one_fields = FieldSet() @@ -55,7 +63,10 @@ def process_file(file_object, constants: Constants, run_timestamp): final_results = pd.DataFrame([{"FILE_NAME": filename}]) ################## PREPROCESS ################## - with timing_utils.timed_block("preprocess", context=filename): + with ( + instrumentation_context.instr_scope(segment="preprocessing"), + timing_utils.timed_block("preprocess", context=filename), + ): contract_text = preprocess.clean_text(contract_text) text_dict, top_sheet_dict = preprocess.split_text(contract_text) if not text_dict: @@ -76,7 +87,10 @@ def process_file(file_object, constants: Constants, run_timestamp): # Create pages_dict from text_dict (headers/footers already applied) pages_dict = prep_funcs.split_large_tables(text_dict) - with timing_utils.timed_block("one_to_n_exhibit_chunking", context=filename): + with ( + instrumentation_context.instr_scope(segment="preprocessing"), + timing_utils.timed_block("one_to_n_exhibit_chunking", context=filename), + ): # Use pages_dict for exhibit chunking (preferred) or fall back to text_dict exhibit_header_dict = preprocess.one_to_n_exhibit_chunking( pages_dict=pages_dict, @@ -122,7 +136,10 @@ def process_file(file_object, constants: Constants, run_timestamp): if process_one_to_one: # Get specific fields to extract (if configured) specific_fields = config.get_specific_fields_list() - with timing_utils.timed_block("one_to_one_extraction", context=filename): + with ( + instrumentation_context.instr_scope(segment="one_to_one"), + timing_utils.timed_block("one_to_one_extraction", context=filename), + ): one_to_one_results = run_one_to_one_prompts( filename, contract_text, @@ -159,13 +176,19 @@ def process_file(file_object, constants: Constants, run_timestamp): # APPLY CODES IF ONE_TO_N WAS PROCESSED if not one_to_n_results.empty: - with timing_utils.timed_block("code_processing", context=filename): + with ( + instrumentation_context.instr_scope(segment="code_breakout"), + 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.debug(f"{datetime_str()} Codes Complete - {filename}") # POSTPROCESS - with timing_utils.timed_block("postprocess", context=filename): + with ( + instrumentation_context.instr_scope(segment="postprocess"), + timing_utils.timed_block("postprocess", context=filename), + ): cc_df, dashboard_df = postprocess.postprocess(final_results, constants) logging.debug(f"{datetime_str()} Postprocessing Complete - {filename}") @@ -178,6 +201,11 @@ def process_file(file_object, constants: Constants, run_timestamp): logging.debug(f"{datetime_str()} Writing Complete - {filename}") + instrumentation.log( + "file_end", + filename=filename, + extra_json=f'{{"final_rows": {len(cc_df)}}}', + ) return cc_df, dashboard_df @@ -358,64 +386,116 @@ def process_page_reimbursements( ) return [], [] - ############################### Reimbursement Primary ############################### - # Run reimbursement_level per-chunk: each relevant chunk is processed individually - # so the LLM sees focused, relevant text one chunk at a time - reimbursement_level_answers = [] - for chunk in relevant_chunks: - chunk_answers = one_to_n_funcs.reimbursement_level( - chunk.text, constants, filename + instrumentation.log("page_start", filename=filename, page_num=str(page_num)) + with instrumentation_context.instr_scope(page_num=str(page_num)): + ############################### Reimbursement Primary ############################### + # Run reimbursement_level per-chunk: each relevant chunk is processed individually + # so the LLM sees focused, relevant text one chunk at a time + reimbursement_level_answers = [] + for chunk in relevant_chunks: + with instrumentation_context.instr_scope(chunk_id=chunk.chunk_id): + before = len(reimbursement_level_answers) + chunk_answers = one_to_n_funcs.reimbursement_level( + chunk.text, constants, filename + ) + if chunk_answers: + reimbursement_level_answers.extend(chunk_answers) + instrumentation.log( + "row_count", + filename=filename, + stage="reimbursement_level_per_chunk", + page_num=str(page_num), + chunk_id=chunk.chunk_id, + n_in=before, + n_out=len(reimbursement_level_answers), + ) + + if not reimbursement_level_answers: + logging.debug( + f"{datetime_str()} Page {page_num}: No reimbursement answers found, skipping - {filename}" + ) + instrumentation.log("page_end", filename=filename, page_num=str(page_num)) + return [], [] + + # Page-level simplification of exhibit for downstream steps + exhibit_text_simplified = preprocessing_funcs.simplify_exhibit( + exhibit=exhibit, + current_page_num=page_num, ) - if chunk_answers: - reimbursement_level_answers.extend(chunk_answers) - if not reimbursement_level_answers: - logging.debug( - f"{datetime_str()} Page {page_num}: No reimbursement answers found, skipping - {filename}" + ################################ Carveouts and Special Case ################################ + _n_before_carveout = len(reimbursement_level_answers) + reimbursement_level_answers, special_case_answers = ( + one_to_n_funcs.carveout_and_special_case( + reimbursement_level_answers, constants, filename + ) + ) + instrumentation.log( + "row_count", + filename=filename, + stage="carveout", + page_num=str(page_num), + n_in=_n_before_carveout, + n_out=len(reimbursement_level_answers), + extra_json=f'{{"special_case_rows": {len(special_case_answers)}}}', ) - return [], [] - # Page-level simplification of exhibit for downstream steps - exhibit_text_simplified = preprocessing_funcs.simplify_exhibit( - exhibit=exhibit, - current_page_num=page_num, - ) - - ################################ Carveouts and Special Case ################################ - reimbursement_level_answers, special_case_answers = ( - one_to_n_funcs.carveout_and_special_case( + ################################ Dynamic Code Assignment ################################ + _n_before_dca = len(reimbursement_level_answers) + reimbursement_level_answers = one_to_n_funcs.dynamic_code_assignment( reimbursement_level_answers, constants, filename ) - ) + instrumentation.log( + "row_count", + filename=filename, + stage="dynamic_code_assignment", + page_num=str(page_num), + n_in=_n_before_dca, + n_out=len(reimbursement_level_answers), + ) - ################################ Dynamic Code Assignment ################################ - reimbursement_level_answers = one_to_n_funcs.dynamic_code_assignment( - reimbursement_level_answers, constants, filename - ) + ############################### Lesser of Distribution ############################### + # Note: We run lesser_of without dynamic fields in Step 1 + _n_before_lesser = len(reimbursement_level_answers) + reimbursement_level_answers = one_to_n_funcs.lesser_of_distribution( + reimbursement_level_answers, + exhibit_text_simplified, + page_num, + constants, + filename, + exhibit, + ) + instrumentation.log( + "row_count", + filename=filename, + stage="lesser_of_distribution", + page_num=str(page_num), + n_in=_n_before_lesser, + n_out=len(reimbursement_level_answers), + ) - ############################### Lesser of Distribution ############################### - # 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, - exhibit, - ) + ################################ Get Breakouts ############################### + _n_before_breakout = len(reimbursement_level_answers) + reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout( + reimbursement_level_answers, + special_case_answers, + filename, + constants, + ) + instrumentation.log( + "row_count", + filename=filename, + stage="breakout", + page_num=str(page_num), + n_in=_n_before_breakout, + n_out=len(reimbursement_level_answers), + ) - ################################ Get Breakouts ############################### - reimbursement_level_answers, special_case_answers = one_to_n_funcs.breakout( - reimbursement_level_answers, - special_case_answers, - filename, - constants, - ) - - ################################ Assign REIMB_PAGE ############################### - for answer_dict in reimbursement_level_answers: - answer_dict["REIMB_PAGE"] = page_num + ################################ Assign REIMB_PAGE ############################### + for answer_dict in reimbursement_level_answers: + answer_dict["REIMB_PAGE"] = page_num + instrumentation.log("page_end", filename=filename, page_num=str(page_num)) return reimbursement_level_answers, special_case_answers @@ -465,147 +545,214 @@ def run_one_to_n_prompts( # Process exhibits serially; within each exhibit, process pages in parallel one_to_n_results = [] - for exhibit in exhibits: - # Search for reimbursement-related chunks, then group by page - relevant_chunks = exhibit.get_relevant_chunks( - threshold=ESC_CONFIG.CHUNK_RELEVANCE_THRESHOLD - ) - page_chunks_map = ( - exhibit.group_chunks_by_page(relevant_chunks) if relevant_chunks else {} + for i, exhibit in enumerate(exhibits): + instrumentation.log( + "exhibit_start", + filename=filename, + exhibit_idx=i, + exhibit_header=( + exhibit.exhibit_header[:80] if exhibit.exhibit_header else "" + ), + exhibit_page=exhibit.exhibit_page, + extra_json=f'{{"num_pages": {len(exhibit.exhibit_page_nums)}}}', ) + with instrumentation_context.instr_scope( + exhibit_idx=i, + exhibit_header=( + exhibit.exhibit_header[:80] if exhibit.exhibit_header else "" + ), + exhibit_page=exhibit.exhibit_page, + ): + # Search for reimbursement-related chunks, then group by page + relevant_chunks = exhibit.get_relevant_chunks( + threshold=ESC_CONFIG.CHUNK_RELEVANCE_THRESHOLD + ) + page_chunks_map = ( + exhibit.group_chunks_by_page(relevant_chunks) if relevant_chunks else {} + ) - logging.debug( - f"{datetime_str()} Processing exhibit {exhibits.index(exhibit) + 1}/{len(exhibits)}: " - f"page={exhibit.exhibit_page}, header='{exhibit.exhibit_header}' - {filename}" - ) - - pages_to_process = exhibit.exhibit_page_nums - - with concurrent.futures.ThreadPoolExecutor( - max_workers=min(len(pages_to_process), 20) - ) as executor: - page_futures = { - executor.submit( - process_page_reimbursements, - exhibit, - page_num, - constants, - filename, - relevant_chunks=page_chunks_map.get( - page_num - ), # None for pages without relevant chunks - ): page_num - for page_num in pages_to_process - } - - 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 - ) - except Exception as e: - logging.error(f"Error processing page {page_num}: {str(e)}") - - if not exhibit.has_reimbursements: logging.debug( - f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: No reimbursements found, skipping steps 2-3 - {filename}" + f"{datetime_str()} Processing exhibit {i + 1}/{len(exhibits)}: " + f"page={exhibit.exhibit_page}, header='{exhibit.exhibit_header}' - {filename}" ) - continue - logging.debug( - f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: " - f"{len(exhibit.reimbursement_rows)} reimbursement row(s), " - f"{len(exhibit.special_case_rows)} special case row(s) - {filename}" - ) + pages_to_process = exhibit.exhibit_page_nums - # 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, - specific_fields=specific_fields, - ) - ) - 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}" - ) + with concurrent.futures.ThreadPoolExecutor( + max_workers=min(len(pages_to_process), 20) + ) as executor: + page_futures = { + instrumentation_context.submit_with_context( + executor, + process_page_reimbursements, + exhibit, + page_num, + constants, + filename, + relevant_chunks=page_chunks_map.get( + page_num + ), # None for pages without relevant chunks + ): page_num + for page_num in pages_to_process + } - # Check if we need reimbursement-level processing - # If specific_fields is set and doesn't include reimbursement fields, skip Step 3 - reimbursement_fields = {"SERVICE_TERM", "REIMB_TERM", "REIMB_PAGE"} - skip_reimbursement_processing = specific_fields is not None and not any( - f in reimbursement_fields for f in specific_fields - ) + 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 + ) + except Exception as e: + logging.error(f"Error processing page {page_num}: {str(e)}") - if skip_reimbursement_processing: - # For exhibit-level-only extraction, just create minimal rows - # with exhibit_level_answers to preserve the data - if exhibit.exhibit_level_answers: - minimal_row = exhibit.exhibit_level_answers.copy() - minimal_row["REIMB_PAGE"] = exhibit.exhibit_page - # Apply crosswalk mapping for derived fields (e.g., CLAIM_TYPE_CD -> AARETE_DERIVED_CLAIM_TYPE_CD) - minimal_rows_with_crosswalk = aarete_derived.get_crosswalk_fields( - [minimal_row], constants + if not exhibit.has_reimbursements: + logging.debug( + f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: No reimbursements found, skipping steps 2-3 - {filename}" ) - exhibit.final_rows = minimal_rows_with_crosswalk - one_to_n_results.extend(minimal_rows_with_crosswalk) - else: - # 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, - pages_dict, - text_dict, - exhibit, - constants, - filename, + instrumentation.log( + "exhibit_gate_skip", + filename=filename, + exhibit_idx=i, + exhibit_page=exhibit.exhibit_page, + extra_json='{"reason": "no_reimbursements_after_step1"}', ) + continue - # 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) logging.debug( f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: " - f"{len(combined_rows)} final row(s) after combine & clean - {filename}" + f"{len(exhibit.reimbursement_rows)} reimbursement row(s), " + f"{len(exhibit.special_case_rows)} special case row(s) - {filename}" ) + # STEP 2: exhibit_level for this exhibit + instrumentation.log( + "stage_transition", + filename=filename, + stage="step2_exhibit_level_start", + exhibit_idx=i, + ) + exhibit_level_answers, dynamic_reimbursement_fields = ( + one_to_n_funcs.exhibit_level( + exhibit.exhibit_text, + exhibit.exhibit_header, + exhibit.exhibit_page, + constants, + filename, + specific_fields=specific_fields, + ) + ) + exhibit.set_exhibit_level_data( + exhibit_level_answers, dynamic_reimbursement_fields + ) + instrumentation.log( + "stage_transition", + filename=filename, + stage="step2_exhibit_level_done", + exhibit_idx=i, + ) + logging.debug( + f"Exhibit Level Answers for {filename}, Page {exhibit.exhibit_page}: {exhibit_level_answers}" + ) + + # Check if we need reimbursement-level processing + # If specific_fields is set and doesn't include reimbursement fields, skip Step 3 + reimbursement_fields = {"SERVICE_TERM", "REIMB_TERM", "REIMB_PAGE"} + skip_reimbursement_processing = specific_fields is not None and not any( + f in reimbursement_fields for f in specific_fields + ) + + if skip_reimbursement_processing: + # For exhibit-level-only extraction, just create minimal rows + # with exhibit_level_answers to preserve the data + if exhibit.exhibit_level_answers: + minimal_row = exhibit.exhibit_level_answers.copy() + minimal_row["REIMB_PAGE"] = exhibit.exhibit_page + # Apply crosswalk mapping for derived fields (e.g., CLAIM_TYPE_CD -> AARETE_DERIVED_CLAIM_TYPE_CD) + minimal_rows_with_crosswalk = aarete_derived.get_crosswalk_fields( + [minimal_row], constants + ) + exhibit.final_rows = minimal_rows_with_crosswalk + one_to_n_results.extend(minimal_rows_with_crosswalk) + else: + # 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 + instrumentation.log( + "stage_transition", + filename=filename, + stage="step3_start", + exhibit_idx=i, + n_in=len(exhibit.reimbursement_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, + pages_dict, + 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, + ) + instrumentation.log( + "row_count", + filename=filename, + stage="combine_one_to_n", + exhibit_idx=i, + n_in=len(reimbursement_rows_with_dynamic), + n_out=len(combined_rows), + ) + + # Run cleaning + combined_rows = one_to_n_funcs.one_to_n_cleaning( + combined_rows, exhibit.exhibit_text, constants, filename + ) + instrumentation.log( + "row_count", + filename=filename, + stage="one_to_n_cleaning", + exhibit_idx=i, + n_in=len(combined_rows), + n_out=len(combined_rows), + ) + + exhibit.final_rows = combined_rows + one_to_n_results.extend(combined_rows) + instrumentation.log( + "stage_transition", + filename=filename, + stage="step3_end", + exhibit_idx=i, + n_out=len(combined_rows), + ) + logging.debug( + f"{datetime_str()} Exhibit page={exhibit.exhibit_page}: " + f"{len(combined_rows)} final row(s) after combine & clean - {filename}" + ) + logging.debug( f"{datetime_str()} STEP 3 COMPLETE: Combined {len(one_to_n_results)} total rows - {filename}" ) diff --git a/src/pipelines/saas/main.py b/src/pipelines/saas/main.py index c9cdc18..c476458 100644 --- a/src/pipelines/saas/main.py +++ b/src/pipelines/saas/main.py @@ -91,6 +91,9 @@ def safe_process_file(item, constants, run_timestamp): def main(testing=False, test_params={}): + from src.utils import instrumentation + + instrumentation.configure_from_env() client = _resolve_client() return runner.main(client=client, testing=testing, test_params=test_params) diff --git a/src/pipelines/shared/extraction/dynamic_funcs.py b/src/pipelines/shared/extraction/dynamic_funcs.py index 7afe254..46c1a21 100644 --- a/src/pipelines/shared/extraction/dynamic_funcs.py +++ b/src/pipelines/shared/extraction/dynamic_funcs.py @@ -4,7 +4,12 @@ from typing import TYPE_CHECKING, Optional from src.pipelines.shared.prompts import prompt_calls from src.pipelines.shared.preprocessing import preprocessing_funcs -from src.utils import string_utils, timing_utils +from src.utils import ( + instrumentation, + instrumentation_context, + string_utils, + timing_utils, +) from src.constants.constants import Constants from src import config from src.prompts import prompt_templates @@ -197,25 +202,37 @@ def dynamic_assignment( # 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( - pages_dict=pages_dict, - text_dict=text_dict, - exhibit=exhibit, - current_page_num=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, + row_id = updated.get("REIMB_ID") or f"auto_{id(answer_dict):x}" + with instrumentation_context.instr_scope(row_id=row_id, page_num=str(page_num)): + exhibit_text_simplified = preprocessing_funcs.simplify_exhibit( + pages_dict=pages_dict, + text_dict=text_dict, + exhibit=exhibit, + current_page_num=page_num, ) - updated.update(dynamic_field_answer) + 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, + ) + updated.update(dynamic_field_answer) + instrumentation.log( + "row_count", + filename=filename, + stage="dynamic_assignment_single", + n_in=1, + n_out=1, + extra_json=( + f'{{"fields_assigned": {[f.field_name for f in dynamic_reimbursement_fields.fields]}}}' + ), + ) return updated max_workers = min(len(reimbursement_primary_answers), 8) @@ -223,7 +240,11 @@ def dynamic_assignment( "dynamic_assignment.parallel_processing", context=filename ): with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - results = list(executor.map(_assign_single, reimbursement_primary_answers)) + results = list( + instrumentation_context.map_with_context( + executor, _assign_single, reimbursement_primary_answers + ) + ) logging.debug( f"Dynamic assignment complete: {len(results)} rows processed with {len(dynamic_reimbursement_fields.fields)} dynamic fields - {filename}" diff --git a/src/pipelines/shared/extraction/exhibit_funcs.py b/src/pipelines/shared/extraction/exhibit_funcs.py index a370aa7..84e1cde 100644 --- a/src/pipelines/shared/extraction/exhibit_funcs.py +++ b/src/pipelines/shared/extraction/exhibit_funcs.py @@ -35,6 +35,7 @@ if TYPE_CHECKING: from src.pipelines.shared.extraction.page_funcs import Page from src.pipelines.shared.prompts import prompt_calls +from src.utils import instrumentation import logging @@ -182,6 +183,25 @@ class Exhibit: self._chunks = _chunk_exhibit_text(self.exhibit_text) # Compute embeddings eagerly self._chunks = _compute_chunk_embeddings(self._chunks) + if instrumentation.is_enabled(): + sizes = sorted(len(c.text) for c in self._chunks) + n = len(sizes) + p50 = sizes[n // 2] if n else 0 + p95 = sizes[int(n * 0.95)] if n else 0 + instrumentation.log( + "chunking_done", + exhibit_header=( + self.exhibit_header[:80] if self.exhibit_header else "" + ), + exhibit_page=self.exhibit_page, + extra_json=( + f'{{"total": {n}, ' + f'"tables": {sum(1 for c in self._chunks if c.is_table)}, ' + f'"sub_chunks_parents": {sum(1 for c in self._chunks if c.has_sub_chunks())}, ' + f'"size_p50": {p50}, ' + f'"size_p95": {p95}}}' + ), + ) return self._chunks @property diff --git a/src/pipelines/shared/extraction/one_to_n_funcs.py b/src/pipelines/shared/extraction/one_to_n_funcs.py index 25397ed..bf5d82d 100644 --- a/src/pipelines/shared/extraction/one_to_n_funcs.py +++ b/src/pipelines/shared/extraction/one_to_n_funcs.py @@ -4,7 +4,13 @@ import pandas as pd from src.pipelines.shared.extraction import dynamic_funcs from src.pipelines.shared.prompts import prompt_calls from src.pipelines.shared.postprocessing import aarete_derived, postprocessing_funcs -from src.utils import llm_utils, string_utils, timing_utils +from src.utils import ( + instrumentation, + instrumentation_context, + llm_utils, + string_utils, + timing_utils, +) import src.prompts.prompt_templates as prompt_templates from src.constants.constants import Constants from src import config @@ -24,8 +30,8 @@ def dynamic_code_assignment( max_workers=min(len(reimbursement_level_answers), 5) ) as executor: futures = [ - executor.submit( - dynamic_code_assignment_single_row, answer_dict, filename + instrumentation_context.submit_with_context( + executor, dynamic_code_assignment_single_row, answer_dict, filename ) for answer_dict in reimbursement_level_answers ] @@ -191,8 +197,7 @@ def clean_reimbursement_primary( 1. Normalize SERVICE_TERM and REIMB_TERM to strings (they should never be lists per-row) 2. Split compound reimbursements 3. Filter out any non-reimbursement content from split answers - 4. Deduplicate against previously seen (SERVICE_TERM, REIMB_TERM) pairs - 5. Apply exhibit lesser-of statement if it exists + 4. Apply exhibit lesser-of statement if it exists Args: reimbursement_primary_answers (list[str]): List of dictionaries containing reimbursement details. @@ -205,12 +210,18 @@ def clean_reimbursement_primary( # If needed, add Reimbursement Splitting here - # If needed, add deduplication here - # Filter out any non-reimbursement content from split answers + _n_in = len(reimbursement_primary_answers) reimbursement_primary_answers = filter_services_without_reimbursements( reimbursement_primary_answers, filename ) + instrumentation.log( + "row_count", + filename=filename, + stage="clean_reimbursement_primary", + n_in=_n_in, + n_out=len(reimbursement_primary_answers), + ) return reimbursement_primary_answers else: @@ -318,7 +329,8 @@ def carveout_and_special_case( max_workers=min(len(reimbursement_breakout_answers), 5) ) as executor: futures = [ - executor.submit( + instrumentation_context.submit_with_context( + executor, process_single_carveout, answer_dict.copy(), carveout_definitions, @@ -370,8 +382,12 @@ def methodology_breakout( max_workers=min(len(reimbursement_primary_answers), 5) ) as executor: futures = [ - executor.submit( - methodology_breakout_single_row, answer_dict, constants, filename + instrumentation_context.submit_with_context( + executor, + methodology_breakout_single_row, + answer_dict, + constants, + filename, ) for answer_dict in reimbursement_primary_answers ] @@ -447,6 +463,13 @@ def methodology_breakout_single_row( results.append( {**answer_dict, **methodology_breakout_dict, **secondary_breakout_dict} ) + instrumentation.log( + "row_count", + filename=filename, + stage="methodology_breakout_single_row", + n_in=1, + n_out=len(results), + ) return results return [answer_dict] # Return original in a list if no reimbursement method found @@ -596,7 +619,8 @@ def special_case_breakout( max_workers=min(len(special_case_answers), 5) ) as executor: futures = [ - executor.submit( + instrumentation_context.submit_with_context( + executor, process_single_special_case_breakout, answer_dict.copy(), special_case_fields, @@ -660,7 +684,9 @@ def filter_services_without_reimbursements( ) as executor: # Create futures with index tracking future_to_index = { - executor.submit(validate_single_reimbursement, answer_dict, filename): i + instrumentation_context.submit_with_context( + executor, validate_single_reimbursement, answer_dict, filename + ): i for i, answer_dict in enumerate(reimbursement_primary_answers) } @@ -685,6 +711,13 @@ def filter_services_without_reimbursements( f"Filtered {filtered_count} services without clear reimbursement methodologies in {filename}." ) + instrumentation.log( + "row_count", + filename=filename, + stage="filter_services_without_reimbursements", + n_in=len(reimbursement_primary_answers), + n_out=len(valid_reimbursements), + ) return valid_reimbursements @@ -821,7 +854,8 @@ def get_lob_relationship(answer_dicts, exhibit_text, filename): max_workers=min(len(answer_dicts), 5) ) as executor: futures = [ - executor.submit( + instrumentation_context.submit_with_context( + executor, process_single_lob_relationship, answer_dict.copy(), exhibit_text, @@ -1086,4 +1120,11 @@ def split_service_terms(one_to_n_results: pd.DataFrame, filename: str) -> pd.Dat result_df = pd.DataFrame(new_rows).reset_index(drop=True) + instrumentation.log( + "row_count", + filename=filename, + stage="split_service_terms", + n_in=len(one_to_n_results), + n_out=len(new_rows), + ) return result_df diff --git a/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py b/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py index bbd8e04..4f1df43 100644 --- a/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py +++ b/src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py @@ -5,10 +5,12 @@ Follows the same pattern as hybrid_smart_chunking_funcs.py (for 1:1 fields). Provides chunking, embedding, and hybrid retrieval functions for exhibit text analysis. """ +import json import re import logging from typing import Optional, List +from src.utils import instrumentation from src.utils.rag_utils import get_embeddings, HSC_CONFIG, cosine_similarity from src.utils.string_utils import keyword_match from src.constants.regex_patterns import SECTION_HEADER_PATTERNS @@ -792,4 +794,17 @@ def search_reimbursement_chunks( result_chunks = [c for c in chunks if c.chunk_id in all_selected_ids] result_chunks.sort(key=lambda c: int(c.chunk_id)) + instrumentation.log( + "retrieval_done", + extra_json=json.dumps( + { + "total_in": len(chunks), + "keyword_hits": len(keyword_matched_ids), + "semantic_hits": len(semantic_matched_ids), + "table_auto": len(table_chunk_ids), + "selected": len(result_chunks), + "selected_ids": [c.chunk_id for c in result_chunks], + } + ), + ) return result_chunks diff --git a/src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py b/src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py index de1b363..bd8e399 100644 --- a/src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py +++ b/src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py @@ -14,7 +14,13 @@ from typing import List, Tuple, Dict from langchain_core.documents import Document from langchain_text_splitters import RecursiveCharacterTextSplitter -from src.utils import llm_utils, string_utils, timing_utils, rag_utils +from src.utils import ( + instrumentation_context, + llm_utils, + string_utils, + timing_utils, + rag_utils, +) import src.pipelines.shared.preprocessing.preprocessing_funcs as preprocessing_funcs import src.pipelines.shared.preprocessing.hybrid_smart_chunking_preprocessing as hybrid_smart_chunking_preprocessing from src import config @@ -316,7 +322,8 @@ def run_hybrid_smart_chunked_fields( max_workers=min(len(fields_to_process), 20) ) as executor: futures = [ - executor.submit( + instrumentation_context.submit_with_context( + executor, prompt_hsc_single_field, field, retriever, diff --git a/src/testbed/instrumentation_metrics.py b/src/testbed/instrumentation_metrics.py new file mode 100644 index 0000000..f5b137d --- /dev/null +++ b/src/testbed/instrumentation_metrics.py @@ -0,0 +1,1383 @@ +"""Instrumentation CSV analyzer — single, batch, and compare modes. + +Usage: + uv run python -m src.testbed.instrumentation_metrics single \\ + --csv path/to/BATCH-INSTRUMENTATION.csv --out-dir my_analysis/ + + uv run python -m src.testbed.instrumentation_metrics batch \\ + --csv path/to/BATCH-INSTRUMENTATION.csv --out-dir my_analysis/ + + uv run python -m src.testbed.instrumentation_metrics compare \\ + --csv-a path/to/run_a/BATCH-INSTRUMENTATION.csv \\ + --csv-b path/to/run_b/BATCH-INSTRUMENTATION.csv \\ + --label-a mar13 --label-b apr08 \\ + --out-dir my_analysis/compare/ +""" + +import argparse +import logging +import os +import sys +from pathlib import Path +from typing import Union + +import pandas as pd + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +_REQUIRED_COLUMNS = {"event", "filename", "segment", "usage_label", "cost_usd"} + + +def load_instrumentation_csv(path: Union[str, Path]) -> pd.DataFrame: + """Read an instrumentation CSV and do basic type coercion. + + Also re-resolves the ``segment`` column from ``usage_label`` using the + current ``USAGE_LABEL_TO_SEGMENT`` mapping. This makes the analyzer + forward-compatible for older CSVs whose ``segment`` column was written + against an outdated mapping: the CSV's recorded segment is only used as + a fallback when the label is not in the current mapping AND no cache_warming + prefix is present. + """ + from src.utils import instrumentation as _instr + + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Instrumentation CSV not found: {path}") + df = pd.read_csv(path, low_memory=False) + missing = _REQUIRED_COLUMNS - set(df.columns) + if missing: + raise ValueError(f"Instrumentation CSV missing required columns: {missing}") + for col in ( + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost_usd", + "n_in", + "n_out", + "backoff_sec", + ): + if col in df.columns: + df[col] = pd.to_numeric(df[col], errors="coerce").fillna(0) + + # Re-resolve segment from usage_label using current mapping. Preserve the + # previously-recorded segment as fallback (e.g. scope-tagged events that + # have no usage_label, like row_count or file_start). Also force + # segment="cache_warming" for any row whose filename starts with + # "cache_warming_" — that's how the pipeline tags priming calls. + if "usage_label" in df.columns: + recorded_segment = df.get("segment") + + def _resolve(r): + fn = r.get("filename") + if isinstance(fn, str) and fn.startswith("cache_warming_"): + return "cache_warming" + label = ( + r.get("usage_label") if isinstance(r.get("usage_label"), str) else None + ) + fallback = ( + recorded_segment.loc[r.name] + if recorded_segment is not None + and isinstance(recorded_segment.loc[r.name], str) + else None + ) + return _instr.resolve_segment(label, fallback) + + df["segment"] = df.apply(_resolve, axis=1) + return df + + +_PROMPT_LABEL_LEAKS = {"date_fix", "derive_term_date", "all_files"} + + +def _is_real_contract(filename: str | float) -> bool: + """True iff this filename is an actual contract (not cache_warming_* or a + known prompt-label leak). Mirrors the filtering we apply to USAGE CSVs.""" + if not isinstance(filename, str) or not filename: + return False + if filename.startswith("cache_warming_"): + return False + if filename in _PROMPT_LABEL_LEAKS: + return False + return True + + +def classify_contracts(df: pd.DataFrame) -> dict[str, set[str]]: + """Partition filenames into successful and errored sets. + + Filters out cache_warming_* pseudo-filenames and prompt-label leaks + (date_fix, derive_term_date, all_files) before classification — these + are not real contracts. + + A contract is errored if: + - It has any ``llm_error`` event, OR + - It has a ``file_start`` event but no matching ``file_end`` event. + """ + all_files = {f for f in df["filename"].dropna().unique() if _is_real_contract(f)} + + errored_via_llm_error = { + f + for f in df.loc[df["event"] == "llm_error", "filename"].dropna().unique() + if _is_real_contract(f) + } + + files_with_start = { + f + for f in df.loc[df["event"] == "file_start", "filename"].dropna().unique() + if _is_real_contract(f) + } + files_with_end = { + f + for f in df.loc[df["event"] == "file_end", "filename"].dropna().unique() + if _is_real_contract(f) + } + errored_missing_end = files_with_start - files_with_end + + errored = errored_via_llm_error | errored_missing_end + successful = all_files - errored + + return {"successful": successful, "errored": errored} + + +def by_segment(df: pd.DataFrame) -> pd.DataFrame: + """Aggregate llm_call events by segment.""" + llm = df[df["event"] == "llm_call"].copy() + if llm.empty: + return pd.DataFrame( + columns=["segment", "n_calls", "cost_usd", "cache_hit_rate"] + ) + total_cost = llm["cost_usd"].sum() + grp = llm.groupby("segment", dropna=False) + result = grp.agg( + n_calls=("event", "count"), + cost_usd=("cost_usd", "sum"), + input_tokens=("input_tokens", "sum"), + cache_creation_tokens=("cache_creation_tokens", "sum"), + cache_read_tokens=("cache_read_tokens", "sum"), + mean_input_tokens=("input_tokens", "mean"), + ).reset_index() + result["pct_of_total"] = ( + result["cost_usd"] / total_cost * 100 if total_cost > 0 else 0.0 + ) + denom = ( + result["input_tokens"] + + result["cache_creation_tokens"] + + result["cache_read_tokens"] + ) + result["cache_hit_rate"] = result["cache_read_tokens"] / denom.replace( + 0, float("nan") + ) + return result + + +def by_label(df: pd.DataFrame) -> pd.DataFrame: + """Aggregate llm_call events by usage_label.""" + llm = df[df["event"] == "llm_call"].copy() + if llm.empty: + return pd.DataFrame( + columns=["usage_label", "segment", "n_calls", "cost_usd", "cache_hit_rate"] + ) + grp = llm.groupby(["usage_label", "segment"], dropna=False) + result = grp.agg( + n_calls=("event", "count"), + cost_usd=("cost_usd", "sum"), + input_tokens=("input_tokens", "sum"), + cache_creation_tokens=("cache_creation_tokens", "sum"), + cache_read_tokens=("cache_read_tokens", "sum"), + ).reset_index() + denom = ( + result["input_tokens"] + + result["cache_creation_tokens"] + + result["cache_read_tokens"] + ) + result["cache_hit_rate"] = result["cache_read_tokens"] / denom.replace( + 0, float("nan") + ) + return result + + +def cache_breakdown(df: pd.DataFrame) -> pd.DataFrame: + """Aggregate llm_call events by (usage_label, cache_miss_reason).""" + llm = df[df["event"] == "llm_call"].copy() + if llm.empty or "cache_miss_reason" not in llm.columns: + return pd.DataFrame( + columns=["usage_label", "cache_miss_reason", "n", "cost_usd"] + ) + grp = llm.groupby(["usage_label", "cache_miss_reason"], dropna=False) + result = grp.agg(n=("event", "count"), cost_usd=("cost_usd", "sum")).reset_index() + return result + + +def _cache_hit_rate(df: pd.DataFrame) -> float: + """Overall cache hit rate across llm_call events.""" + llm = df[df["event"] == "llm_call"] + if llm.empty: + return 0.0 + num = llm["cache_read_tokens"].sum() + denom = ( + llm["input_tokens"].sum() + + llm["cache_creation_tokens"].sum() + + llm["cache_read_tokens"].sum() + ) + return float(num / denom) if denom > 0 else 0.0 + + +def _write(df: pd.DataFrame, out_dir: Path, filename: str) -> None: + df.to_csv(out_dir / filename, index=False) + + +def _print_table(title: str, df: pd.DataFrame, cols: list[str]) -> None: + """Print a compact aligned table to stdout.""" + print(f"\n--- {title} ---") + if df.empty: + print(" (no data)") + return + sub = df[cols].copy() + for col in sub.select_dtypes("float").columns: + sub[col] = sub[col].map(lambda v: f"{v:.4f}" if pd.notna(v) else "") + print(sub.to_string(index=False)) + + +# --------------------------------------------------------------------------- +# Single-file mode +# --------------------------------------------------------------------------- + + +def run_single(csv_path: Union[str, Path], out_dir: Union[str, Path]) -> None: + """Summarize a single contract's (or a small run's) instrumentation trace. + + Outputs written to *out_dir*: + summary.csv, by_segment.csv, by_label.csv, cache_miss_breakdown.csv, + by_exhibit.csv, row_count_funnel.csv, retries_errors.csv, + top10_calls_by_cost.csv + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + df = load_instrumentation_csv(csv_path) + llm = df[df["event"] == "llm_call"] + + # --- summary.csv --- + n_inmem = (df["event"] == "llm_call_inmem_hit").sum() + n_retries = (df["event"] == "llm_retry").sum() + n_errors = (df["event"] == "llm_error").sum() + + file_start_rows = df[df["event"] == "file_start"] + file_end_rows = df[df["event"] == "file_end"] + wall_clock_sec = None + if "ts_rel_sec" in df.columns and not file_end_rows.empty: + wall_clock_sec = float(df["ts_rel_sec"].max()) + + total_cost = float(llm["cost_usd"].sum()) + total_tokens = ( + int(llm["total_tokens"].sum()) if "total_tokens" in llm.columns else 0 + ) + + summary = pd.DataFrame( + [ + { + "total_cost": total_cost, + "total_tokens": total_tokens, + "n_llm_calls": len(llm), + "n_inmem_hits": int(n_inmem), + "cache_hit_rate": _cache_hit_rate(df), + "n_retries": int(n_retries), + "n_errors": int(n_errors), + "wall_clock_sec": wall_clock_sec, + } + ] + ) + _write(summary, out_dir, "summary.csv") + + # --- by_segment.csv --- + seg_df = by_segment(df) + _write(seg_df, out_dir, "by_segment.csv") + + # --- by_label.csv --- + label_df = by_label(df) + if "cache_miss_reason" in df.columns: + dominant = ( + df[df["event"] == "llm_call"] + .groupby("usage_label")["cache_miss_reason"] + .agg(lambda s: s.value_counts().idxmax() if not s.empty else None) + .reset_index() + .rename(columns={"cache_miss_reason": "cache_miss_reason_dominant"}) + ) + label_df = label_df.merge(dominant, on="usage_label", how="left") + _write(label_df, out_dir, "by_label.csv") + + # --- cache_miss_breakdown.csv --- + miss_df = cache_breakdown(df) + _write(miss_df, out_dir, "cache_miss_breakdown.csv") + + # --- by_exhibit.csv --- + exhibit_cols = ["exhibit_idx", "exhibit_header"] + available_exhibit_cols = [c for c in exhibit_cols if c in df.columns] + if available_exhibit_cols and not llm.empty: + grp_cols = available_exhibit_cols + n_rows_out_df = ( + df[df["event"] == "row_count"] + .groupby(available_exhibit_cols[:1], dropna=False)["n_out"] + .max() + .reset_index() + .rename(columns={"n_out": "n_rows_out"}) + if "n_out" in df.columns + else pd.DataFrame(columns=available_exhibit_cols[:1] + ["n_rows_out"]) + ) + exhibit_df = ( + llm.groupby(grp_cols, dropna=False) + .agg(n_calls=("event", "count"), cost_usd=("cost_usd", "sum")) + .reset_index() + ) + if not n_rows_out_df.empty: + exhibit_df = exhibit_df.merge( + n_rows_out_df, on=available_exhibit_cols[:1], how="left" + ) + _write(exhibit_df, out_dir, "by_exhibit.csv") + else: + _write( + pd.DataFrame( + columns=[ + "exhibit_idx", + "exhibit_header", + "n_calls", + "cost_usd", + "n_rows_out", + ] + ), + out_dir, + "by_exhibit.csv", + ) + + # --- row_count_funnel.csv --- + rc = df[df["event"] == "row_count"].copy() + if not rc.empty and "stage" in rc.columns: + # Preserve first-seen order of stages rather than alphabetizing. + stage_order = rc["stage"].drop_duplicates().tolist() + funnel = ( + rc.groupby("stage", dropna=False, sort=False) + .agg(n_in=("n_in", "sum"), n_out=("n_out", "sum")) + .reset_index() + ) + funnel["stage"] = pd.Categorical( + funnel["stage"], categories=stage_order, ordered=True + ) + funnel = funnel.sort_values("stage").reset_index(drop=True) + funnel["stage"] = funnel["stage"].astype(str) + denom = funnel["n_in"].replace(0, float("nan")) + funnel["delta_pct"] = (funnel["n_out"] - funnel["n_in"]) / denom * 100 + _write(funnel, out_dir, "row_count_funnel.csv") + else: + _write( + pd.DataFrame(columns=["stage", "n_in", "n_out", "delta_pct"]), + out_dir, + "row_count_funnel.csv", + ) + + # --- retries_errors.csv --- + retry_df = df[df["event"].isin(["llm_retry", "llm_error"])].copy() + if not retry_df.empty and "usage_label" in retry_df.columns: + re_grp = retry_df.groupby( + ( + ["usage_label", "error_class"] + if "error_class" in retry_df.columns + else ["usage_label"] + ), + dropna=False, + ) + re_agg = re_grp.agg( + n_retries=("event", lambda s: (s == "llm_retry").sum()), + n_errors=("event", lambda s: (s == "llm_error").sum()), + ).reset_index() + if "backoff_sec" in retry_df.columns: + backoff_sum = ( + retry_df[retry_df["event"] == "llm_retry"] + .groupby( + ( + ["usage_label", "error_class"] + if "error_class" in retry_df.columns + else ["usage_label"] + ), + dropna=False, + )["backoff_sec"] + .sum() + .reset_index() + .rename(columns={"backoff_sec": "total_backoff_sec"}) + ) + merge_cols = ( + ["usage_label", "error_class"] + if "error_class" in retry_df.columns + else ["usage_label"] + ) + re_agg = re_agg.merge(backoff_sum, on=merge_cols, how="left") + _write(re_agg, out_dir, "retries_errors.csv") + else: + _write( + pd.DataFrame( + columns=[ + "usage_label", + "error_class", + "n_retries", + "n_errors", + "total_backoff_sec", + ] + ), + out_dir, + "retries_errors.csv", + ) + + # --- top10_calls_by_cost.csv --- + top10_cols = [ + "seq", + "ts_rel_sec", + "segment", + "usage_label", + "cost_usd", + "input_tokens", + ] + available_top10 = [c for c in top10_cols if c in llm.columns] + top10 = llm.nlargest(10, "cost_usd")[available_top10].reset_index(drop=True) + _write(top10, out_dir, "top10_calls_by_cost.csv") + + # --- stdout summary --- + print(f"\n=== SINGLE-FILE INSTRUMENTATION SUMMARY ===") + print(f" Source: {csv_path}") + print(f" Total cost: ${total_cost:.4f}") + print(f" LLM calls: {len(llm)}") + print(f" In-mem hits: {int(n_inmem)}") + print(f" Cache hit rate: {_cache_hit_rate(df):.1%}") + print(f" Retries: {int(n_retries)}") + print(f" Errors: {int(n_errors)}") + if wall_clock_sec is not None: + print(f" Wall clock: {wall_clock_sec:.1f}s") + + seg_print_cols = [ + c + for c in ["segment", "n_calls", "cost_usd", "pct_of_total", "cache_hit_rate"] + if c in seg_df.columns + ] + _print_table( + "Cost by Segment", + seg_df.sort_values("cost_usd", ascending=False), + seg_print_cols, + ) + + print(f"\nOutputs written to: {out_dir}") + logger.info("run_single complete: %s", out_dir) + + +# --------------------------------------------------------------------------- +# Batch mode +# --------------------------------------------------------------------------- + + +def run_batch(csv_path: Union[str, Path], out_dir: Union[str, Path]) -> None: + """Summarize a full batch run's instrumentation CSV. + + Outputs written to *out_dir*: + batch_summary.csv, per_contract.csv, by_segment.csv, by_label.csv, + top10_offenders.csv, errored_contracts.csv, cache_lever_estimate.csv + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + df = load_instrumentation_csv(csv_path) + classes = classify_contracts(df) + successful = classes["successful"] + errored = classes["errored"] + + all_files = successful | errored + n_attempted = len(all_files) + n_successful = len(successful) + n_errored = len(errored) + + llm = df[df["event"] == "llm_call"] + + batch_hit_rate = _cache_hit_rate(df) + batch_total_cost = float(llm["cost_usd"].sum()) + + # Per-contract cost + per_file_cost = llm.groupby("filename")["cost_usd"].sum() + successful_costs = per_file_cost.reindex(list(successful)).dropna() + mean_cost = float(successful_costs.mean()) if not successful_costs.empty else 0.0 + median_cost = ( + float(successful_costs.median()) if not successful_costs.empty else 0.0 + ) + p95_cost = ( + float(successful_costs.quantile(0.95)) if not successful_costs.empty else 0.0 + ) + + # --- batch_summary.csv --- + batch_summary = pd.DataFrame( + [ + { + "n_contracts_attempted": n_attempted, + "n_successful": n_successful, + "n_errored": n_errored, + "batch_total_cost": batch_total_cost, + "mean_cost_per_successful": mean_cost, + "median_cost_per_successful": median_cost, + "p95_cost": p95_cost, + "cache_hit_rate": batch_hit_rate, + } + ] + ) + _write(batch_summary, out_dir, "batch_summary.csv") + + # --- per_contract.csv --- + per_file_calls = llm.groupby("filename")["event"].count().rename("n_calls") + per_file_retries = ( + df[df["event"] == "llm_retry"] + .groupby("filename")["event"] + .count() + .rename("n_retries") + ) + per_file_errors = ( + df[df["event"] == "llm_error"] + .groupby("filename")["event"] + .count() + .rename("n_errors") + ) + + def _file_hit_rate(sub): + num = sub["cache_read_tokens"].sum() + denom = ( + sub["input_tokens"].sum() + + sub["cache_creation_tokens"].sum() + + sub["cache_read_tokens"].sum() + ) + return float(num / denom) if denom > 0 else 0.0 + + per_file_hit = ( + llm.groupby("filename").apply(_file_hit_rate).rename("cache_hit_rate") + ) + + per_contract = ( + per_file_cost.rename("total_cost") + .to_frame() + .join(per_file_calls, how="outer") + .join(per_file_hit, how="outer") + .join(per_file_retries, how="outer") + .join(per_file_errors, how="outer") + .fillna(0) + .reset_index() + ) + # Restrict to real contracts (exclude cache_warming_* and prompt-label leaks) + per_contract = per_contract[per_contract["filename"].apply(_is_real_contract)] + per_contract["status"] = per_contract["filename"].apply( + lambda fn: "errored" if fn in errored else "successful" + ) + + # cost_by_segment_spread: segment -> cost as JSON-ish string + if not llm.empty and "segment" in llm.columns: + seg_spread = ( + llm.groupby(["filename", "segment"])["cost_usd"] + .sum() + .unstack(fill_value=0) + .reset_index() + ) + per_contract = per_contract.merge(seg_spread, on="filename", how="left") + + _write(per_contract, out_dir, "per_contract.csv") + + # --- by_segment.csv --- + seg_df = by_segment(df) + if not seg_df.empty: + seg_df = seg_df.rename(columns={"cost_usd": "total_cost"}) + total_batch = seg_df["total_cost"].sum() + seg_df["pct_of_batch"] = ( + seg_df["total_cost"] / total_batch * 100 if total_batch > 0 else 0.0 + ) + # mean cost per contract (successful) + seg_df["mean_cost_per_contract"] = ( + seg_df["total_cost"] / n_successful if n_successful > 0 else 0.0 + ) + _write(seg_df, out_dir, "by_segment.csv") + + # --- by_label.csv --- + label_df = by_label(df) + if not label_df.empty and "cache_miss_reason" in df.columns: + miss_summary = ( + df[df["event"] == "llm_call"] + .groupby("usage_label")["cache_miss_reason"] + .value_counts() + .rename("n") + .reset_index() + .pivot_table( + index="usage_label", + columns="cache_miss_reason", + values="n", + fill_value=0, + ) + .reset_index() + ) + miss_summary.columns.name = None + miss_summary = miss_summary.add_prefix("miss_") + miss_summary = miss_summary.rename(columns={"miss_usage_label": "usage_label"}) + label_df = label_df.merge(miss_summary, on="usage_label", how="left") + if n_successful > 0: + label_df["n_calls_per_contract_mean"] = label_df["n_calls"] / n_successful + _write(label_df, out_dir, "by_label.csv") + + # --- top10_offenders.csv --- + top10 = per_contract.nlargest(10, "total_cost")[ + ["filename", "total_cost", "n_calls", "status"] + ].reset_index(drop=True) + top10["error_flag"] = top10["status"] == "errored" + _write(top10, out_dir, "top10_offenders.csv") + + # --- errored_contracts.csv --- + if n_errored > 0: + error_events = df[df["event"] == "llm_error"] + first_error = ( + error_events.sort_values( + "seq" if "seq" in error_events.columns else error_events.columns[0] + ) + .groupby("filename") + .first() + .reset_index()[ + ["filename"] + + [c for c in ["segment", "usage_label"] if c in error_events.columns] + ] + ) + first_error = first_error.rename( + columns={ + "segment": "first_error_segment", + "usage_label": "first_error_label", + } + ) + wasted = ( + llm[llm["filename"].isin(errored)] + .groupby("filename")["cost_usd"] + .sum() + .rename("wasted_cost_usd") + .reset_index() + ) + n_retries_before = ( + df[df["event"] == "llm_retry"] + .groupby("filename")["event"] + .count() + .rename("n_retries_before_error") + .reset_index() + ) + errored_df = first_error.merge(wasted, on="filename", how="left").merge( + n_retries_before, on="filename", how="left" + ) + _write(errored_df, out_dir, "errored_contracts.csv") + else: + _write( + pd.DataFrame( + columns=[ + "filename", + "first_error_segment", + "first_error_label", + "wasted_cost_usd", + "n_retries_before_error", + ] + ), + out_dir, + "errored_contracts.csv", + ) + + # --- cache_lever_estimate.csv --- + if "cache_miss_reason" in df.columns: + not_attempted = ( + llm[llm["cache_miss_reason"] == "not_attempted"] + .groupby(["usage_label", "segment"]) + .agg(cost_at_risk=("cost_usd", "sum"), n=("event", "count")) + .reset_index() + ) + # Potential savings estimate: ~90% of cost (cache read is ~10% of input price) + not_attempted["potential_savings_if_cached"] = ( + not_attempted["cost_at_risk"] * 0.9 + ) + # Filter to labels where not_attempted dominates + total_by_label = ( + llm.groupby("usage_label")["event"].count().rename("total_calls") + ) + not_attempted = not_attempted.join(total_by_label, on="usage_label", how="left") + not_attempted["not_attempted_share"] = ( + not_attempted["n"] / not_attempted["total_calls"] + ) + dominant = not_attempted[not_attempted["not_attempted_share"] > 0.5][ + ["usage_label", "segment", "cost_at_risk", "potential_savings_if_cached"] + ].rename(columns={"usage_label": "label"}) + _write(dominant, out_dir, "cache_lever_estimate.csv") + else: + _write( + pd.DataFrame( + columns=[ + "label", + "segment", + "cost_at_risk", + "potential_savings_if_cached", + ] + ), + out_dir, + "cache_lever_estimate.csv", + ) + + # --- stdout summary --- + print(f"\n=== BATCH INSTRUMENTATION SUMMARY ===") + print(f" Source: {csv_path}") + print(f" Contracts attempted: {n_attempted}") + print(f" Successful: {n_successful}") + print(f" Errored: {n_errored}") + print(f" Batch total cost: ${batch_total_cost:.4f}") + print(f" Mean cost/success: ${mean_cost:.4f}") + print(f" Median cost/success: ${median_cost:.4f}") + print(f" P95 cost/success: ${p95_cost:.4f}") + print(f" Cache hit rate: {batch_hit_rate:.1%}") + + seg_print = ( + seg_df.rename(columns={"total_cost": "cost_usd"}) + if "total_cost" in seg_df.columns + else seg_df + ) + seg_print_cols = [ + c + for c in ["segment", "cost_usd", "pct_of_batch", "cache_hit_rate"] + if c in seg_print.columns + ] + _print_table( + "Cost by Segment", + ( + seg_print.sort_values("cost_usd", ascending=False) + if not seg_print.empty + else seg_print + ), + seg_print_cols, + ) + + print(f"\nOutputs written to: {out_dir}") + logger.info("run_batch complete: %s", out_dir) + + +# --------------------------------------------------------------------------- +# Compare mode +# --------------------------------------------------------------------------- + + +def run_compare( + csv_a: Union[str, Path], + csv_b: Union[str, Path], + out_dir: Union[str, Path], + label_a: str = "a", + label_b: str = "b", +) -> None: + """Compare two batch instrumentation CSVs with robust contract-set handling. + + Partitions contracts into four disjoint sets: + successful_in_both, successful_in_a_only, successful_in_b_only, + errored_in_both. + + All cost / hit-rate diffs are computed ONLY on successful_in_both. + + Outputs written to *out_dir*: + corpus_diff.csv, batch_summary_compare.csv, by_segment_diff.csv, + by_label_diff.csv, regression_watchlist.csv, corpus_change_summary.csv + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + df_a = load_instrumentation_csv(csv_a) + df_b = load_instrumentation_csv(csv_b) + + classes_a = classify_contracts(df_a) + classes_b = classify_contracts(df_b) + + succ_a = classes_a["successful"] + succ_b = classes_b["successful"] + err_a = classes_a["errored"] + err_b = classes_b["errored"] + + all_a = succ_a | err_a + all_b = succ_b | err_b + + successful_in_both = succ_a & succ_b + errored_in_both = err_a & err_b + # in_a but not successful in b (absent or errored) + successful_in_a_only = succ_a - succ_b + # in_b but not successful in a (absent or errored) + successful_in_b_only = succ_b - succ_a + + # Regressed: was successful in a, errored in b (subset of a_only) + regressed = succ_a & err_b + # Recovered: was errored in a, successful in b + recovered = err_a & succ_b + # New in b: not in a at all + new_in_b = all_b - all_a + # Dropped from a: not in b at all + dropped_from_a = all_a - all_b + + def _category(fn): + if fn in successful_in_both: + return "in_both" + if fn in regressed: + return "regressed" + if fn in recovered: + return "recovered" + if fn in new_in_b: + return "only_b" + if fn in dropped_from_a: + return "only_a" + if fn in errored_in_both: + return "errored_in_both" + # fallback for errored in one + if fn in succ_a: + return "only_a" + if fn in succ_b: + return "only_b" + return "errored_in_both" + + all_files = all_a | all_b + llm_a = df_a[df_a["event"] == "llm_call"] + llm_b = df_b[df_b["event"] == "llm_call"] + + cost_a_by_file = llm_a.groupby("filename")["cost_usd"].sum() + cost_b_by_file = llm_b.groupby("filename")["cost_usd"].sum() + + def _status(fn, succ, err): + if fn in succ: + return "successful" + if fn in err: + return "errored" + return "absent" + + corpus_rows = [] + for fn in sorted(all_files): + corpus_rows.append( + { + "filename": fn, + f"status_{label_a}": _status(fn, succ_a, err_a), + f"status_{label_b}": _status(fn, succ_b, err_b), + f"cost_{label_a}": float(cost_a_by_file.get(fn, 0.0)), + f"cost_{label_b}": float(cost_b_by_file.get(fn, 0.0)), + "category": _category(fn), + } + ) + corpus_diff = pd.DataFrame(corpus_rows) + # rename to canonical column names + corpus_diff = corpus_diff.rename( + columns={ + f"status_{label_a}": "status_a", + f"status_{label_b}": "status_b", + f"cost_{label_a}": "cost_a", + f"cost_{label_b}": "cost_b", + } + ) + _write(corpus_diff, out_dir, "corpus_diff.csv") + + # --- corpus_change_summary.csv --- + corpus_change = pd.DataFrame( + [ + { + "n_in_both": len(successful_in_both), + "n_regressed": len(regressed), + "n_recovered": len(recovered), + "n_new_in_b": len(new_in_b), + "n_dropped_from_a": len(dropped_from_a), + } + ] + ) + _write(corpus_change, out_dir, "corpus_change_summary.csv") + + # Apples-to-apples slice + df_a_both = df_a[df_a["filename"].isin(successful_in_both)] + df_b_both = df_b[df_b["filename"].isin(successful_in_both)] + llm_a_both = df_a_both[df_a_both["event"] == "llm_call"] + llm_b_both = df_b_both[df_b_both["event"] == "llm_call"] + n_both = len(successful_in_both) + + def _run_summary(df_run, llm_run, label, succ_set, n_contracts_total): + return { + "run": label, + "n_contracts": n_contracts_total, + "n_successful": len(succ_set), + "batch_total": float(llm_run["cost_usd"].sum()), + "mean_per_success": ( + float( + llm_run[llm_run["filename"].isin(succ_set)]["cost_usd"].sum() + / len(succ_set) + ) + if succ_set + else 0.0 + ), + "cache_hit_rate": _cache_hit_rate(df_run), + } + + summary_rows = [ + _run_summary(df_a, llm_a, label_a, succ_a, len(all_a)), + _run_summary(df_b, llm_b, label_b, succ_b, len(all_b)), + ] + batch_summary_compare = pd.DataFrame(summary_rows) + _write(batch_summary_compare, out_dir, "batch_summary_compare.csv") + + # --- by_segment_diff.csv (apples-to-apples only) --- + def _seg_cost(llm_df): + if llm_df.empty: + return pd.Series(dtype=float) + return llm_df.groupby("segment")["cost_usd"].sum() + + def _seg_hit_rate(llm_df): + if llm_df.empty: + return pd.Series(dtype=float) + g = llm_df.groupby("segment") + num = g["cache_read_tokens"].sum() + denom = ( + g["input_tokens"].sum() + + g["cache_creation_tokens"].sum() + + g["cache_read_tokens"].sum() + ) + return num / denom.replace(0, float("nan")) + + seg_cost_a = _seg_cost(llm_a_both).rename("cost_a") + seg_cost_b = _seg_cost(llm_b_both).rename("cost_b") + seg_diff = pd.concat([seg_cost_a, seg_cost_b], axis=1).fillna(0).reset_index() + seg_diff.columns = ["segment", "cost_a", "cost_b"] + seg_diff["delta_abs"] = seg_diff["cost_b"] - seg_diff["cost_a"] + seg_diff["delta_pct"] = ( + seg_diff["delta_abs"] / seg_diff["cost_a"].replace(0, float("nan")) * 100 + ) + seg_diff["n_contracts_compared"] = n_both + _write(seg_diff, out_dir, "by_segment_diff.csv") + + # --- by_label_diff.csv (apples-to-apples only) --- + def _label_agg(llm_df): + if llm_df.empty: + return pd.DataFrame( + columns=["usage_label", "segment", "cost", "hit_rate", "n_calls"] + ) + g = llm_df.groupby(["usage_label", "segment"], dropna=False) + cost = g["cost_usd"].sum().rename("cost") + n_calls = g["event"].count().rename("n_calls") + num = g["cache_read_tokens"].sum() + denom = ( + g["input_tokens"].sum() + + g["cache_creation_tokens"].sum() + + g["cache_read_tokens"].sum() + ) + hit_rate = (num / denom.replace(0, float("nan"))).rename("hit_rate") + return pd.concat([cost, n_calls, hit_rate], axis=1).reset_index() + + lab_a = _label_agg(llm_a_both) + lab_b = _label_agg(llm_b_both) + + label_diff = lab_a.merge( + lab_b, + on=["usage_label", "segment"], + how="outer", + suffixes=(f"_{label_a}", f"_{label_b}"), + ).fillna(0) + label_diff = label_diff.rename( + columns={ + f"cost_{label_a}": "cost_a", + f"cost_{label_b}": "cost_b", + f"hit_rate_{label_a}": "hit_rate_a", + f"hit_rate_{label_b}": "hit_rate_b", + } + ) + label_diff["delta_abs"] = label_diff["cost_b"] - label_diff["cost_a"] + label_diff["delta_pct"] = ( + label_diff["delta_abs"] / label_diff["cost_a"].replace(0, float("nan")) * 100 + ) + label_diff["hit_rate_delta"] = label_diff["hit_rate_b"] - label_diff["hit_rate_a"] + _write(label_diff, out_dir, "by_label_diff.csv") + + # --- regression_watchlist.csv --- + # Cost grew >10% OR hit_rate dropped >10pp on apples-to-apples slice + watchlist_rows = [] + for _, row in seg_diff.iterrows(): + if pd.notna(row["delta_pct"]) and row["delta_pct"] > 10: + watchlist_rows.append( + { + "label_or_segment": row["segment"], + "type": "segment", + "metric": "cost", + "a_value": row["cost_a"], + "b_value": row["cost_b"], + "delta": row["delta_pct"], + } + ) + + for _, row in label_diff.iterrows(): + if pd.notna(row.get("delta_pct")) and row["delta_pct"] > 10: + watchlist_rows.append( + { + "label_or_segment": row["usage_label"], + "type": "label", + "metric": "cost", + "a_value": row["cost_a"], + "b_value": row["cost_b"], + "delta": row["delta_pct"], + } + ) + hit_delta = row.get("hit_rate_delta", 0) + if pd.notna(hit_delta) and hit_delta < -0.10: + watchlist_rows.append( + { + "label_or_segment": row["usage_label"], + "type": "label", + "metric": "hit_rate", + "a_value": row.get("hit_rate_a", 0), + "b_value": row.get("hit_rate_b", 0), + "delta": hit_delta * 100, + } + ) + + regression_watchlist = pd.DataFrame( + watchlist_rows, + columns=["label_or_segment", "type", "metric", "a_value", "b_value", "delta"], + ) + _write(regression_watchlist, out_dir, "regression_watchlist.csv") + + # --- stdout summary --- + print(f"\n=== COMPARE: {label_a} vs {label_b} ===") + print(f" Corpus partitioning (apples-to-apples = {n_both} contracts):") + print(f" successful_in_both: {len(successful_in_both)}") + print(f" regressed (a→err b): {len(regressed)}") + print(f" recovered (err a→b): {len(recovered)}") + print(f" new in {label_b}: {len(new_in_b)}") + print(f" dropped from {label_a}: {len(dropped_from_a)}") + + for row in summary_rows: + print( + f"\n [{row['run']}] contracts={row['n_contracts']} successful={row['n_successful']} " + f"batch_total=${row['batch_total']:.4f} cache_hit={row['cache_hit_rate']:.1%}" + ) + + if not regression_watchlist.empty: + print(f"\n Regression watchlist ({len(regression_watchlist)} items):") + _print_table( + "Regression Watchlist", + regression_watchlist, + [ + c + for c in ["label_or_segment", "metric", "a_value", "b_value", "delta"] + if c in regression_watchlist.columns + ], + ) + else: + print("\n No regressions detected (>10% cost growth or >10pp hit-rate drop).") + + print(f"\nOutputs written to: {out_dir}") + logger.info("run_compare complete: %s", out_dir) + + +# --------------------------------------------------------------------------- +# Print-only mode (no file writes — dumps all tables to stdout) +# --------------------------------------------------------------------------- + + +def run_print(csv_path: Union[str, Path], view: str = "batch", top_n: int = 15) -> None: + """Print all relevant tables for an instrumentation CSV to stdout. + + Args: + csv_path: Path to an instrumentation CSV. + view: "batch" for cross-contract aggregation (default) or "single" + for per-call detail on a single contract / small run. + top_n: How many rows to show in each top-N table (default 15). + """ + df = load_instrumentation_csv(csv_path) + llm = df[df["event"] == "llm_call"] + + if view == "batch": + classes = classify_contracts(df) + successful = classes["successful"] + errored = classes["errored"] + n_attempted = len(successful | errored) + + # Per-contract cost (real contracts only) + real_mask = llm["filename"].apply(_is_real_contract) + per_file_cost = llm[real_mask].groupby("filename")["cost_usd"].sum() + successful_costs = per_file_cost.reindex(list(successful)).dropna() + + print(f"\n=== BATCH INSTRUMENTATION SUMMARY ===") + print(f" Source: {csv_path}") + print(f" Contracts attempted: {n_attempted}") + print(f" Successful: {len(successful)}") + print(f" Errored: {len(errored)}") + print(f" Batch total cost: ${llm['cost_usd'].sum():.4f}") + if not successful_costs.empty: + print(f" Mean $/successful: ${successful_costs.mean():.4f}") + print(f" Median $/successful: ${successful_costs.median():.4f}") + print(f" P95 $/successful: ${successful_costs.quantile(0.95):.4f}") + print(f" Cache hit rate: {_cache_hit_rate(df):.1%}") + + # --- by_segment --- + seg_df = by_segment(df) + if not seg_df.empty: + seg_df = seg_df.sort_values("cost_usd", ascending=False) + _print_table( + "Cost by Segment", + seg_df, + ["segment", "n_calls", "cost_usd", "pct_of_total", "cache_hit_rate"], + ) + + # --- per_contract (successful + errored, restricted to real contracts) --- + per_contract = ( + per_file_cost.rename("total_cost") + .to_frame() + .reset_index() + .sort_values("total_cost", ascending=False) + ) + per_contract["status"] = per_contract["filename"].apply( + lambda fn: "errored" if fn in errored else "successful" + ) + _print_table( + "Per-Contract Cost", + per_contract, + ["filename", "total_cost", "status"], + ) + + # --- by_label top N by cost --- + label_df = by_label(df) + if not label_df.empty: + label_df = label_df.sort_values("cost_usd", ascending=False).head(top_n) + _print_table( + f"Top {top_n} Labels by Cost", + label_df, + ["usage_label", "segment", "n_calls", "cost_usd", "cache_hit_rate"], + ) + + # --- cache_miss_breakdown --- + miss_df = cache_breakdown(df) + if not miss_df.empty: + miss_df = miss_df.sort_values("cost_usd", ascending=False) + _print_table( + "Cache Miss Reason Breakdown", + miss_df, + ["cache_miss_reason", "usage_label", "n", "cost_usd"], + ) + + # --- cache lever (not_attempted) --- + if "cache_miss_reason" in llm.columns: + not_attempted = llm[llm["cache_miss_reason"] == "not_attempted"] + if not not_attempted.empty: + lever = ( + not_attempted.groupby(["usage_label", "segment"])["cost_usd"] + .sum() + .reset_index() + .sort_values("cost_usd", ascending=False) + ) + lever["potential_savings_if_cached"] = lever["cost_usd"] * 0.9 + _print_table( + "Cache-Lever Candidates (cache_miss_reason='not_attempted')", + lever.head(top_n), + [ + "usage_label", + "segment", + "cost_usd", + "potential_savings_if_cached", + ], + ) + + # --- retries + errors --- + retry_df = df[df["event"].isin(["llm_retry", "llm_error"])] + if not retry_df.empty: + re_agg = ( + retry_df.groupby(["event", "usage_label", "error_class"]) + .size() + .rename("n") + .reset_index() + .sort_values("n", ascending=False) + ) + _print_table( + "Retries + Errors", + re_agg, + ["event", "usage_label", "error_class", "n"], + ) + else: + print("\n--- Retries + Errors ---\n (none)") + + # --- errored contracts --- + if errored: + print(f"\n--- Errored Contracts ---") + for fn in sorted(errored): + wasted = float(per_file_cost.get(fn, 0.0)) + print(f" {fn} wasted=${wasted:.4f}") + else: + print(f"\n--- Errored Contracts ---\n (none)") + + elif view == "single": + n_inmem = (df["event"] == "llm_call_inmem_hit").sum() + n_retries = (df["event"] == "llm_retry").sum() + n_errors = (df["event"] == "llm_error").sum() + total_cost = float(llm["cost_usd"].sum()) + wall = float(df["ts_rel_sec"].max()) if "ts_rel_sec" in df.columns else None + + print(f"\n=== SINGLE-FILE INSTRUMENTATION SUMMARY ===") + print(f" Source: {csv_path}") + print(f" Total cost: ${total_cost:.4f}") + print(f" LLM calls: {len(llm)}") + print(f" In-mem hits: {int(n_inmem)}") + print(f" Cache hit rate: {_cache_hit_rate(df):.1%}") + print(f" Retries: {int(n_retries)}") + print(f" Errors: {int(n_errors)}") + if wall is not None: + print(f" Wall clock: {wall:.1f}s") + + seg_df = by_segment(df) + if not seg_df.empty: + seg_df = seg_df.sort_values("cost_usd", ascending=False) + _print_table( + "Cost by Segment", + seg_df, + ["segment", "n_calls", "cost_usd", "pct_of_total", "cache_hit_rate"], + ) + + label_df = by_label(df) + if not label_df.empty: + label_df = label_df.sort_values("cost_usd", ascending=False).head(top_n) + _print_table( + f"Top {top_n} Labels by Cost", + label_df, + ["usage_label", "segment", "n_calls", "cost_usd", "cache_hit_rate"], + ) + + miss_df = cache_breakdown(df) + if not miss_df.empty: + miss_df = miss_df.sort_values("cost_usd", ascending=False) + _print_table( + "Cache Miss Reason Breakdown", + miss_df.head(top_n), + ["cache_miss_reason", "usage_label", "n", "cost_usd"], + ) + + rc = df[df["event"] == "row_count"] + if not rc.empty and "stage" in rc.columns: + stage_order = rc["stage"].drop_duplicates().tolist() + funnel = ( + rc.groupby("stage", sort=False) + .agg(n_in=("n_in", "sum"), n_out=("n_out", "sum")) + .reindex(stage_order) + .reset_index() + ) + funnel["delta"] = funnel["n_out"] - funnel["n_in"] + _print_table( + "Row Count Funnel (event order)", + funnel, + ["stage", "n_in", "n_out", "delta"], + ) + + retry_df = df[df["event"].isin(["llm_retry", "llm_error"])] + if retry_df.empty: + print("\n--- Retries + Errors ---\n (none)") + else: + re_agg = ( + retry_df.groupby(["event", "usage_label", "error_class"]) + .size() + .rename("n") + .reset_index() + .sort_values("n", ascending=False) + ) + _print_table( + "Retries + Errors", + re_agg, + ["event", "usage_label", "error_class", "n"], + ) + + if not llm.empty: + top = llm.nlargest(min(10, len(llm)), "cost_usd")[ + [ + "seq", + "ts_rel_sec", + "segment", + "usage_label", + "cost_usd", + "input_tokens", + ] + ] + _print_table("Top 10 Calls by Cost", top, list(top.columns)) + + else: + raise ValueError(f"Unknown view: {view!r}. Use 'batch' or 'single'.") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Analyze instrumentation CSVs from doczy pipeline runs.", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + sub = parser.add_subparsers(dest="mode", required=True) + + # single + p_single = sub.add_parser("single", help="Summarize one contract/run trace.") + p_single.add_argument("--csv", required=True, help="Path to instrumentation CSV.") + p_single.add_argument("--out-dir", required=True, help="Output directory.") + + # batch + p_batch = sub.add_parser("batch", help="Summarize a full batch run.") + p_batch.add_argument("--csv", required=True, help="Path to instrumentation CSV.") + p_batch.add_argument("--out-dir", required=True, help="Output directory.") + + # compare + p_compare = sub.add_parser("compare", help="Compare two batch CSVs.") + p_compare.add_argument( + "--csv-a", required=True, help="Path to run-A instrumentation CSV." + ) + p_compare.add_argument( + "--csv-b", required=True, help="Path to run-B instrumentation CSV." + ) + p_compare.add_argument( + "--label-a", default="a", help="Label for run A (default: a)." + ) + p_compare.add_argument( + "--label-b", default="b", help="Label for run B (default: b)." + ) + p_compare.add_argument("--out-dir", required=True, help="Output directory.") + + # print (no file writes — dumps all tables to stdout) + p_print = sub.add_parser( + "print", + help="Print all analysis tables to stdout; write nothing to disk.", + ) + p_print.add_argument("--csv", required=True, help="Path to instrumentation CSV.") + p_print.add_argument( + "--view", + choices=["batch", "single"], + default="batch", + help="batch (default) aggregates across filenames; single is per-call detail.", + ) + p_print.add_argument( + "--top-n", + type=int, + default=15, + help="How many rows to show in each top-N table (default 15).", + ) + + return parser + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + parser = _build_parser() + args = parser.parse_args() + + if args.mode == "single": + run_single(args.csv, args.out_dir) + elif args.mode == "batch": + run_batch(args.csv, args.out_dir) + elif args.mode == "compare": + run_compare( + csv_a=args.csv_a, + csv_b=args.csv_b, + out_dir=args.out_dir, + label_a=args.label_a, + label_b=args.label_b, + ) + elif args.mode == "print": + run_print(args.csv, view=args.view, top_n=args.top_n) + else: + parser.print_help() + sys.exit(1) diff --git a/src/tests/test_instrumentation.py b/src/tests/test_instrumentation.py new file mode 100644 index 0000000..3a012d0 --- /dev/null +++ b/src/tests/test_instrumentation.py @@ -0,0 +1,697 @@ +"""Tests for src/utils/instrumentation.py and src/utils/instrumentation_context.py.""" + +import concurrent.futures +import csv +import os +import threading +import time + +import pandas as pd +import pytest + +import src.utils.instrumentation as instrumentation +import src.utils.instrumentation_context as instrumentation_context + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _reset_logger(): + """Force close and reset the module-level logger between tests.""" + instrumentation.close() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestDisabledIsNoop: + def test_disabled_is_noop(self, tmp_path, monkeypatch): + # Instrumentation is on by default; opt out explicitly. + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "0") + monkeypatch.delenv("DOCZY_INSTRUMENTATION_CSV", raising=False) + _reset_logger() + + # log should not raise and should not create any file + instrumentation.log("file_start", filename="x.txt") + + assert not instrumentation.is_enabled() + csv_path = tmp_path / "should_not_exist.csv" + assert not csv_path.exists() + + +class TestEnabledWritesHeaderAndRows: + def test_enabled_writes_header_and_rows(self, tmp_path, monkeypatch): + csv_path = tmp_path / "test_run.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + try: + instrumentation.log("file_start", filename="a.txt") + instrumentation.log( + "file_end", filename="a.txt", extra_json='{"final_rows": 5}' + ) + finally: + instrumentation.close() + + df = pd.read_csv(csv_path) + assert len(df) == 2 + assert list(df.columns) == instrumentation.FIELDNAMES + assert df["event"].tolist() == ["file_start", "file_end"] + assert df.loc[1, "extra_json"] == '{"final_rows": 5}' + + +class TestThreadSafetySequenceMonotonic: + def test_thread_safety_sequence_monotonic(self, tmp_path, monkeypatch): + csv_path = tmp_path / "thread_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + n_threads = 20 + calls_per_thread = 100 + barrier = threading.Barrier(n_threads) + + def worker(): + barrier.wait() + for _ in range(calls_per_thread): + instrumentation.log("row_count", stage="test", n_in=0, n_out=1) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + instrumentation.close() + + df = pd.read_csv(csv_path) + expected_total = n_threads * calls_per_thread + assert len(df) == expected_total + seq_values = df["seq"].tolist() + # All seq values must be unique + assert len(set(seq_values)) == expected_total + # seq values must span 1..N without gaps + assert sorted(seq_values) == list(range(1, expected_total + 1)) + + +class TestContextScopePropagation: + def test_context_scope_propagation(self, tmp_path, monkeypatch): + # Asserts that submit_with_context correctly propagates the caller's + # instr_scope context into the worker thread, so worker threads see + # the same filename (and other context vars) set by the caller. + csv_path = tmp_path / "ctx_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + results = [] + + def worker(): + ctx = instrumentation_context.get_ctx() + results.append(ctx.get("filename")) + instrumentation.log( + "file_start", **{k: v for k, v in ctx.items() if not k.startswith("_")} + ) + + with instrumentation_context.instr_scope(filename="f1.txt"): + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex: + future = instrumentation_context.submit_with_context(ex, worker) + future.result() + + instrumentation.close() + + df = pd.read_csv(csv_path) + assert len(df) == 1 + assert df.loc[0, "filename"] == "f1.txt" + assert results[0] == "f1.txt" + + +class TestSubmitWithContextHelperIsolation: + def test_submit_with_context_helper_isolation(self, tmp_path, monkeypatch): + # Asserts that two parallel scopes with different filename values each + # propagate only their own context to their respective workers — i.e. + # scopes do not bleed into each other when using submit_with_context. + csv_path = tmp_path / "isolation_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + results = {} + + def worker(expected_filename): + ctx = instrumentation_context.get_ctx() + results[expected_filename] = ctx.get("filename") + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex: + with instrumentation_context.instr_scope(filename="scope_a.txt"): + future_a = instrumentation_context.submit_with_context( + ex, worker, "scope_a.txt" + ) + with instrumentation_context.instr_scope(filename="scope_b.txt"): + future_b = instrumentation_context.submit_with_context( + ex, worker, "scope_b.txt" + ) + future_a.result() + future_b.result() + + instrumentation.close() + + assert results["scope_a.txt"] == "scope_a.txt" + assert results["scope_b.txt"] == "scope_b.txt" + + +class TestMapWithContextHelper: + def test_map_with_context_parallel_concurrent_workers(self, tmp_path, monkeypatch): + """Tests that map_with_context correctly propagates context into each + worker task with true concurrency. Uses 4 workers and 20 items with + sleep to force overlapping worker execution and trigger the bug if + fix is reverted (RuntimeError: cannot enter context: ... is already entered). + """ + csv_path = tmp_path / "map_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + results = [] + + def worker_fn(item): + # Sleep to force overlapping execution across workers + time.sleep(0.01) + ctx = instrumentation_context.get_ctx() + filename = ctx.get("filename") + results.append((item, filename)) + return item * 2 + + with instrumentation_context.instr_scope(filename="map_test.txt"): + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as ex: + mapped_results = list( + instrumentation_context.map_with_context( + ex, worker_fn, list(range(20)) + ) + ) + + instrumentation.close() + + # Verify results are correct + assert mapped_results == [i * 2 for i in range(20)] + # Verify all workers saw the correct filename (no context bleeding) + assert len(results) == 20 + assert all(filename == "map_test.txt" for _, filename in results) + + def test_map_with_context_scope_isolation(self, tmp_path, monkeypatch): + """Tests that two sequential map_with_context calls in different + instr_scope contexts each see only their own scope's context (no bleed-over). + """ + csv_path = tmp_path / "map_isolation_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + results_a = [] + results_b = [] + + def worker_fn_a(item): + ctx = instrumentation_context.get_ctx() + filename = ctx.get("filename") + results_a.append((item, filename)) + return item + + def worker_fn_b(item): + ctx = instrumentation_context.get_ctx() + filename = ctx.get("filename") + results_b.append((item, filename)) + return item + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as ex: + # First scope + with instrumentation_context.instr_scope(filename="scope_a.txt"): + list( + instrumentation_context.map_with_context(ex, worker_fn_a, [1, 2, 3]) + ) + + # Second scope (sequential, not concurrent) + with instrumentation_context.instr_scope(filename="scope_b.txt"): + list( + instrumentation_context.map_with_context(ex, worker_fn_b, [4, 5, 6]) + ) + + instrumentation.close() + + # Verify first scope's workers saw correct filename + assert len(results_a) == 3 + assert all(filename == "scope_a.txt" for _, filename in results_a) + + # Verify second scope's workers saw correct filename (no bleed-over from first) + assert len(results_b) == 3 + assert all(filename == "scope_b.txt" for _, filename in results_b) + + +class TestUsageLabelThreadsThroughInvokeClaude: + def test_usage_label_threads_through_invoke_claude(self, tmp_path, monkeypatch): + import src.utils.llm_utils as llm_utils + from src import config + from src.utils import usage_tracking + + csv_path = tmp_path / "llm_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + fake_response_body = { + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + }, + "content": [{"type": "text", "text": "hello"}], + } + + def fake_local_claude(prompt, model_id, filename, max_tokens, **kwargs): + # Simulate what local_claude_3_and_up does: call _track_usage_from_response + llm_utils._track_usage_from_response(fake_response_body, filename, model_id) + return "hello" + + model_id = config.MODEL_ID_CLAUDE35_SONNET + llm_utils.claude_cache.clear() + + original_run_mode = config.RUN_MODE + try: + config.RUN_MODE = "local" + monkeypatch.setattr( + "src.utils.llm_utils.local_claude_3_and_up", fake_local_claude + ) + + llm_utils.invoke_claude( + "test prompt", + model_id, + "test_file.txt", + usage_label="FOO", + ) + finally: + config.RUN_MODE = original_run_mode + llm_utils.claude_cache.clear() + + instrumentation.close() + + df = pd.read_csv(csv_path) + llm_rows = df[df["event"] == "llm_call"] + assert len(llm_rows) == 1 + assert llm_rows.iloc[0]["usage_label"] == "FOO" + + +# --------------------------------------------------------------------------- +# Segment column — mapping-first resolution with scope fallback +# --------------------------------------------------------------------------- + + +class TestSegmentResolution: + def test_segment_from_mapping_wins_over_scope(self): + # Known label should resolve via mapping even when scope says otherwise. + result = instrumentation.resolve_segment( + "DYNAMIC_ASSIGNMENT", ctx_segment="preprocessing" + ) + assert result == "one_to_n_enrichment" + + def test_segment_falls_back_to_scope_when_label_unknown(self): + result = instrumentation.resolve_segment( + "completely_new_prompt_label", ctx_segment="preprocessing" + ) + assert result == "preprocessing" + + def test_segment_unknown_when_neither_matches(self): + assert instrumentation.resolve_segment("bogus_label", None) == "unknown" + assert instrumentation.resolve_segment(None, None) == "unknown" + + def test_cache_warming_prefix_maps_to_cache_warming(self): + assert ( + instrumentation.resolve_segment("cache_warming_exhibit_header", None) + == "cache_warming" + ) + + +# --------------------------------------------------------------------------- +# cache_miss_reason classifier +# --------------------------------------------------------------------------- + + +class TestCacheMissReason: + def _reset_seen(self): + import src.utils.llm_utils as llm_utils + + with llm_utils._cache_key_seen_lock: + llm_utils._cache_key_seen.clear() + + def test_not_attempted_when_cache_flag_false(self): + import src.utils.llm_utils as llm_utils + + self._reset_seen() + assert ( + llm_utils._classify_cache_outcome( + has_context_cache=False, + cache_read_tokens=0, + cache_creation_tokens=0, + context_chars=0, + cache_key=None, + ) + == "not_attempted" + ) + + def test_hit_when_cache_read_tokens_positive(self): + import src.utils.llm_utils as llm_utils + + self._reset_seen() + assert ( + llm_utils._classify_cache_outcome( + has_context_cache=True, + cache_read_tokens=500, + cache_creation_tokens=0, + context_chars=10000, + cache_key="k1", + ) + == "hit" + ) + + def test_first_call_then_ttl_expired(self): + import src.utils.llm_utils as llm_utils + + self._reset_seen() + # First call — key never seen. + first = llm_utils._classify_cache_outcome( + has_context_cache=True, + cache_read_tokens=0, + cache_creation_tokens=5000, + context_chars=10000, + cache_key="k_shared", + ) + assert first == "first_call" + # Same key re-created later → TTL expired. + second = llm_utils._classify_cache_outcome( + has_context_cache=True, + cache_read_tokens=0, + cache_creation_tokens=5000, + context_chars=10000, + cache_key="k_shared", + ) + assert second == "ttl_expired" + + def test_under_min_tokens(self): + import src.utils.llm_utils as llm_utils + + self._reset_seen() + # ~500 tokens (1750 chars / 3.5) — below the 1024 minimum. + assert ( + llm_utils._classify_cache_outcome( + has_context_cache=True, + cache_read_tokens=0, + cache_creation_tokens=0, + context_chars=1750, + cache_key="k_small", + ) + == "under_min_tokens" + ) + + def test_silent_miss_when_above_threshold_but_no_tokens(self): + import src.utils.llm_utils as llm_utils + + self._reset_seen() + # ~5000 tokens but neither read nor creation — unexplained. + assert ( + llm_utils._classify_cache_outcome( + has_context_cache=True, + cache_read_tokens=0, + cache_creation_tokens=0, + context_chars=17500, + cache_key="k_silent", + ) + == "silent_miss" + ) + + +# --------------------------------------------------------------------------- +# llm_retry / llm_error emit helpers +# --------------------------------------------------------------------------- + + +class TestLlmRetryAndErrorEvents: + def test_llm_retry_event_written(self, tmp_path, monkeypatch): + import src.utils.llm_utils as llm_utils + + csv_path = tmp_path / "retry_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + with instrumentation_context.instr_scope( + _usage_label="REIMBURSEMENT_PRIMARY", + filename="retry_file.txt", + ): + llm_utils._log_llm_retry( + "retry_file.txt", + "sonnet-test", + attempt=2, + error_class="ThrottlingException", + error_msg="rate exceeded", + backoff_sec=1.234, + ) + + instrumentation.close() + df = pd.read_csv(csv_path) + retry_rows = df[df["event"] == "llm_retry"] + assert len(retry_rows) == 1 + row = retry_rows.iloc[0] + assert row["attempt"] == 2 + assert row["error_class"] == "ThrottlingException" + assert row["error_msg"] == "rate exceeded" + assert float(row["backoff_sec"]) == 1.234 + assert row["segment"] == "one_to_n_primary" + assert row["usage_label"] == "REIMBURSEMENT_PRIMARY" + + def test_llm_error_event_truncates_long_message(self, tmp_path, monkeypatch): + import src.utils.llm_utils as llm_utils + + csv_path = tmp_path / "error_test.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + + long_msg = "x" * 500 + with instrumentation_context.instr_scope( + _usage_label="DYNAMIC_CODE_ASSIGNMENT", + filename="err.txt", + ): + llm_utils._log_llm_error( + "err.txt", + "sonnet-test", + attempt=5, + error_class="ValidationException", + error_msg=long_msg, + ) + + instrumentation.close() + df = pd.read_csv(csv_path) + err_rows = df[df["event"] == "llm_error"] + assert len(err_rows) == 1 + row = err_rows.iloc[0] + assert row["attempt"] == 5 + assert row["error_class"] == "ValidationException" + assert len(row["error_msg"]) == 200 # truncated + assert row["segment"] == "one_to_n_enrichment" + + def test_emit_helpers_are_noop_when_disabled(self, tmp_path, monkeypatch): + import src.utils.llm_utils as llm_utils + + # Instrumentation is on by default; opt out explicitly. + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "0") + _reset_logger() + # Must not raise; must not create any file. + llm_utils._log_llm_retry("f.txt", "m", 1, "E", "msg", 0.1) + llm_utils._log_llm_error("f.txt", "m", 1, "E", "msg") + assert not instrumentation.is_enabled() + + +# --------------------------------------------------------------------------- +# Regression: emit helpers must not collide with scope-set public keys +# (caught during instrumented smoke run — `segment` key collision between +# instr_scope(segment=...) context and explicit segment=... kwarg) +# --------------------------------------------------------------------------- + + +class TestSegmentKwargNoCollision: + def _mk_logger(self, tmp_path, monkeypatch): + csv_path = tmp_path / "collision.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(csv_path)) + _reset_logger() + instrumentation.configure_from_env() + return csv_path + + def test_llm_retry_under_segment_scope_does_not_collide( + self, tmp_path, monkeypatch + ): + import src.utils.llm_utils as llm_utils + + csv_path = self._mk_logger(tmp_path, monkeypatch) + # Simulate file_processing.py wrapping a phase: public segment key in ctx. + with instrumentation_context.instr_scope( + segment="preprocessing", + filename="x.txt", + _usage_label="REIMBURSEMENT_PRIMARY", + ): + # Would raise TypeError: got multiple values for 'segment' if ctx + # splat isn't filtered. + llm_utils._log_llm_retry( + "x.txt", "sonnet", 1, "ThrottlingException", "", 0.5 + ) + + instrumentation.close() + df = pd.read_csv(csv_path) + assert len(df) == 1 + # Mapping wins over scope for known labels + assert df.loc[0, "segment"] == "one_to_n_primary" + + def test_inmem_hit_under_segment_scope_does_not_collide( + self, tmp_path, monkeypatch + ): + """Regression for the bug caught in the 2026-04-24 smoke run: + llm_call_inmem_hit was splatting ctx without stripping 'segment', + colliding with the explicit segment= kwarg.""" + import src.utils.llm_utils as llm_utils + from src import config + + csv_path = self._mk_logger(tmp_path, monkeypatch) + # Ensure inmem hit path fires: pre-seed the cache for a known key. + llm_utils.claude_cache.clear() + model_id = config.MODEL_ID_CLAUDE35_SONNET + cache_key = llm_utils.get_cache_key( + "prompt-x", model_id, instruction=None, max_tokens=4096, cache=False + ) + llm_utils.claude_cache[cache_key] = "cached response" + + try: + with instrumentation_context.instr_scope( + segment="preprocessing", + filename="x.txt", + ): + # Would raise TypeError before fix. + result = llm_utils.invoke_claude( + "prompt-x", model_id, "x.txt", usage_label="some_label" + ) + assert result == "cached response" + finally: + llm_utils.claude_cache.clear() + + instrumentation.close() + df = pd.read_csv(csv_path) + assert (df["event"] == "llm_call_inmem_hit").sum() == 1 + + +# --------------------------------------------------------------------------- +# configure_for_run — batch-aware runtime hookup +# --------------------------------------------------------------------------- + + +class TestConfigureForRun: + def test_noop_when_disabled(self, tmp_path, monkeypatch): + """When DOCZY_INSTRUMENTATION is explicitly opted out, configure_for_run is a no-op.""" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "0") + monkeypatch.delenv("DOCZY_INSTRUMENTATION_CSV", raising=False) + _reset_logger() + + instrumentation.configure_for_run("BATCH99", "run_ts_disabled") + + assert not instrumentation.is_enabled() + # No CSV should appear anywhere under tmp_path + assert list(tmp_path.iterdir()) == [] + + def test_writes_to_canonical_path_when_enabled(self, tmp_path, monkeypatch): + """With instrumentation enabled, configure_for_run targets the canonical path + and a subsequent log+close produces a CSV there.""" + import src.config as config + + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.delenv("DOCZY_INSTRUMENTATION_CSV", raising=False) + monkeypatch.setattr(config, "CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)) + _reset_logger() + + instrumentation.configure_for_run("BATCH42", "run_ts_X") + + try: + instrumentation.log("file_start", filename="contract.txt") + finally: + instrumentation.close() + + expected = tmp_path / "run_ts_X" / "tracking" / "BATCH42-INSTRUMENTATION.csv" + assert expected.exists(), f"Expected CSV at {expected}" + + def test_respects_explicit_csv_env_override(self, tmp_path, monkeypatch): + """If DOCZY_INSTRUMENTATION_CSV is already set, configure_for_run must + honour it rather than computing the canonical path.""" + import src.config as config + + explicit_csv = tmp_path / "explicit_override.csv" + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.setenv("DOCZY_INSTRUMENTATION_CSV", str(explicit_csv)) + # Point CONSOLIDATED_OUTPUT_DIRECTORY somewhere else so we can confirm + # the canonical path is NOT created. + canonical_dir = tmp_path / "canonical" + canonical_dir.mkdir() + monkeypatch.setattr(config, "CONSOLIDATED_OUTPUT_DIRECTORY", str(canonical_dir)) + _reset_logger() + + instrumentation.configure_for_run("BATCH42", "run_ts_X") + + try: + instrumentation.log("file_start", filename="contract.txt") + finally: + instrumentation.close() + + # Explicit path must exist + assert explicit_csv.exists() + # Canonical path must NOT exist + canonical_csv = ( + canonical_dir / "run_ts_X" / "tracking" / "BATCH42-INSTRUMENTATION.csv" + ) + assert not canonical_csv.exists() + + def test_schema_parity(self, tmp_path, monkeypatch): + """CSV produced after configure_for_run must have exactly instrumentation.FIELDNAMES.""" + import src.config as config + + monkeypatch.setenv("DOCZY_INSTRUMENTATION", "1") + monkeypatch.delenv("DOCZY_INSTRUMENTATION_CSV", raising=False) + monkeypatch.setattr(config, "CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)) + _reset_logger() + + instrumentation.configure_for_run("BATCHSCHEMA", "run_schema_ts") + + try: + instrumentation.log( + "llm_call", + filename="x.txt", + usage_label="DYNAMIC_ASSIGNMENT", + cost_usd=0.001, + input_tokens=500, + ) + finally: + instrumentation.close() + + csv_path = ( + tmp_path / "run_schema_ts" / "tracking" / "BATCHSCHEMA-INSTRUMENTATION.csv" + ) + df = pd.read_csv(csv_path) + assert list(df.columns) == instrumentation.FIELDNAMES diff --git a/src/tests/test_instrumentation_metrics.py b/src/tests/test_instrumentation_metrics.py new file mode 100644 index 0000000..f03c0f8 --- /dev/null +++ b/src/tests/test_instrumentation_metrics.py @@ -0,0 +1,744 @@ +"""Tests for src/testbed/instrumentation_metrics.py — manual analyzer. + +Tests are written against the spec in: + .claude/plans/instrumentation_runtime_and_analyzer_20260424.md + +The implementation is exercised via: + from src.testbed import instrumentation_metrics as im + +Public API assumed (per plan §2): + im.run_single(csv_path, out_dir) + im.run_batch(csv_path, out_dir) + im.run_compare(csv_a, csv_b, out_dir, label_a="a", label_b="b") + +Fixture CSVs are built with csv.DictWriter against instrumentation.FIELDNAMES so +that schema parity with the real logger is guaranteed. +""" + +import csv +import os +from pathlib import Path + +import pandas as pd +import pytest + +import src.utils.instrumentation as instrumentation +from src.testbed import instrumentation_metrics as im + + +# --------------------------------------------------------------------------- +# Fixture helpers +# --------------------------------------------------------------------------- + +FIELDNAMES = instrumentation.FIELDNAMES + +_CACHE_MISS_REASONS = [ + "hit", + "first_call", + "ttl_expired", + "under_min_tokens", + "silent_miss", + "not_attempted", +] + + +def _blank_row(**overrides): + """Return a dict with every FIELDNAMES key set to empty string, then apply overrides.""" + row = {f: "" for f in FIELDNAMES} + row.update(overrides) + return row + + +def _write_fixture(path: Path, rows: list[dict]) -> Path: + """Write rows to a CSV at path using FIELDNAMES header.""" + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=FIELDNAMES, extrasaction="ignore") + writer.writeheader() + for row in rows: + writer.writerow(_blank_row(**row)) + return path + + +def _llm_call( + filename="contract.txt", + segment="one_to_n_enrichment", + usage_label="DYNAMIC_ASSIGNMENT", + cost_usd=0.01, + input_tokens=500, + cache_miss_reason="hit", + seq=1, +): + return _blank_row( + event="llm_call", + filename=filename, + segment=segment, + usage_label=usage_label, + cost_usd=str(cost_usd), + input_tokens=str(input_tokens), + cache_miss_reason=cache_miss_reason, + seq=str(seq), + ) + + +def _file_start(filename="contract.txt", seq=100): + return _blank_row(event="file_start", filename=filename, seq=str(seq)) + + +def _file_end(filename="contract.txt", seq=200): + return _blank_row(event="file_end", filename=filename, seq=str(seq)) + + +def _llm_error( + filename="contract.txt", + usage_label="DYNAMIC_ASSIGNMENT", + segment="one_to_n_enrichment", + error_class="ThrottlingException", + error_msg="rate exceeded", + attempt=3, + seq=300, +): + return _blank_row( + event="llm_error", + filename=filename, + usage_label=usage_label, + segment=segment, + error_class=error_class, + error_msg=error_msg, + attempt=str(attempt), + seq=str(seq), + ) + + +def _row_count( + filename="contract.txt", stage="dedup_across_chunks", n_in=10, n_out=8, seq=50 +): + return _blank_row( + event="row_count", + filename=filename, + stage=stage, + n_in=str(n_in), + n_out=str(n_out), + seq=str(seq), + ) + + +# --------------------------------------------------------------------------- +# Single-mode fixtures and tests +# --------------------------------------------------------------------------- + + +def _make_single_fixture(path: Path) -> Path: + """20-row fixture for a single contract.""" + rows = [] + rows.append(_file_start("akron.txt", seq=1)) + + # 12 llm_call events across 3 segments + for i in range(4): + rows.append( + _llm_call( + filename="akron.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.02, + input_tokens=800, + cache_miss_reason="hit", + seq=10 + i, + ) + ) + for i in range(5): + rows.append( + _llm_call( + filename="akron.txt", + segment="one_to_n_enrichment", + usage_label="DYNAMIC_ASSIGNMENT", + cost_usd=0.01, + input_tokens=400, + cache_miss_reason="first_call", + seq=20 + i, + ) + ) + for i in range(3): + rows.append( + _llm_call( + filename="akron.txt", + segment="one_to_one", + usage_label="prompt_contract_dates", + cost_usd=0.005, + input_tokens=300, + cache_miss_reason="not_attempted", + seq=30 + i, + ) + ) + + # cache_miss_reason variety rows (one per reason) + for idx, reason in enumerate(_CACHE_MISS_REASONS): + rows.append( + _llm_call( + filename="akron.txt", + segment="one_to_n_enrichment", + usage_label=f"label_reason_{idx}", + cost_usd=0.001, + cache_miss_reason=reason, + seq=50 + idx, + ) + ) + + # row_count events — deliberate non-alphabetical stage order + stage_order = [ + "initial_extraction", + "dedup_across_chunks", + "methodology_breakout_single_row", + "split_service_terms", + ] + for idx, stage in enumerate(stage_order): + rows.append( + _row_count( + "akron.txt", + stage=stage, + n_in=10 + idx, + n_out=10 + idx + 2, + seq=60 + idx, + ) + ) + + rows.append(_file_end("akron.txt", seq=200)) + return _write_fixture(path, rows) + + +class TestSingleMode: + def test_summary_csv_has_expected_columns_and_types(self, tmp_path): + fixture = tmp_path / "single_fixture.csv" + _make_single_fixture(fixture) + out_dir = tmp_path / "single_out" + + im.run_single(fixture, out_dir) + + summary_path = out_dir / "summary.csv" + assert summary_path.exists(), "summary.csv not produced" + df = pd.read_csv(summary_path) + assert len(df) == 1, "summary.csv must have exactly 1 row" + required_cols = { + "total_cost", + "n_llm_calls", + "cache_hit_rate", + } + missing = required_cols - set(df.columns) + assert not missing, f"summary.csv missing columns: {missing}" + # total_cost must be numeric and non-negative + assert pd.api.types.is_numeric_dtype( + df["total_cost"] + ), "total_cost must be numeric" + assert float(df["total_cost"].iloc[0]) >= 0 + + def test_by_segment_aggregates_cost_correctly(self, tmp_path): + fixture = tmp_path / "single_fixture.csv" + _make_single_fixture(fixture) + out_dir = tmp_path / "single_out" + + im.run_single(fixture, out_dir) + + by_seg_path = out_dir / "by_segment.csv" + assert by_seg_path.exists(), "by_segment.csv not produced" + df = pd.read_csv(by_seg_path) + + # The fixture has llm_call events with known segment/cost + # one_to_n_primary: 4 × 0.02 = 0.08 + # one_to_n_enrichment: 5 × 0.01 + 6 × 0.001 = 0.056 + # one_to_one: 3 × 0.005 = 0.015 + # Totals per segment must match (allowing floating point tolerance) + seg_costs = df.set_index("segment")["cost_usd"] + assert "one_to_n_primary" in seg_costs.index + assert abs(float(seg_costs["one_to_n_primary"]) - 0.08) < 1e-6 + + # Sum across all segments must equal total llm_call cost (excluding reason rows for simplicity) + total_from_segments = float(df["cost_usd"].sum()) + # Read fixture directly to compute expected total + raw = pd.read_csv(fixture) + llm_rows = raw[raw["event"] == "llm_call"] + expected_total = float(llm_rows["cost_usd"].astype(float).sum()) + assert abs(total_from_segments - expected_total) < 1e-6 + + def test_cache_miss_breakdown_groups_by_reason(self, tmp_path): + fixture = tmp_path / "single_fixture.csv" + _make_single_fixture(fixture) + out_dir = tmp_path / "single_out" + + im.run_single(fixture, out_dir) + + breakdown_path = out_dir / "cache_miss_breakdown.csv" + assert breakdown_path.exists(), "cache_miss_breakdown.csv not produced" + df = pd.read_csv(breakdown_path) + + # Fixture includes one row per cache_miss_reason value + found_reasons = set(df["cache_miss_reason"].unique()) + for reason in _CACHE_MISS_REASONS: + assert ( + reason in found_reasons + ), f"cache_miss_reason '{reason}' missing from breakdown" + + def test_retries_errors_csv_written_even_if_empty(self, tmp_path): + """A fixture with zero retries/errors must still produce retries_errors.csv with a header.""" + # Minimal fixture: only llm_call events, no llm_retry/llm_error + rows = [ + _file_start("clean.txt", seq=1), + _llm_call("clean.txt", seq=2), + _file_end("clean.txt", seq=3), + ] + fixture = tmp_path / "clean_fixture.csv" + _write_fixture(fixture, rows) + out_dir = tmp_path / "retries_out" + + im.run_single(fixture, out_dir) + + retries_path = out_dir / "retries_errors.csv" + assert ( + retries_path.exists() + ), "retries_errors.csv not produced even for zero-retry fixture" + # Must at least have a header row (file is not empty) + assert retries_path.stat().st_size > 0, "retries_errors.csv is completely empty" + + def test_row_count_funnel_preserves_order(self, tmp_path): + """row_count_funnel.csv must reflect stage insertion order, not alphabetical order.""" + fixture = tmp_path / "single_fixture.csv" + _make_single_fixture(fixture) + out_dir = tmp_path / "funnel_out" + + im.run_single(fixture, out_dir) + + funnel_path = out_dir / "row_count_funnel.csv" + assert funnel_path.exists(), "row_count_funnel.csv not produced" + df = pd.read_csv(funnel_path) + + # Fixture stages (in insertion order): initial_extraction, dedup_across_chunks, + # methodology_breakout_single_row, split_service_terms + # Alphabetical order would be: dedup_across_chunks, initial_extraction, + # methodology_breakout_single_row, split_service_terms + # So positions of initial_extraction and dedup_across_chunks must not be swapped + stages_in_output = df["stage"].tolist() + if ( + "initial_extraction" in stages_in_output + and "dedup_across_chunks" in stages_in_output + ): + idx_initial = stages_in_output.index("initial_extraction") + idx_dedup = stages_in_output.index("dedup_across_chunks") + assert idx_initial < idx_dedup, ( + "row_count_funnel.csv appears alphabetically sorted; " + "expected insertion/event order" + ) + + +# --------------------------------------------------------------------------- +# Batch-mode fixtures and tests +# --------------------------------------------------------------------------- + + +def _make_batch_fixture(path: Path) -> Path: + """Fixture with 3 contracts: 2 successful, 1 errored via llm_error.""" + rows = [] + + # Contract A — successful, high cost + rows.append(_file_start("contract_a.txt", seq=1)) + for i in range(5): + rows.append( + _llm_call( + filename="contract_a.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.05, + seq=10 + i, + ) + ) + rows.append(_file_end("contract_a.txt", seq=20)) + + # Contract B — successful, lower cost + rows.append(_file_start("contract_b.txt", seq=30)) + for i in range(3): + rows.append( + _llm_call( + filename="contract_b.txt", + segment="one_to_n_enrichment", + usage_label="DYNAMIC_ASSIGNMENT", + cost_usd=0.01, + seq=40 + i, + ) + ) + rows.append(_file_end("contract_b.txt", seq=50)) + + # Contract C — errored (has llm_error event, no file_end) + rows.append(_file_start("contract_c.txt", seq=60)) + for i in range(2): + rows.append( + _llm_call( + filename="contract_c.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.03, + seq=70 + i, + ) + ) + rows.append( + _llm_error( + filename="contract_c.txt", + usage_label="DYNAMIC_ASSIGNMENT", + segment="one_to_n_enrichment", + seq=80, + ) + ) + # No file_end for contract_c — simulates a crash + + # not_attempted rows to test cache lever estimate + for i in range(3): + rows.append( + _llm_call( + filename="contract_a.txt", + segment="one_to_one", + usage_label="prompt_tin_npi", + cost_usd=0.002, + cache_miss_reason="not_attempted", + seq=90 + i, + ) + ) + + return _write_fixture(path, rows) + + +class TestBatchMode: + def test_batch_summary_counts_successful_vs_errored(self, tmp_path): + fixture = tmp_path / "batch_fixture.csv" + _make_batch_fixture(fixture) + out_dir = tmp_path / "batch_out" + + im.run_batch(fixture, out_dir) + + summary_path = out_dir / "batch_summary.csv" + assert summary_path.exists(), "batch_summary.csv not produced" + df = pd.read_csv(summary_path) + assert len(df) == 1 + + # 2 successful (a, b), 1 errored (c) + assert int(df["n_successful"].iloc[0]) == 2 + assert int(df["n_errored"].iloc[0]) == 1 + assert int(df["n_contracts_attempted"].iloc[0]) == 3 + + def test_per_contract_splits_by_filename(self, tmp_path): + fixture = tmp_path / "batch_fixture.csv" + _make_batch_fixture(fixture) + out_dir = tmp_path / "batch_out" + + im.run_batch(fixture, out_dir) + + per_contract_path = out_dir / "per_contract.csv" + assert per_contract_path.exists(), "per_contract.csv not produced" + df = pd.read_csv(per_contract_path) + + # Exactly 3 contracts + assert len(df) == 3 + filenames = set(df["filename"].tolist()) + assert filenames == {"contract_a.txt", "contract_b.txt", "contract_c.txt"} + + # contract_a total cost: 5×0.05 + 3×0.002 = 0.256 + row_a = df[df["filename"] == "contract_a.txt"].iloc[0] + expected_cost_a = 5 * 0.05 + 3 * 0.002 + assert abs(float(row_a["total_cost"]) - expected_cost_a) < 1e-6 + + # contract_b total cost: 3×0.01 = 0.03 + row_b = df[df["filename"] == "contract_b.txt"].iloc[0] + assert abs(float(row_b["total_cost"]) - 0.03) < 1e-6 + + def test_top10_offenders_ranked_descending_by_cost(self, tmp_path): + fixture = tmp_path / "batch_fixture.csv" + _make_batch_fixture(fixture) + out_dir = tmp_path / "batch_out" + + im.run_batch(fixture, out_dir) + + top10_path = out_dir / "top10_offenders.csv" + assert top10_path.exists(), "top10_offenders.csv not produced" + df = pd.read_csv(top10_path) + assert len(df) >= 1 + + # Costs must be in descending order + costs = df["total_cost"].astype(float).tolist() + assert costs == sorted( + costs, reverse=True + ), "top10_offenders.csv not sorted descending" + + def test_cache_lever_estimate_uses_not_attempted_rows(self, tmp_path): + fixture = tmp_path / "batch_fixture.csv" + _make_batch_fixture(fixture) + out_dir = tmp_path / "batch_out" + + im.run_batch(fixture, out_dir) + + cache_lever_path = out_dir / "cache_lever_estimate.csv" + assert cache_lever_path.exists(), "cache_lever_estimate.csv not produced" + df = pd.read_csv(cache_lever_path) + + # Only labels with not_attempted rows should appear + # Fixture has 3 not_attempted rows for "prompt_tin_npi" + assert "prompt_tin_npi" in df["label"].tolist() + + # Labels that never had not_attempted should NOT appear + # (prompt_reimbursement_primary used "hit" in the fixture) + assert "prompt_reimbursement_primary" not in df["label"].tolist() + + +# --------------------------------------------------------------------------- +# Compare-mode fixtures and tests +# --------------------------------------------------------------------------- + + +def _make_compare_fixture_a(path: Path) -> Path: + """Run A: contracts W (successful), X (successful), Y (successful).""" + rows = [] + + # W — successful + rows.append(_file_start("w.txt", seq=1)) + rows.append(_llm_call("w.txt", segment="one_to_n_primary", cost_usd=0.10, seq=2)) + rows.append(_file_end("w.txt", seq=3)) + + # X — successful in A, cost $0.50 (deliberately large to test exclusion in compare) + rows.append(_file_start("x.txt", seq=10)) + rows.append( + _llm_call( + "x.txt", + segment="one_to_n_enrichment", + usage_label="DYNAMIC_ASSIGNMENT", + cost_usd=0.50, + seq=11, + ) + ) + rows.append(_file_end("x.txt", seq=12)) + + # Y — only in A + rows.append(_file_start("y.txt", seq=20)) + rows.append(_llm_call("y.txt", segment="one_to_one", cost_usd=0.05, seq=21)) + rows.append(_file_end("y.txt", seq=22)) + + return _write_fixture(path, rows) + + +def _make_compare_fixture_b(path: Path) -> Path: + """Run B: contracts W (successful), X (errored), Z (successful, new). + + X has llm_error in B — deliberately small cost so the diff math is clearly + affected only by W and Z in the apples-to-apples slice. + """ + rows = [] + + # W — successful, slightly lower cost (10% drop on one segment) + rows.append(_file_start("w.txt", seq=1)) + rows.append( + _llm_call( + "w.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.09, # 10% cheaper than A + seq=2, + ) + ) + rows.append( + _llm_call( + "w.txt", + segment="preprocessing", + usage_label="prompt_exhibit_header", + cost_usd=0.08, # 15% more expensive than A (for regression watchlist test) + seq=3, + ) + ) + rows.append(_file_end("w.txt", seq=4)) + + # X — errored in B (has llm_error, no file_end) + rows.append(_file_start("x.txt", seq=10)) + rows.append( + _llm_call("x.txt", segment="one_to_n_enrichment", cost_usd=0.001, seq=11) + ) + rows.append(_llm_error("x.txt", seq=12)) + + # Z — only in B + rows.append(_file_start("z.txt", seq=20)) + rows.append(_llm_call("z.txt", segment="one_to_one", cost_usd=0.04, seq=21)) + rows.append(_file_end("z.txt", seq=22)) + + return _write_fixture(path, rows) + + +class TestCompareMode: + def test_compare_partitions_into_four_disjoint_sets(self, tmp_path): + """corpus_diff.csv must categorise W/X/Y/Z into the correct partitions, + and no filename must appear in two categories at once.""" + fixture_a = tmp_path / "a.csv" + fixture_b = tmp_path / "b.csv" + _make_compare_fixture_a(fixture_a) + _make_compare_fixture_b(fixture_b) + out_dir = tmp_path / "compare_out" + + im.run_compare(fixture_a, fixture_b, out_dir, label_a="run_a", label_b="run_b") + + corpus_path = out_dir / "corpus_diff.csv" + assert corpus_path.exists(), "corpus_diff.csv not produced" + df = pd.read_csv(corpus_path) + + w_row = df[df["filename"] == "w.txt"] + x_row = df[df["filename"] == "x.txt"] + y_row = df[df["filename"] == "y.txt"] + z_row = df[df["filename"] == "z.txt"] + + assert not w_row.empty, "w.txt missing from corpus_diff" + assert not x_row.empty, "x.txt missing from corpus_diff" + assert not y_row.empty, "y.txt missing from corpus_diff" + assert not z_row.empty, "z.txt missing from corpus_diff" + + assert w_row.iloc[0]["category"] == "in_both", "w.txt should be in_both" + assert ( + x_row.iloc[0]["category"] == "regressed" + ), "x.txt should be regressed (success→error)" + assert y_row.iloc[0]["category"] == "only_a", "y.txt should be only_a" + assert z_row.iloc[0]["category"] == "only_b", "z.txt should be only_b" + + # Disjointness: each filename appears exactly once + assert ( + df["filename"].duplicated().sum() == 0 + ), "corpus_diff.csv has duplicate filenames" + + def test_compare_cost_diff_uses_only_successful_in_both(self, tmp_path): + """CRITICAL INVARIANT: by_segment_diff.csv cost computations must exclude + contracts not in the successful_in_both set. + + X has $0.50 cost in A (one_to_n_enrichment) but is errored in B. + If X is included, the one_to_n_enrichment segment in A will appear + inflated by $0.50. The test confirms it is excluded entirely. + """ + fixture_a = tmp_path / "a.csv" + fixture_b = tmp_path / "b.csv" + _make_compare_fixture_a(fixture_a) + _make_compare_fixture_b(fixture_b) + out_dir = tmp_path / "compare_out" + + im.run_compare(fixture_a, fixture_b, out_dir, label_a="run_a", label_b="run_b") + + diff_path = out_dir / "by_segment_diff.csv" + assert diff_path.exists(), "by_segment_diff.csv not produced" + df = pd.read_csv(diff_path) + + # Only W is in successful_in_both. + # W's one_to_n_primary in A = $0.10, in B = $0.09. + # X's one_to_n_enrichment ($0.50 in A) must be EXCLUDED. + # Therefore cost_a for one_to_n_enrichment on the apples-to-apples slice + # must NOT include $0.50 from X. + if "one_to_n_enrichment" in df["segment"].values: + seg_row = df[df["segment"] == "one_to_n_enrichment"].iloc[0] + # W has no one_to_n_enrichment rows in B — only X does in A. + # Depending on implementation, this segment may appear with cost_a + # close to 0 or be omitted entirely. Either way, cost_a must be + # well below 0.50 (which would only be present if X were included). + cost_a = float(seg_row.get("cost_a", 0)) + assert cost_a < 0.50, ( + f"one_to_n_enrichment cost_a={cost_a} includes X's $0.50 cost — " + "X should have been excluded from the apples-to-apples slice" + ) + + def test_compare_corpus_change_summary_numbers(self, tmp_path): + """corpus_change_summary.csv must report the exact partition counts.""" + fixture_a = tmp_path / "a.csv" + fixture_b = tmp_path / "b.csv" + _make_compare_fixture_a(fixture_a) + _make_compare_fixture_b(fixture_b) + out_dir = tmp_path / "compare_out" + + im.run_compare(fixture_a, fixture_b, out_dir, label_a="run_a", label_b="run_b") + + summary_path = out_dir / "corpus_change_summary.csv" + assert summary_path.exists(), "corpus_change_summary.csv not produced" + df = pd.read_csv(summary_path) + assert len(df) == 1 + + row = df.iloc[0] + assert int(row["n_in_both"]) == 1, "n_in_both should be 1 (only W)" + assert ( + int(row["n_regressed"]) == 1 + ), "n_regressed should be 1 (X: success→error)" + assert int(row["n_recovered"]) == 0, "n_recovered should be 0" + # n_dropped_from_a or n_only_a — plan uses both names; accept either + only_a_val = int(row.get("n_only_a", row.get("n_dropped_from_a", -1))) + assert only_a_val == 1, "n_only_a should be 1 (Y)" + only_b_val = int(row.get("n_only_b", row.get("n_new_in_b", -1))) + assert only_b_val == 1, "n_only_b should be 1 (Z)" + + def test_compare_regression_watchlist_flags_gt_10pct_drift(self, tmp_path): + """regression_watchlist.csv must flag segments/labels with >10% cost growth + and must NOT flag those at ≤10% drift. + + Fixture: W is the only successful_in_both contract. + - W one_to_n_primary: A=$0.10, B=$0.09 → -10% (no flag — decrease, not increase) + - W preprocessing: A=not present, B=$0.08 — edge case; implementation may vary + + To produce a clear +15% case we need A and B to both have the same segment. + The fixture already has one_to_n_primary in both: + A: $0.10, B: $0.09 → -10% drift (decrease, not a regression). + + We extend the fixture with a second segment where cost grows >10%: + A preprocessing: $0.04, B preprocessing: $0.08 → +100% (flag expected). + """ + # Build custom fixtures with clear >10% growth for preprocessing + rows_a = [ + _file_start("w.txt", seq=1), + _llm_call( + "w.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.10, + seq=2, + ), + _llm_call( + "w.txt", + segment="preprocessing", + usage_label="prompt_exhibit_header", + cost_usd=0.04, + seq=3, + ), + _file_end("w.txt", seq=4), + ] + rows_b = [ + _file_start("w.txt", seq=1), + _llm_call( + "w.txt", + segment="one_to_n_primary", + usage_label="prompt_reimbursement_primary", + cost_usd=0.105, # +5% — below 10% threshold, no flag + seq=2, + ), + _llm_call( + "w.txt", + segment="preprocessing", + usage_label="prompt_exhibit_header", + cost_usd=0.046, # +15% — above 10% threshold, flag expected + seq=3, + ), + _file_end("w.txt", seq=4), + ] + + fixture_a = tmp_path / "a_watchlist.csv" + fixture_b = tmp_path / "b_watchlist.csv" + _write_fixture(fixture_a, rows_a) + _write_fixture(fixture_b, rows_b) + out_dir = tmp_path / "watchlist_out" + + im.run_compare(fixture_a, fixture_b, out_dir, label_a="run_a", label_b="run_b") + + watchlist_path = out_dir / "regression_watchlist.csv" + assert watchlist_path.exists(), "regression_watchlist.csv not produced" + df = pd.read_csv(watchlist_path) + + watchlist_items = set(df["label_or_segment"].tolist()) + + # preprocessing grew +15% → must appear on watchlist + assert "preprocessing" in watchlist_items or any( + "preprocessing" in str(v) for v in watchlist_items + ), "preprocessing (+15%) should appear on regression watchlist" + + # one_to_n_primary grew only +5% → must NOT appear on watchlist + assert "one_to_n_primary" not in watchlist_items and not any( + "one_to_n_primary" == str(v) for v in watchlist_items + ), "one_to_n_primary (+5%) should NOT appear on regression watchlist" diff --git a/src/tests/test_io_utils.py b/src/tests/test_io_utils.py index ec5040a..aec4e4b 100644 --- a/src/tests/test_io_utils.py +++ b/src/tests/test_io_utils.py @@ -778,3 +778,76 @@ class TestIOUtils: call[1]["Key"] for call in mock_s3_client.put_object.call_args_list ] assert any("QC-QA-SPLIT" in key for key in call_keys) + + +class TestUploadInstrumentationCsv: + """Tests for io_utils.upload_instrumentation_csv(run_timestamp).""" + + def test_no_op_when_file_missing(self, mocker, tmp_path): + """When the local instrumentation CSV doesn't exist, return None and skip upload.""" + from src.utils.io_utils import upload_instrumentation_csv + + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "BATCHX") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "bucket-y") + mocker.patch("src.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)) + + result = upload_instrumentation_csv("run_ts_missing") + + assert result is None + mock_s3_client.upload_file.assert_not_called() + + def test_uploads_to_canonical_s3_key(self, mocker, tmp_path): + """When the local file exists, upload_instrumentation_csv calls + S3_CLIENT.upload_file with the canonical key.""" + from src.utils.io_utils import upload_instrumentation_csv + + # Place the local file at the canonical location + local_dir = tmp_path / "run_ts_upload" / "tracking" + local_dir.mkdir(parents=True) + local_file = local_dir / "BATCHX-INSTRUMENTATION.csv" + local_file.write_text("seq,ts_iso\n1,2026-04-01T00:00:00\n") + + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "BATCHX") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "bucket-y") + mocker.patch("src.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)) + mocker.patch("src.config.WRITE_TO_S3", True) + + upload_instrumentation_csv("run_ts_upload") + + expected_s3_key = "BATCHX/run_ts_upload/tracking/BATCHX-INSTRUMENTATION.csv" + mock_s3_client.upload_file.assert_called_once_with( + str(local_file), "bucket-y", expected_s3_key + ) + + def test_respects_write_to_s3_flag_if_applicable(self, mocker, tmp_path): + """If the implementation gates on WRITE_TO_S3, verify both True and False paths. + + When WRITE_TO_S3=True: upload is attempted (if file exists). + When WRITE_TO_S3=False: upload_file is NOT called (the plan's upload helper + is meant to be invoked only from the S3 block in runner.py, so the helper + itself may or may not check the flag — both behaviours are valid). + This test documents the gating behaviour; adjust the assertion below + if the implementation does not gate on WRITE_TO_S3 internally. + """ + from src.utils.io_utils import upload_instrumentation_csv + + # Place the local file + local_dir = tmp_path / "run_ts_flag" / "tracking" + local_dir.mkdir(parents=True) + local_file = local_dir / "BATCHFLAG-INSTRUMENTATION.csv" + local_file.write_text("seq\n1\n") + + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "BATCHFLAG") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "bucket-flag") + mocker.patch("src.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path)) + + # WRITE_TO_S3=True: the helper should upload + mocker.patch("src.config.WRITE_TO_S3", True) + upload_instrumentation_csv("run_ts_flag") + assert mock_s3_client.upload_file.call_count >= 1 diff --git a/src/utils/instrumentation.py b/src/utils/instrumentation.py new file mode 100644 index 0000000..b0d9768 --- /dev/null +++ b/src/utils/instrumentation.py @@ -0,0 +1,220 @@ +import atexit +import csv +import itertools +import os +import threading +import time +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + + +def _env_disabled(var: str) -> bool: + """Return True only when the env var is *explicitly* set to a falsy value. + + Instrumentation is on by default; users opt out with + DOCZY_INSTRUMENTATION=0 (or false/no/off). + """ + return os.environ.get(var, "").lower() in ("0", "false", "no", "off") + + +FIELDNAMES = [ + "seq", + "ts_iso", + "ts_rel_sec", + "thread", + "event", + "filename", + "stage", + "segment", + "usage_label", + "model", + "exhibit_idx", + "exhibit_header", + "exhibit_page", + "page_num", + "chunk_id", + "row_id", + "n_in", + "n_out", + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost_usd", + "cache_hit_inmem", + "cache_miss_reason", + "has_context_cache", + "context_chars", + "prompt_chars", + "attempt", + "error_class", + "error_msg", + "backoff_sec", + "extra_json", +] + + +# Maps fine-grained usage_label values to coarse pipeline-phase segments. +# Used as a fallback when a call escapes any instr_scope(segment=...) block. +# Anything unmapped lands in "unknown" — a visible signal to add new prompts here. +USAGE_LABEL_TO_SEGMENT: dict[str, str] = { + # --- preprocessing (exhibit boundary detection, runs before one_to_n) --- + "EXHIBIT_HEADER": "preprocessing", # src/pipelines/saas/prompts/prompt_calls.py:605 + "EXHIBIT_HEADER_DEDUP": "preprocessing", # prompt_calls.py:649 + "EXHIBIT_LINKAGE": "preprocessing", # preprocessing_funcs.py:198 (cross-page exhibit match) + # --- one_to_n primary (exhibit_level + reimbursement-primary extraction) --- + "REIMBURSEMENT_PRIMARY": "one_to_n_primary", # prompt_calls.py:175 + "validate_reimbursements_for_llm": "one_to_n_primary", # one_to_n_funcs.py + "EXHIBIT_LEVEL": "one_to_n_primary", # prompt_calls.py:47 (prompt_exhibit_level) + "EXHIBIT_TITLE_MATCH": "one_to_n_primary", # exhibit_funcs.py:646 + # --- one_to_n enrichment (per-row cascade + breakouts) --- + "DYNAMIC_ASSIGNMENT": "one_to_n_enrichment", # prompt_calls.py:768 + "DYNAMIC_CODE_ASSIGNMENT": "one_to_n_enrichment", # prompt_calls.py:202 + "DYNAMIC_PRIMARY": "one_to_n_enrichment", # prompt_calls.py:119,149 + "METHODOLOGY_BREAKOUT": "one_to_n_enrichment", # prompt_calls.py:233 + "FEE_SCHEDULE_BREAKOUT": "one_to_n_enrichment", # prompt_calls.py:271 + "GROUPER_BREAKOUT": "one_to_n_enrichment", # prompt_calls.py:302 + "SPECIAL_CASE_BREAKOUT": "one_to_n_enrichment", # prompt_calls.py:374 + "SPECIAL_CASE_ASSIGNMENT": "one_to_n_enrichment", # prompt_calls.py:450 + "SPLIT_SERVICE_TERM": "one_to_n_enrichment", # prompt_calls.py:1318 + "SPLIT_REIMB_DATES": "one_to_n_enrichment", # prompt_calls.py:1095 + "LESSER_OF_CHECK": "one_to_n_enrichment", # prompt_calls.py:921 + "LESSER_OF_DISTRIBUTION": "one_to_n_enrichment", # prompt_calls.py:856 + "CARVEOUT_CHECK": "one_to_n_enrichment", # prompt_calls.py:333 + "prompt_lob_relationship": "one_to_n_enrichment", # prompt_calls.py:380 (inspect.stack) + # --- one_to_one (contract-level field extraction via RAG + date normalization) --- + "prompt_hsc_single_field": "one_to_one", # hybrid_smart_chunking_funcs.py:421 + "prompt_provider_info": "one_to_one", + "CHECK_PROVIDER_NAME_MATCH": "one_to_one", # prompt_calls.py:1042 + "EXTRACT_AMENDMENT_NUM_FROM_FILENAME": "one_to_one", # hybrid_smart_chunking_funcs.py:404 + # DATE_FIX / DERIVED_TERM_DATE run inside run_one_to_one_prompts, deriving + # AARETE_DERIVED_EFFECTIVE_DT / AARETE_DERIVED_TERMINATION_DT from raw dates. + # Classified by *when* they fire (one_to_one phase), not their postprocess-like role. + "DATE_FIX": "one_to_one", # prompt_calls.py:690 + "DERIVED_TERM_DATE": "one_to_one", # prompt_calls.py:724 + # Client-specific one_to_one extraction paths + "LA_CARE_SINGLE_FIELD": "one_to_one", # vendors/la_care/prompts/prompt_calls.py:92 + "LA_CARE_MULTI_FIELD": "one_to_one", # vendors/la_care/prompts/prompt_calls.py:103 + "LA_CARE_FULL_CONTEXT": "one_to_one", # vendors/la_care/prompts/prompt_calls.py:161 + # --- code breakout (all live in src/codes/code_funcs.py) --- + "CODE_EXPLICIT": "code_breakout", # code_funcs.py:114 + "SERVICE_ENRICHMENT": "code_breakout", # code_funcs.py:81 + "code_implicit_rag": "code_breakout", # code_funcs.py:352 (inspect.stack) + "code_last_check": "code_breakout", # code_funcs.py:596 (inspect.stack) + "fill_bill_type": "code_breakout", # code_funcs.py:626 (inspect.stack) + # --- postprocess (runs after all per-contract extraction completes) --- + "AARETE_DERIVED_PAYER_NAME": "postprocess", # runner.py:290 (post-run clustering) + "AARETE_DERIVED_PROVIDER_NAME": "postprocess", # runner.py:300 (post-run clustering) +} + + +def segment_for(usage_label: str | None) -> str | None: + """Look up the coarse pipeline segment for a usage_label. + + Returns "cache_warming" for any cache_warming_* label, the mapped segment + for known labels, and None when the label is unmapped (lets callers fall + back to scope-based tagging before defaulting to "unknown"). + """ + if not usage_label: + return None + if usage_label.startswith("cache_warming_"): + return "cache_warming" + return USAGE_LABEL_TO_SEGMENT.get(usage_label) + + +def resolve_segment(usage_label: str | None, ctx_segment: str | None) -> str: + """Resolve final segment: label mapping wins, scope is fallback, else 'unknown'.""" + return segment_for(usage_label) or ctx_segment or "unknown" + + +class InstrumentationLogger: + def __init__(self, csv_path: Path) -> None: + csv_path.parent.mkdir(parents=True, exist_ok=True) + self._fh = csv_path.open("w", newline="", encoding="utf-8") + self._writer = csv.DictWriter( + self._fh, fieldnames=FIELDNAMES, extrasaction="ignore" + ) + self._writer.writeheader() + self._fh.flush() + self._lock = threading.Lock() + self._seq = itertools.count(1) + self._t0 = time.monotonic() + + def log(self, event: str, **fields) -> None: + now = datetime.now(tz=timezone.utc) + row = { + "seq": next(self._seq), + "ts_iso": now.isoformat(), + "ts_rel_sec": round(time.monotonic() - self._t0, 4), + "thread": threading.current_thread().name, + "event": event, + } + row.update(fields) + with self._lock: + self._writer.writerow(row) + self._fh.flush() + + def close(self) -> None: + with self._lock: + self._fh.flush() + self._fh.close() + + +_logger: Optional[InstrumentationLogger] = None + + +def configure_from_env() -> None: + global _logger + if _env_disabled("DOCZY_INSTRUMENTATION"): + return + csv_path_str = os.environ.get("DOCZY_INSTRUMENTATION_CSV", "") + if csv_path_str: + csv_path = Path(csv_path_str) + else: + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + csv_path = Path(f"output_individual/instrumentation/run_{ts}.csv") + _logger = InstrumentationLogger(csv_path) + atexit.register(close) + + +def log(event: str, **fields) -> None: + if _logger is None: + return + _logger.log(event, **fields) + + +def is_enabled() -> bool: + return _logger is not None + + +def close() -> None: + global _logger + if _logger is not None: + _logger.close() + _logger = None + + +def configure_for_run(batch_id: str, run_timestamp: str) -> None: + """Hook called once at batch start from runner.py. + + Instrumentation is on by default; this is a no-op only when + DOCZY_INSTRUMENTATION is explicitly set to a falsy value (0/false/no/off). + When enabled, computes the canonical local path under + CONSOLIDATED_OUTPUT_DIRECTORY/{run_timestamp}/tracking/ and configures the + module-level logger to write there. Respects any pre-existing + DOCZY_INSTRUMENTATION_CSV override (for ad-hoc single-file runs). + """ + if _env_disabled("DOCZY_INSTRUMENTATION"): + return + if os.environ.get("DOCZY_INSTRUMENTATION_CSV"): + configure_from_env() + return + from src import config + + out_dir = Path(config.CONSOLIDATED_OUTPUT_DIRECTORY) / run_timestamp / "tracking" + out_dir.mkdir(parents=True, exist_ok=True) + csv_path = out_dir / f"{batch_id}-INSTRUMENTATION.csv" + os.environ["DOCZY_INSTRUMENTATION_CSV"] = str(csv_path) + configure_from_env() diff --git a/src/utils/instrumentation_context.py b/src/utils/instrumentation_context.py new file mode 100644 index 0000000..f7cfd57 --- /dev/null +++ b/src/utils/instrumentation_context.py @@ -0,0 +1,46 @@ +from contextvars import ContextVar, copy_context +from typing import Any + +_ctx: ContextVar[dict] = ContextVar("instr_ctx", default={}) + + +class instr_scope: + def __init__(self, **kwargs: Any) -> None: + self.kwargs = kwargs + + def __enter__(self) -> None: + self._token = _ctx.set({**_ctx.get(), **self.kwargs}) + + def __exit__(self, *_: Any) -> None: + _ctx.reset(self._token) + + +def get_ctx() -> dict: + return _ctx.get() + + +def submit_with_context(executor, fn, *args, **kwargs): + """Submit a task that preserves the current contextvars into the worker. + Use in place of executor.submit(fn, *args, **kwargs) when the task + should see the caller's instr_scope context. + """ + ctx = copy_context() + return executor.submit(ctx.run, fn, *args, **kwargs) + + +def map_with_context(executor, fn, *iterables): + """Like executor.map(fn, *iterables) but preserves the instrumentation + ContextVar (_ctx) into each worker task. Safe for parallel workers + because each worker sets/resets the var on its own thread, without + sharing a Context object. + """ + parent_vars = _ctx.get() + + def _wrapped(*args): + token = _ctx.set(parent_vars) + try: + return fn(*args) + finally: + _ctx.reset(token) + + return executor.map(_wrapped, *iterables) diff --git a/src/utils/io_utils.py b/src/utils/io_utils.py index 383dec9..42f09b4 100644 --- a/src/utils/io_utils.py +++ b/src/utils/io_utils.py @@ -1342,6 +1342,46 @@ def upload_local_file_to_s3( return None +def upload_instrumentation_csv(run_timestamp: str) -> Optional[str]: + """Upload the instrumentation CSV for a run to S3. + + The CSV is streamed to disk during the run; this function uploads the + already-existing local file. No-op when WRITE_TO_S3 is False or the + local file does not exist (instrumentation disabled). + + Args: + run_timestamp: Run timestamp (e.g. run_20260212_13-53_test_batch). + + Returns: + S3 URI on success, None otherwise. + """ + local_path = ( + Path(config.CONSOLIDATED_OUTPUT_DIRECTORY) + / run_timestamp + / "tracking" + / f"{config.BATCH_ID}-INSTRUMENTATION.csv" + ) + if not local_path.exists(): + logging.debug( + f"Instrumentation CSV not found, skipping S3 upload: {local_path}" + ) + return None + if not config.WRITE_TO_S3: + return None + s3_key = ( + f"{config.BATCH_ID}/{run_timestamp}/tracking/" + f"{config.BATCH_ID}-INSTRUMENTATION.csv" + ) + try: + config.S3_CLIENT.upload_file(str(local_path), config.S3_OUTPUT_BUCKET, s3_key) + s3_uri = f"s3://{config.S3_OUTPUT_BUCKET}/{s3_key}" + logging.info(f"Uploaded instrumentation CSV to S3: {s3_uri}") + return s3_uri + except ClientError as e: + logging.error(f"Failed to upload instrumentation CSV to S3: {e}") + return None + + def save_result_to_json(filename, result, json_folder): """Save a result dictionary to an individual JSON file in the specified folder. diff --git a/src/utils/llm_utils.py b/src/utils/llm_utils.py index a6bb2fb..0019832 100644 --- a/src/utils/llm_utils.py +++ b/src/utils/llm_utils.py @@ -119,6 +119,110 @@ def _supports_prompt_cache(model_id: str) -> bool: return model_id in {m for m in supported_ids if m is not None} +def _classify_cache_outcome( + *, + has_context_cache: bool, + cache_read_tokens: int, + cache_creation_tokens: int, + context_chars: int, + cache_key: str | None, +) -> str: + """Classify why a cache-attempting call did or didn't hit the Anthropic cache. + + Values: hit, first_call, ttl_expired, under_min_tokens, silent_miss, not_attempted. + Anthropic's minimum cacheable block is 1024 tokens; we approximate as + context_chars / 3.5. Updates _cache_key_seen on every classify for future calls. + """ + if not has_context_cache: + return "not_attempted" + if cache_read_tokens > 0: + if cache_key is not None: + with _cache_key_seen_lock: + _cache_key_seen[cache_key] = time.monotonic() + return "hit" + if cache_creation_tokens > 0: + # Fresh write — either this is the first call with this key, or a prior + # entry expired (TTL ~5 min) and had to be re-written. + if cache_key is not None: + with _cache_key_seen_lock: + prior_seen = cache_key in _cache_key_seen + _cache_key_seen[cache_key] = time.monotonic() + return "ttl_expired" if prior_seen else "first_call" + return "first_call" + # has_context_cache=True but both token counts are 0 — silent failure mode. + approx_tokens = int(context_chars / 3.5) if context_chars else 0 + if approx_tokens and approx_tokens < 1024: + return "under_min_tokens" + return "silent_miss" + + +def _ctx_minus_explicit_keys(ctx: dict, *explicit_keys: str) -> dict: + """Filter a ctx dict so explicit kwargs aren't duplicated via **ctx splat.""" + skip = {"_usage_label", "segment", *explicit_keys} + return {k: v for k, v in ctx.items() if not k.startswith("_") and k not in skip} + + +def _log_llm_retry( + filename: str, + model_id: str, + attempt: int, + error_class: str, + error_msg: str, + backoff_sec: float, +) -> None: + """Emit one llm_retry instrumentation event (no-op when disabled).""" + from src.utils import instrumentation, instrumentation_context + + if not instrumentation.is_enabled(): + return + ctx = dict(instrumentation_context.get_ctx()) + usage_label = ctx.get("_usage_label") + ctx_segment = ctx.get("segment") + instrumentation.log( + "llm_retry", + filename=filename, + model=model_id, + usage_label=usage_label, + segment=instrumentation.resolve_segment(usage_label, ctx_segment), + attempt=attempt, + error_class=error_class, + error_msg=(error_msg or "")[:200], + backoff_sec=round(backoff_sec, 4), + **_ctx_minus_explicit_keys(ctx, "filename", "model", "attempt"), + ) + + +def _log_llm_error( + filename: str, + model_id: str, + attempt: int, + error_class: str, + error_msg: str, +) -> None: + """Emit one llm_error instrumentation event (no-op when disabled). + + Fires when retries are exhausted and the call is about to raise. + """ + from src.utils import instrumentation, instrumentation_context + + if not instrumentation.is_enabled(): + return + ctx = dict(instrumentation_context.get_ctx()) + usage_label = ctx.get("_usage_label") + ctx_segment = ctx.get("segment") + instrumentation.log( + "llm_error", + filename=filename, + model=model_id, + usage_label=usage_label, + segment=instrumentation.resolve_segment(usage_label, ctx_segment), + attempt=attempt, + error_class=error_class, + error_msg=(error_msg or "")[:200], + **_ctx_minus_explicit_keys(ctx, "filename", "model", "attempt"), + ) + + def _track_usage_from_response( response_body: dict, filename: str, model_id: str ) -> None: @@ -153,6 +257,81 @@ def _track_usage_from_response( cache_creation_tokens, cache_read_tokens, ) + # One llm_call event per Bedrock response with non-zero tokens. + from src.utils import instrumentation, instrumentation_context + + if instrumentation.is_enabled(): + from src.utils.usage_tracking import get_cost_per_token + + input_cpt, output_cpt, cache_read_cpt = get_cost_per_token(model_id) + if "sonnet" in model_id.lower(): + cache_write_cpt = input_cpt * 1.25 + elif "haiku" in model_id.lower(): + cache_write_cpt = input_cpt * 1.20 + else: + cache_write_cpt = input_cpt + cost = ( + input_tokens * input_cpt + + output_tokens * output_cpt + + cache_read_tokens * cache_read_cpt + + cache_creation_tokens * cache_write_cpt + ) + ctx = dict(instrumentation_context.get_ctx()) + usage_label = ctx.pop("_usage_label", None) + has_context_cache = ctx.pop("_has_context_cache", False) + context_chars = ctx.pop("_context_chars", 0) + prompt_chars = ctx.pop("_prompt_chars", 0) + cache_key = ctx.pop("_cache_key", None) + segment = instrumentation.resolve_segment( + usage_label, ctx.pop("segment", None) + ) + cache_miss_reason = _classify_cache_outcome( + has_context_cache=has_context_cache, + cache_read_tokens=cache_read_tokens, + cache_creation_tokens=cache_creation_tokens, + context_chars=context_chars, + cache_key=cache_key, + ) + instrumentation.log( + "llm_call", + filename=filename, + model=model_id, + usage_label=usage_label, + segment=segment, + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_creation_tokens=cache_creation_tokens, + cache_read_tokens=cache_read_tokens, + total_tokens=( + input_tokens + + output_tokens + + cache_creation_tokens + + cache_read_tokens + ), + cost_usd=cost, + cache_hit_inmem=False, + cache_miss_reason=cache_miss_reason, + has_context_cache=has_context_cache, + context_chars=context_chars, + prompt_chars=prompt_chars, + **_ctx_minus_explicit_keys( + ctx, + "filename", + "model", + "usage_label", + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost_usd", + "cache_hit_inmem", + "cache_miss_reason", + "has_context_cache", + "context_chars", + "prompt_chars", + ), + ) def _extract_text_from_claude_response(response_body: dict) -> str: @@ -279,6 +458,12 @@ claude_cache = OrderedDict() CACHE_LIMIT = 20000 _cache_lock = threading.Lock() # Thread safety for cache operations +# Tracks which Anthropic cache-key hashes have been seen in this process and +# when (monotonic seconds). Used to classify cache_miss_reason=first_call vs +# ttl_expired on llm_call events. O(distinct cache keys) — tiny. +_cache_key_seen: dict[str, float] = {} +_cache_key_seen_lock = threading.Lock() + # Model routing set (derived from _SUPPORTED_MODELS for consistency) _CLAUDE_3_AND_UP_MODELS = {m for m in _SUPPORTED_MODELS if m is not None} @@ -376,6 +561,37 @@ def invoke_claude( with _cache_lock: if cache_key in claude_cache: claude_cache.move_to_end(cache_key) # Mark as recently used + # One llm_call_inmem_hit event per in-memory cache return. + from src.utils import instrumentation, instrumentation_context + + if instrumentation.is_enabled(): + ctx = instrumentation_context.get_ctx() + segment = instrumentation.resolve_segment( + usage_label, ctx.get("segment") + ) + instrumentation.log( + "llm_call_inmem_hit", + filename=filename, + usage_label=usage_label, + segment=segment, + model=model_id, + prompt_chars=len(prompt) if prompt else 0, + context_chars=( + len(context_for_caching) if context_for_caching else 0 + ), + has_context_cache=bool(context_for_caching), + cache_hit_inmem=True, + **_ctx_minus_explicit_keys( + dict(ctx), + "filename", + "usage_label", + "model", + "prompt_chars", + "context_chars", + "has_context_cache", + "cache_hit_inmem", + ), + ) return claude_cache[cache_key] # Resolve alias to actual model ID @@ -384,31 +600,43 @@ def invoke_claude( # Validate model support (checks for deprecated and unsupported models) _validate_model_support(model_id, resolved_model_id) - if config.RUN_MODE == "local": - response = local_claude_3_and_up( - prompt, - resolved_model_id, - filename, - max_tokens, - cache=cache, - instruction=instruction, - usage_label=usage_label, - context_for_caching=context_for_caching, - ) - elif config.RUN_MODE == "ec2": - response = ec2_claude_3_and_up( - prompt, - resolved_model_id, - filename, - max_tokens, - cache=cache, - instruction=instruction, - usage_label=usage_label, - context_for_caching=context_for_caching, - ) - else: - logging.error("Usage: python local_main.py <-local> OR python main.py ") - raise + # Push call metadata into context so _track_usage_from_response can read it. + from src.utils.instrumentation_context import instr_scope + + with instr_scope( + _usage_label=usage_label, + _has_context_cache=bool(context_for_caching), + _context_chars=len(context_for_caching) if context_for_caching else 0, + _prompt_chars=len(prompt) if prompt else 0, + _cache_key=cache_key, + ): + if config.RUN_MODE == "local": + response = local_claude_3_and_up( + prompt, + resolved_model_id, + filename, + max_tokens, + cache=cache, + instruction=instruction, + usage_label=usage_label, + context_for_caching=context_for_caching, + ) + elif config.RUN_MODE == "ec2": + response = ec2_claude_3_and_up( + prompt, + resolved_model_id, + filename, + max_tokens, + cache=cache, + instruction=instruction, + usage_label=usage_label, + context_for_caching=context_for_caching, + ) + else: + logging.error( + "Usage: python local_main.py <-local> OR python main.py " + ) + raise # Store response in cache with LRU eviction - thread-safe _store_in_cache(cache_key, response) @@ -589,25 +817,35 @@ def local_claude_3_and_up( except ValueError as e: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, model_id, attempt + 1, "ValueError", str(e), delay + ) logging.warning( f"Attempt {attempt + 1} failed due to malformed Claude response. " f"Retrying in {delay:.2f} seconds. Error: {e}" ) time.sleep(delay) else: + _log_llm_error(filename, model_id, attempt + 1, "ValueError", str(e)) raise except ClientError as e: - if e.response["Error"]["Code"] in [ + error_code = e.response["Error"]["Code"] + if error_code in [ "ThrottlingException", "TooManyRequestsException", ]: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, model_id, attempt + 1, error_code, str(e), delay + ) time.sleep(delay) else: + _log_llm_error(filename, model_id, attempt + 1, error_code, str(e)) raise else: + _log_llm_error(filename, model_id, attempt + 1, error_code, str(e)) raise raise Exception("Max retries exceeded") @@ -820,12 +1058,23 @@ def ec2_claude_3_and_up( except ValueError as e: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, model_id, attempt + 1, "ValueError", str(e), delay + ) logging.warning( f"Attempt {attempt + 1} failed due to malformed Claude response. " f"Retrying in {delay:.2f} seconds... Error: {e}" ) time.sleep(delay) else: + _log_llm_retry( + filename, + model_id, + attempt + 1, + f"runtime_switch:{runtime_key}:ValueError", + str(e), + 0.0, + ) logging.warning( f"Switching to next runtime due to malformed Claude response: {e}." ) @@ -839,20 +1088,46 @@ def ec2_claude_3_and_up( ]: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, + model_id, + attempt + 1, + error_code, + str(e), + delay, + ) logging.warning( f"Attempt {attempt + 1} failed. Retrying in {delay:.2f} seconds..." ) time.sleep(delay) else: + _log_llm_retry( + filename, + model_id, + attempt + 1, + f"runtime_switch:{runtime_key}:{error_code}", + str(e), + 0.0, + ) logging.warning( f"Switching to next runtime due to {error_code}: {str(e)}." ) break # Switch to the next runtime else: + _log_llm_error( + filename, model_id, attempt + 1, error_code, str(e) + ) logging.error(f"Unexpected error: {str(e)}") raise # Reraise other exceptions except Exception as e: + _log_llm_error( + filename, + model_id, + attempt + 1, + type(e).__name__, + str(e), + ) logging.error(f"Unexpected error: {str(e)}") raise else: # Runtime rotation is disabled @@ -875,40 +1150,76 @@ def ec2_claude_3_and_up( except ValueError as e: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, model_id, attempt + 1, "ValueError", str(e), delay + ) logging.warning( f"Attempt {attempt + 1} failed due to malformed Claude response. " f"Retrying in {delay:.2f} seconds... Error: {e}" ) time.sleep(delay) else: + _log_llm_error( + filename, model_id, attempt + 1, "ValueError", str(e) + ) raise except ReadTimeoutError as e: # Timeout errors should probably retry if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, + model_id, + attempt + 1, + "ReadTimeoutError", + str(e), + delay, + ) logging.warning( f"Attempt {attempt + 1} failed due to timeout. Retrying in {delay:.2f} seconds..." ) time.sleep(delay) else: + _log_llm_error( + filename, + model_id, + attempt + 1, + "ReadTimeoutError", + str(e), + ) raise except ClientError as e: error_code = e.response["Error"]["Code"] if error_code in ["ThrottlingException", "TooManyRequestsException"]: if attempt < max_retries - 1: delay = (2**attempt + random.uniform(0, 1)) * initial_delay + _log_llm_retry( + filename, + model_id, + attempt + 1, + error_code, + str(e), + delay, + ) logging.warning( f"Attempt {attempt + 1} failed due to rate limiting. Retrying in {delay:.2f} seconds..." ) time.sleep(delay) else: + _log_llm_error( + filename, model_id, attempt + 1, error_code, str(e) + ) raise else: + _log_llm_error(filename, model_id, attempt + 1, error_code, str(e)) logging.error(f"Unexpected error: {str(e)}") raise # Non-throttling client errors should fail immediately except Exception as e: + _log_llm_error( + filename, model_id, attempt + 1, type(e).__name__, str(e) + ) logging.error(f"Unexpected error: {str(e)}") raise raise Exception(f"Max retries ({max_retries}) exceeded for model {model_id}.")