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

Merged in bugfix/exhibit-smart-chunking-cost-improvements (pull request #985)

Browse Source 
Bugfix/exhibit smart chunking cost improvements

* Add opt-in instrumentation for per-call token and row-count tracing

Introduce src/utils/instrumentation.py (thread-safe CSV logger) and
src/utils/instrumentation_context.py (ContextVar scope plus
submit_with_context / map_with_context helpers for propagating context
into ThreadPoolExecutor workers).

Emit events at every Bedrock call in llm_utils.invoke_claude, including
in-memory claude_cache hits, with full input/output/cache-read/cache-write
token breakdown. Emit row-count events at each row-mutating stage in the
one-to-N pipeline (clean_reimbursement_primary,
filter_services_without_reimbursements, methodology_breakout,
split_service_terms, carveout, dynamic_code_assignment,
lesser_of_distribution, dynamic_assignment) and chunking / retrieval
events in exhibit smart chunking (chunking_done, retrieval_done) plus
exhibit lifecycle events (exhibit_start, exhibit_gate_skip,
stage_transition).

All hooks are no-ops unless DOCZY_INSTRUMENTATION=1; production defaults
unchanged. Runner and inspector scripts in …
* Extend instrumentation: segment, cache_miss_reason, retry/error events, runtime hookup, analyzer

Schema expansion and new event types:
- Add segment column to every llm_call / llm_call_inmem_hit, resolved from
  USAGE_LABEL_TO_SEGMENT (authoritative) with scope fallback. Mapping built
  from grep + smoke-run ground truth; earlier guessed entries removed.
- Add cache_miss_reason classifier: hit / first_call / ttl_expired /
  under_min_tokens / silent_miss / not_attempted. Uses a per-process
  _cache_key_seen dict guarded by its own lock.
- Emit llm_retry on each retry attempt (attempt, error_class, error_msg,
  backoff_sec) and llm_error on exhausted retries in ec2_claude_3_and_up
  (rotation and non-rotation branches) plus local_claude_3_and_up.
- Add six new CSV columns: segment, cache_miss_reason, attempt,
  error_class, error_msg, backoff_sec. Total schema now 34 columns.

Defensive kwarg hygiene:
- _ctx_minus_explicit_keys filter on all emit sites to prevent
  segment / filename kwarg collisions between …
* Merged dev into bugfix/exhibit-smart-chunking-cost-improvements

* Merged dev into bugfix/exhibit-smart-chunking-cost-improvements

* Make instrumentation tracking on by default

Flip the gate: instrumentation is now enabled unless DOCZY_INSTRUMENTATION
is explicitly set to a falsy value (0/false/no/off). Replaces the prior
opt-in behaviour where it was off unless DOCZY_INSTRUMENTATION=1 was set.

* Fix missing WRITE_TO_S3 patch in upload_instrumentation_csv test

* Merged dev into bugfix/exhibit-smart-chunking-cost-improvements


Approved-by: Katon Minhas
This commit is contained in:
Praneel Panchigar
2026-04-28 15:58:21 +00:00
committed by Katon Minhas
parent 7a39dabcb7
commit f1df468587
17 changed files with 4028 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config.py
+10
View File
@@ -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", "")
src/pipelines/runner.py
+9
View File
@@ -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)
src/pipelines/saas/file_processing.py
+327 -180
View File
@@ -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}"
)
src/pipelines/saas/main.py
+3
View File
@@ -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)
src/pipelines/shared/extraction/dynamic_funcs.py
+41 -20
View File
@@ -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}"
src/pipelines/shared/extraction/exhibit_funcs.py
+20
View File
@@ -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
src/pipelines/shared/extraction/one_to_n_funcs.py
+54 -13
View File
@@ -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
src/pipelines/shared/preprocessing/exhibit_smart_chunking_funcs.py
+15
View File
@@ -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
src/pipelines/shared/preprocessing/hybrid_smart_chunking_funcs.py
+9 -2
View File
@@ -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,
src/testbed/instrumentation_metrics.py
+1383
View File
File diff suppressed because it is too large Load Diff
src/tests/test_instrumentation.py
+697
View File
@@ -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
src/tests/test_instrumentation_metrics.py
+744
View File
@@ -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"
src/tests/test_io_utils.py
+73
View File
@@ -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
src/utils/instrumentation.py
+220
View File
@@ -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()
src/utils/instrumentation_context.py
+46
View File
@@ -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)
src/utils/io_utils.py
+40
View File
@@ -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.
src/utils/llm_utils.py
+337 -26
View File
@@ -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 <ec2>")
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 <ec2>"
)
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}.")