From 7056c1c6878a4be2d7ff1a2aa378309827c1c1b2 Mon Sep 17 00:00:00 2001 From: Praneel Panchigar Date: Mon, 10 Nov 2025 19:29:45 +0000 Subject: [PATCH] Merged in feature/automated-cost-logging-csv (pull request #768) Feature/automated cost logging csv * feat: Add usage and cost tracking with CSV export - Add usage_tracking.py module for thread-safe token and cost tracking - Extract actual tokens from Bedrock API responses (replaces word count estimation) - Integrate usage tracking into llm_utils.py for all LLM invocations - Add CSV export functionality (per-file/per-model and batch summary) - Integrate CSV export into main.py for local runs - Update .gitignore to exclude PRD Phase 1: Data collection and local CSV export implemented - Tracks all LLM calls including utility functions (date_fix, derive_term_date) - Calculates costs per model using centralized cost constants - Generates two CSV files: USAGE.csv and USAGE-SUMMARY.csv - Exports to same directory as final results when write_to_s3=False * - Removed some temporary logging code - Removed unused imports (uuid, datetime) from llm_utils.py - Added cache token columns (cache_creation_tokens, cache_read_tokens) to CSV exports - Added tracking for cache warming - Updated usage_tracking.py to track cache tokens separately in data structure All token tracking functionality remains intact. * Split input tokens into fresh and cache in summary section - Add total_fresh_input_tokens and total_cache_tokens to GLOBAL_USAGE - Update summary CSV to include split token columns for better tractability - Add averages for fresh_input_tokens and cache_tokens - Maintain total_input_tokens = fresh_input_tokens + cache_tokens relationship - All calculations verified and match per-file totals correctly * upload files functionality * Add comprehensive unit tests for usage_tracking module - Implement Phase 1-4 tests covering all 10 functions in usage_tracking.py - Add 98 test cases with parametrized tests for comprehensive coverage - Fix extract_usage_from_response to handle None input gracefully - Add python-dotenv dependency to pyproject.toml - Tests cover: core functions, data access, export functions, and utilities - All tests passing (98/98) * Add S3 upload tests for main.py and fix f-string syntax error - Add tests/test_investment_main_s3.py with 3 test cases: * test_main_writes_final_and_error_to_s3: validates both final and error DataFrames uploaded when processing has errors * test_main_writes_only_final_when_no_errors: validates only final DataFrame uploaded when all files succeed * test_main_writes_usage_when_present: validates usage tracking files uploaded when usage data exists - Tests use monkeypatching and module stubs to isolate S3 write logic without external dependencies - Fix syntax error in src/codes/code_funcs.py: nested f-string quotes (level_dict['level_suffix']) - All tests passing * Add comprehensive test for usage_tracking.get_usage_dataframes S3 uploads - Add test_usage_dataframes_written_to_s3 to validate per_file_usage_df and batch_summary_df - Test verifies both dataframes from usage_tracking.get_usage_dataframes() are written to S3 - Validates 'usage' suffix writes per_file_usage_df with correct schema (file_name, tokens, cost) - Validates 'usage_summary' suffix writes batch_summary_df with correct schema (totals, averages, region_mode) - Uses pd.testing.assert_frame_equal to ensure exact dataframes are uploaded - Fix module caching issue between tests by clearing sys.modules cache - All 4 tests passing * Add cleanup fixture to prevent test pollution - Add autouse pytest fixture to clean up stubbed modules after each test - Remove manual cache clearing from individual tests (now handled by fixture) - Prevents our module stubs from affecting other test files in CI/CD pipeline - Fixes Bitbucket pipeline failures where other tests couldn't find llm_utils attributes * Refactor tests: add comprehensive io_utils coverage; disable usage_tracking suite - Added 22 focused tests for write_local and write_s3 plus existing IO behaviors - Introduced FakeS3Client to avoid boto3 network/credential dependency - Added preserve_config fixture (no monkeypatch) to isolate config side effects - Marked test_usage_tracking.py skipped per new consolidation approach - Verified 26 tests (io_utils + investment_main) all pass locally without AWS exceptions * Fix S3_CLIENT mocking: use mocker.patch instead of direct assignment - Changed preserve_config fixture to NOT save/restore S3_CLIENT - Updated all 5 write_s3 tests to use mocker.patch('src.config.S3_CLIENT', fake) - This prevents pollution of config.S3_CLIENT across test modules - Fixes NoCredentialsError in other tests that were importing config after our tests set FakeS3Client - All 22 io_utils tests pass + 4 investment_main_s3 tests pass - test_llm_utils::test_invoke_claude_local_mode now passes * Add tests for write_s3 function * Resolve merge conflict: keep mocker-based write_s3 tests * Remove rogue skip mark from test_usage_tracking.py The skip mark was incorrectly added by remote branch claiming tests were migrated to io_utils. However, test_io_utils.py only tests io_utils functions (read/write operations), not usage_tracking functions. Restore the comprehensive usage_tracking tests from commit 26a556f0 which includes 98 test cases covering all 10 functions in usage_tracking.py. * Fix test pollution: add config cleanup fixture and restore RUN_MODE - Add autouse fixture in test_io_utils.py to restore config values after each test - Fix test_llm_utils.py to restore RUN_MODE after test_invoke_claude_local_mode - Prevents config patches from leaking into subsequent tests causing NoCredentialsError * Remove redundant test_investment_main_s3.py - File was testing same write_s3 functionality already covered in test_io_utils.py - Was causing module pollution by stubbing src.utils.llm_utils - Caused test failures in other test files due to module cache issues - Integration testing was minimal and mocked write_s3 anyway * Merged main into feature/automated-cost-logging-csv Approved-by: Karan Desai Approved-by: Katon Minhas --- .gitignore | 10 +- fieldExtraction/src/codes/code_funcs.py | 4 +- fieldExtraction/src/config.py | 1 + fieldExtraction/src/investment/main.py | 20 + fieldExtraction/src/utils/io_utils.py | 28 +- fieldExtraction/src/utils/llm_utils.py | 175 +- fieldExtraction/src/utils/usage_tracking.py | 476 ++++++ fieldExtraction/tests/test_io_utils.py | 116 +- fieldExtraction/tests/test_llm_utils.py | 14 +- fieldExtraction/tests/test_usage_tracking.py | 1600 ++++++++++++++++++ poetry.lock | 79 +- pyproject.toml | 1 + 12 files changed, 2487 insertions(+), 37 deletions(-) create mode 100644 fieldExtraction/src/utils/usage_tracking.py create mode 100644 fieldExtraction/tests/test_usage_tracking.py diff --git a/.gitignore b/.gitignore index e71e4b4..365cb50 100644 --- a/.gitignore +++ b/.gitignore @@ -23,8 +23,9 @@ dist/ target/ dist/ -# JetBrains IDE +# IDE files .idea/ +.cursor/ # Unit test reports TEST*.xml @@ -98,4 +99,9 @@ __pycache__/ new/ # Embeddings -fieldExtraction/embeddings/ \ No newline at end of file +fieldExtraction/embeddings/ + +# PRD Documents (local only, living documents) +*.prd +*prd.md +usage-cost-monitoring-prd.md \ No newline at end of file diff --git a/fieldExtraction/src/codes/code_funcs.py b/fieldExtraction/src/codes/code_funcs.py index df38605..386db2b 100644 --- a/fieldExtraction/src/codes/code_funcs.py +++ b/fieldExtraction/src/codes/code_funcs.py @@ -356,13 +356,13 @@ def code_implicit_rag(service, implicit_run_dict, filename, constants): code_answer_dict["PROCEDURE_CD"] = "|".join(proc_codes) code_answer_dict["PROCEDURE_CD_DESC"] = "|".join(proc_descs) code_answer_dict["CODE_METHODOLOGY"] = ( - f"Implicit - Level {level_dict["level_suffix"]}" + f"Implicit - Level {level_dict['level_suffix']}" ) if rev_codes: code_answer_dict["REVENUE_CD"] = "|".join(rev_codes) code_answer_dict["REVENUE_CD_DESC"] = "|".join(rev_descs) code_answer_dict["CODE_METHODOLOGY"] = ( - f"Implicit - Level {level_dict["level_suffix"]}" + f"Implicit - Level {level_dict['level_suffix']}" ) if code_answer_dict: diff --git a/fieldExtraction/src/config.py b/fieldExtraction/src/config.py index 86d7767..f5a6513 100644 --- a/fieldExtraction/src/config.py +++ b/fieldExtraction/src/config.py @@ -259,6 +259,7 @@ def resolve_model_id(model_identifier: str) -> str: ENABLE_ANSWER_TRACKING = False ANSWER_TRACKING_DICT = {} + ######################################## POSTPROCESSING SETTINGS ######################################## FUZZY_MATCH_THRESHOLD = 0.8 FILE_NAME_COLUMN = "Contract Name" diff --git a/fieldExtraction/src/investment/main.py b/fieldExtraction/src/investment/main.py index 02a2caa..999428e 100644 --- a/fieldExtraction/src/investment/main.py +++ b/fieldExtraction/src/investment/main.py @@ -14,6 +14,7 @@ import src.prompts.prompt_templates as prompt_templates import src.utils.io_utils as io_utils import src.utils.llm_utils as llm_utils import src.utils.logging_utils as logging_utils +import src.utils.usage_tracking as usage_tracking from constants.constants import Constants from src import config @@ -151,18 +152,37 @@ def main(testing=False, test_params={}): else: ERROR_RESULT_DF = pd.DataFrame() + if not testing: + # Get usage tracking dataframes + per_file_usage_df, batch_summary_df = usage_tracking.get_usage_dataframes( + config.BATCH_ID, run_timestamp + ) + if config.WRITE_TO_S3: io_utils.write_s3(FINAL_RESULT_DF, "", run_timestamp, "final") if not ERROR_RESULT_DF.empty: io_utils.write_s3(ERROR_RESULT_DF, "", run_timestamp, "error") + + # Export usage and cost tracking data to S3 only if both dataframes have data + if not per_file_usage_df.empty and not batch_summary_df.empty: + logging.info("Exporting usage and cost tracking data to S3...") + io_utils.write_s3(per_file_usage_df, "", run_timestamp, "usage") + io_utils.write_s3(batch_summary_df, "", run_timestamp, "usage_summary") else: io_utils.write_local(FINAL_RESULT_DF, "", run_timestamp, "final") if not ERROR_RESULT_DF.empty: io_utils.write_local(ERROR_RESULT_DF, "", run_timestamp, "error") + + # Export usage and cost tracking data locally only if both dataframes have data + if not per_file_usage_df.empty and not batch_summary_df.empty: + logging.info("Exporting usage and cost tracking data locally...") + io_utils.write_local(per_file_usage_df, "", run_timestamp, "usage") + io_utils.write_local(batch_summary_df, "", run_timestamp, "usage_summary") else: return FINAL_RESULT_DF, ERROR_RESULT_DF if __name__ == "__main__": main() + diff --git a/fieldExtraction/src/utils/io_utils.py b/fieldExtraction/src/utils/io_utils.py index b00ba7a..f3a4fce 100644 --- a/fieldExtraction/src/utils/io_utils.py +++ b/fieldExtraction/src/utils/io_utils.py @@ -380,9 +380,27 @@ def write_local( quoting=1, ) return None + elif output_type == "usage": + output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "tracking") + os.makedirs(output_dir, exist_ok=True) + df.to_csv( + os.path.join(output_dir, f"{config.BATCH_ID}-USAGE.csv"), + index=False, + quoting=1, + ) + return None + elif output_type == "usage_summary": + output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp, "tracking") + os.makedirs(output_dir, exist_ok=True) + df.to_csv( + os.path.join(output_dir, f"{config.BATCH_ID}-USAGE-SUMMARY.csv"), + index=False, + quoting=1, + ) + return None else: logging.error( - f"Unknown output_type: {output_type}; output_type must be 'final', 'individual', or 'error'" + f"Unknown output_type: {output_type}; output_type must be 'final', 'individual', 'error', 'usage', or 'usage_summary'" ) return None @@ -423,6 +441,10 @@ def write_s3( ) elif output_type == "error": output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-ERRORS.csv" + elif output_type == "usage": + output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE.csv" + elif output_type == "usage_summary": + output_path = f"{config.BATCH_ID}/{run_timestamp}/tracking/{config.BATCH_ID}-USAGE-SUMMARY.csv" else: # Fallback path for unknown output types output_path = f"{config.BATCH_ID}/{run_timestamp}/{config.BATCH_ID}-unknown-{output_type}.csv" @@ -438,6 +460,10 @@ def write_s3( logging.info(f"Saved individual output file: {output_path}") elif output_type == "error": logging.info(f"Saved error output file: {output_path}") + elif output_type == "usage": + logging.info(f"Saved usage tracking file: {output_path}") + elif output_type == "usage_summary": + logging.info(f"Saved usage summary file: {output_path}") else: logging.info(f"Saved unknown output type file: {output_path}") return None diff --git a/fieldExtraction/src/utils/llm_utils.py b/fieldExtraction/src/utils/llm_utils.py index 36345d0..10ed98e 100644 --- a/fieldExtraction/src/utils/llm_utils.py +++ b/fieldExtraction/src/utils/llm_utils.py @@ -4,6 +4,7 @@ import inspect import json import logging import math +import os import random import threading import time @@ -11,6 +12,7 @@ from collections import OrderedDict import anthropic import src.utils.string_utils as string_utils +import src.utils.usage_tracking as usage_tracking from botocore.exceptions import ClientError, ReadTimeoutError from src import config @@ -374,6 +376,29 @@ def local_claude_2(prompt, model_id, filename, max_tokens): response_body = json.loads(response.get("body").read()) response_text = response_body["completion"] + # Track usage and costs from API response (Claude 2 does not support cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + if config.ENABLE_ANSWER_TRACKING: config.ANSWER_TRACKING_DICT[prompt] = response_text @@ -458,11 +483,31 @@ def local_claude_3_and_up( contentType="application/json", ) response_body = json.loads(response.get("body").read()) - # Print usage stats if available for debugging (cache creation/read tokens, etc.) - if "usage" in response_body: - label = usage_label or "UNLABELED" response_text = response_body["content"][0]["text"] + # Track usage and costs from API response (including cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + if config.ENABLE_ANSWER_TRACKING: config.ANSWER_TRACKING_DICT[prompt] = response_text @@ -552,11 +597,31 @@ def local_claude_3_and_up_with_image_array( contentType="application/json", ) response_body = json.loads(response.get("body").read()) - # Print usage stats if available for debugging (cache creation/read tokens, etc.) - if "usage" in response_body: - label = usage_label or "UNLABELED" response_text = response_body["content"][0]["text"] + # Track usage and costs from API response (including cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + if config.ENABLE_ANSWER_TRACKING: config.ANSWER_TRACKING_DICT[prompt] = response_text @@ -598,6 +663,30 @@ def ec2_claude2(prompt, model_id, filename, max_tokens): ) response_body = json.loads(response.get("body").read()) response_text = response_body["completion"] + + # Track usage and costs from API response (Claude 2 does not support cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + elapsed_time = time.time() - start_time tokens_sent = len(prompt.split()) tokens_received = len(response_text.split()) @@ -714,11 +803,31 @@ def ec2_claude_3_and_up( contentType="application/json", ) response_body = json.loads(response.get("body").read()) - # Print usage stats if available for debugging (cache creation/read tokens, etc.) - if "usage" in response_body: - label = usage_label or "UNLABELED" response_text = response_body["content"][0]["text"] + # Track usage and costs from API response (including cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + elapsed_time = time.time() - start_time tokens_sent = len(prompt.split()) tokens_received = len(response_text.split()) @@ -780,11 +889,31 @@ def ec2_claude_3_and_up( contentType="application/json", ) response_body = json.loads(response.get("body").read()) - # Print usage stats if available for debugging (cache creation/read tokens, etc.) - if "usage" in response_body: - label = usage_label or "UNLABELED" response_text = response_body["content"][0]["text"] + # Track usage and costs from API response (including cache tokens) + if filename: + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + filename, + model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + elapsed_time = time.time() - start_time tokens_sent = len(prompt.split()) tokens_received = len(response_text.split()) @@ -922,6 +1051,28 @@ def warm_prompt_cache(instruction: str, cache_key: str, model_id: str = "sonnet_ response_body = json.loads(response.get("body").read()) logging.info(f"Successfully warmed cache for {cache_key}") + # Track usage for cache warming call (use cache_key as filename identifier) + ( + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) = usage_tracking.extract_usage_from_response(response_body) + if ( + input_tokens > 0 + or output_tokens > 0 + or cache_creation_tokens > 0 + or cache_read_tokens > 0 + ): + usage_tracking.track_usage( + f"cache_warming_{cache_key}", + resolved_model_id, + input_tokens, + output_tokens, + cache_creation_tokens, + cache_read_tokens, + ) + # Log cache usage info if available if "usage" in response_body: usage = response_body["usage"] diff --git a/fieldExtraction/src/utils/usage_tracking.py b/fieldExtraction/src/utils/usage_tracking.py new file mode 100644 index 0000000..273c26b --- /dev/null +++ b/fieldExtraction/src/utils/usage_tracking.py @@ -0,0 +1,476 @@ +"""Usage and cost tracking for LLM API calls. + +This module provides thread-safe tracking of token usage and costs per file and model. +Phase 1: Data collection only. +""" + +import logging +import os +import threading +from collections import defaultdict +from datetime import datetime + +import pandas as pd + +from src import config + + +# Cost constants per 1k tokens (input, output, cache_read) +# Cache creation tokens are charged at the same rate as input tokens +# Cache read tokens are charged at ~10% of input token rate +# Update these values as pricing changes +COST_PER_1K_TOKENS = { + # Claude 2 + "anthropic.claude-instant-v1": (0.008, 0.024, 0.0008), + # Claude 3 Haiku + "anthropic.claude-3-haiku-20240307-v1:0": (0.00025, 0.00125, 0.000025), + # Claude 3.5 Sonnet + "anthropic.claude-3-5-sonnet-20240620-v1:0": (0.003, 0.015, 0.0003), + "anthropic.claude-3-5-sonnet-20241022-v2:0": (0.003, 0.015, 0.0003), + # Claude 3.7 Sonnet + "us.anthropic.claude-3-7-sonnet-20250219-v1:0": (0.003, 0.015, 0.0003), + # Claude 4 Sonnet + "us.anthropic.claude-sonnet-4-20250514-v1:0": (0.003, 0.015, 0.0003), + # Claude 4.5 Sonnet + "us.anthropic.claude-sonnet-4-5-20250929-v1:0": (0.003, 0.015, 0.0003), + # Default fallback (Claude 3.5 Sonnet pricing) + "_default": (0.003, 0.015, 0.0003), +} + +# Thread-safe lock for usage data updates +_usage_lock = threading.Lock() + +# Per-file, per-model usage data +# Structure: USAGE_DATA[filename][model_id] = { +# "input_tokens": int, +# "output_tokens": int, +# "cache_creation_tokens": int, +# "cache_read_tokens": int, +# "total_tokens": int, +# "total_cost": float, +# "request_count": int +# } +USAGE_DATA = defaultdict(lambda: defaultdict(lambda: { + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "total_tokens": 0, + "total_cost": 0.0, + "request_count": 0, +})) + +# Global usage totals +# Note: total_input_tokens = fresh_input_tokens + cache_tokens +# where fresh_input_tokens = input_tokens (non-cached) +# and cache_tokens = cache_creation_tokens + cache_read_tokens +GLOBAL_USAGE = { + "total_input_tokens": 0, # input_tokens + cache_creation_tokens + cache_read_tokens + "total_output_tokens": 0, # output tokens (always fresh, no cache) + "total_tokens": 0, + "total_cost": 0.0, + "total_requests": 0, + "total_fresh_input_tokens": 0, # input_tokens only (non-cached) + "total_cache_tokens": 0, # cache_creation_tokens + cache_read_tokens +} + + +def get_model_name(model_id: str) -> str: + """Convert model ID to human-readable model name. + + Args: + model_id: Full model ID (e.g., "anthropic.claude-3-5-sonnet-20240620-v1:0") + + Returns: + Human-readable model name (e.g., "Claude 3.5 Sonnet") + """ + model_id_lower = model_id.lower() + if "claude-3-haiku" in model_id_lower or "claude-3.5-haiku" in model_id_lower: + return "Claude 3 Haiku" + elif "claude-3-5-sonnet" in model_id_lower or "claude-3.5-sonnet" in model_id_lower: + return "Claude 3.5 Sonnet" + elif "claude-3-7-sonnet" in model_id_lower or "claude-3.7-sonnet" in model_id_lower: + return "Claude 3.7 Sonnet" + elif "claude-sonnet-4" in model_id_lower or "claude-4-sonnet" in model_id_lower: + if "4-5" in model_id_lower or "4.5" in model_id_lower: + return "Claude 4.5 Sonnet" + else: + return "Claude 4 Sonnet" + elif "claude-instant" in model_id_lower or "claude-2" in model_id_lower: + return "Claude 2" + else: + # Fallback: extract version from model ID + return f"Claude {model_id.split('.')[-1].split('-')[0] if '.' in model_id else 'Unknown'}" + + +def get_cost_per_token(model_id: str) -> tuple[float, float, float]: + """Get input, output, and cache read cost per token for a given model. + + Args: + model_id: Resolved model ID + + Returns: + Tuple of (input_cost_per_token, output_cost_per_token, cache_read_cost_per_token) + """ + # Resolve model alias to actual model ID + resolved_model_id = config.resolve_model_id(model_id) + + # Look up cost from centralized constants + cost_per_1k = COST_PER_1K_TOKENS.get(resolved_model_id, COST_PER_1K_TOKENS["_default"]) + + if cost_per_1k == COST_PER_1K_TOKENS["_default"] and resolved_model_id not in COST_PER_1K_TOKENS: + logging.warning( + f"Unknown model {resolved_model_id} (original: {model_id}), " + f"using default pricing" + ) + + input_cost_per_1k, output_cost_per_1k, cache_read_cost_per_1k = cost_per_1k + + # Convert per-1k to per-token + input_cost_per_token = input_cost_per_1k / 1000.0 + output_cost_per_token = output_cost_per_1k / 1000.0 + cache_read_cost_per_token = cache_read_cost_per_1k / 1000.0 + + return input_cost_per_token, output_cost_per_token, cache_read_cost_per_token + + +def track_usage( + filename: str, + model_id: str, + input_tokens: int, + output_tokens: int, + cache_creation_tokens: int = 0, + cache_read_tokens: int = 0, +) -> None: + """Track token usage and calculate cost for a single LLM API call. + + This function is thread-safe and updates both per-file/per-model data + and global totals. Includes cache token tracking with proper cost calculation. + + Args: + filename: Name of the file being processed + model_id: Full model ID used for the API call + input_tokens: Regular input tokens (non-cached) from API response + output_tokens: Output tokens from API response + cache_creation_tokens: Tokens used to create cache (charged at input rate) + cache_read_tokens: Tokens read from cache (charged at cache read rate) + """ + if ( + not filename + or input_tokens < 0 + or output_tokens < 0 + or cache_creation_tokens < 0 + or cache_read_tokens < 0 + ): + logging.warning( + f"Invalid usage data: filename={filename}, " + f"input_tokens={input_tokens}, output_tokens={output_tokens}, " + f"cache_creation={cache_creation_tokens}, cache_read={cache_read_tokens}" + ) + return + + # Total tokens includes all types + total_tokens = input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens + + # Get cost per token (cache creation uses input rate) + input_cost_per_token, output_cost_per_token, cache_read_cost_per_token = get_cost_per_token( + model_id + ) + + # Calculate total cost: + # - Regular input tokens: input rate + # - Cache creation tokens: input rate (same as regular input) + # - Cache read tokens: cache read rate (much lower) + # - Output tokens: output rate + total_cost = ( + (input_tokens * input_cost_per_token) + + (cache_creation_tokens * input_cost_per_token) + + (cache_read_tokens * cache_read_cost_per_token) + + (output_tokens * output_cost_per_token) + ) + + # Thread-safe update + with _usage_lock: + # Update per-file, per-model data + file_data = USAGE_DATA[filename][model_id] + file_data["input_tokens"] += input_tokens + file_data["output_tokens"] += output_tokens + file_data["cache_creation_tokens"] += cache_creation_tokens + file_data["cache_read_tokens"] += cache_read_tokens + file_data["total_tokens"] += total_tokens + file_data["total_cost"] += total_cost + file_data["request_count"] += 1 + + # Update global totals (include all token types) + cache_tokens = cache_creation_tokens + cache_read_tokens + GLOBAL_USAGE["total_fresh_input_tokens"] += input_tokens + GLOBAL_USAGE["total_cache_tokens"] += cache_tokens + GLOBAL_USAGE["total_input_tokens"] += input_tokens + cache_tokens + GLOBAL_USAGE["total_output_tokens"] += output_tokens + GLOBAL_USAGE["total_tokens"] += total_tokens + GLOBAL_USAGE["total_cost"] += total_cost + GLOBAL_USAGE["total_requests"] += 1 + + +def extract_usage_from_response(response_body: dict) -> tuple[int, int, int, int]: + """Extract input, output, and cache tokens from AWS Bedrock API response. + + Args: + response_body: Parsed JSON response from Bedrock API + + Returns: + Tuple of (input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens). + Returns (0, 0, 0, 0) if usage data not found. + """ + if response_body is None or "usage" not in response_body: + logging.debug("No usage data in API response") + return 0, 0, 0, 0 + + usage = response_body["usage"] + # Regular input tokens (non-cached) + input_tokens = usage.get("input_tokens", 0) + # Output tokens + output_tokens = usage.get("output_tokens", 0) + # Cache creation tokens (charged at input rate when cache is first created) + cache_creation_tokens = usage.get("cache_creation_input_tokens", 0) + # Cache read tokens (charged at lower cache read rate) + cache_read_tokens = usage.get("cache_read_input_tokens", 0) + + return ( + int(input_tokens), + int(output_tokens), + int(cache_creation_tokens), + int(cache_read_tokens), + ) + + +def get_usage_data() -> dict: + """Get current usage data snapshot. + + Returns: + Dictionary containing USAGE_DATA and GLOBAL_USAGE + """ + with _usage_lock: + # Deep copy the nested defaultdict structure to regular dicts + per_file_data = { + filename: {model_id: dict(usage_info) for model_id, usage_info in models.items()} + for filename, models in USAGE_DATA.items() + } + return { + "per_file_per_model": per_file_data, + "global": dict(GLOBAL_USAGE), + } + + +def reset_usage_data() -> None: + """Reset all usage tracking data. Useful for testing.""" + global USAGE_DATA, GLOBAL_USAGE + with _usage_lock: + USAGE_DATA = defaultdict(lambda: defaultdict(lambda: { + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "total_tokens": 0, + "total_cost": 0.0, + "request_count": 0, + })) + GLOBAL_USAGE = { + "total_input_tokens": 0, + "total_output_tokens": 0, + "total_tokens": 0, + "total_cost": 0.0, + "total_requests": 0, + "total_fresh_input_tokens": 0, + "total_cache_tokens": 0, + } + + +def get_usage_dataframes(run_id: str, timestamp: str) -> tuple[pd.DataFrame, pd.DataFrame]: + """Get usage and cost DataFrames ready for export. + + Returns the formatted DataFrames that can be passed to S3 upload functions + or saved locally. This is the interface for other developers handling S3 uploads. + + Args: + run_id: Unique identifier for this run (typically BATCH_ID) + timestamp: Timestamp string for this run + + Returns: + Tuple of (per_file_per_model_df, batch_summary_df) + - per_file_per_model_df: DataFrame with per-file, per-model usage data + - batch_summary_df: DataFrame with single row containing batch summary + + Example: + ```python + per_file_df, summary_df = usage_tracking.get_usage_dataframes( + config.BATCH_ID, run_timestamp + ) + # Pass to S3 upload function + write_s3(per_file_df, "USAGE", run_timestamp, "usage") + write_s3(summary_df, "USAGE-SUMMARY", run_timestamp, "usage") + ``` + """ + per_file_df = _generate_per_file_per_model_df(run_id, timestamp) + summary_df = _generate_batch_summary_df(run_id, timestamp) + return per_file_df, summary_df + + +def _generate_per_file_per_model_df(run_id: str, timestamp: str) -> pd.DataFrame: + """Generate per-file, per-model usage DataFrame. + + Args: + run_id: Unique identifier for this run (typically BATCH_ID) + timestamp: Timestamp string for this run + + Returns: + DataFrame with columns: run_id, timestamp, file_name, model_name, + input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, + total_tokens, cost_per_token, total_cost + """ + rows = [] + usage_snapshot = get_usage_data() + + for filename, models_dict in usage_snapshot["per_file_per_model"].items(): + for model_id, usage_info in models_dict.items(): + if usage_info["request_count"] == 0: + continue # Skip models with no usage + + model_name = get_model_name(model_id) + input_tokens = usage_info["input_tokens"] + output_tokens = usage_info["output_tokens"] + cache_creation_tokens = usage_info["cache_creation_tokens"] + cache_read_tokens = usage_info["cache_read_tokens"] + total_tokens = usage_info["total_tokens"] + total_cost = usage_info["total_cost"] + request_count = usage_info["request_count"] + + # Calculate average cost per token (for this file/model combination) + cost_per_token = total_cost / total_tokens if total_tokens > 0 else 0.0 + + rows.append({ + "run_id": run_id, + "timestamp": timestamp, + "file_name": filename, + "model_name": model_name, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_creation_tokens": cache_creation_tokens, + "cache_read_tokens": cache_read_tokens, + "total_tokens": total_tokens, + "cost_per_token": round(cost_per_token, 8), + "total_cost": round(total_cost, 6), + }) + + return pd.DataFrame(rows) + + +def _generate_batch_summary_df(run_id: str, timestamp: str) -> pd.DataFrame: + """Generate batch-level summary DataFrame. + + Args: + run_id: Unique identifier for this run (typically BATCH_ID) + timestamp: Timestamp string for this run + + Returns: + DataFrame with single row containing batch summary statistics + """ + usage_snapshot = get_usage_data() + global_data = usage_snapshot["global"] + per_file_data = usage_snapshot["per_file_per_model"] + + # Calculate averages + total_files = len(per_file_data) + total_requests = global_data["total_requests"] + + # Extract token counts for clearer calculation + fresh_input_tokens = global_data["total_fresh_input_tokens"] + cache_tokens = global_data["total_cache_tokens"] + total_input_tokens = global_data["total_input_tokens"] + total_output_tokens = global_data["total_output_tokens"] + total_tokens = global_data["total_tokens"] + total_cost = global_data["total_cost"] + + # Calculate averages + average_fresh_input_tokens = ( + fresh_input_tokens / total_files if total_files > 0 else 0 + ) + average_cache_tokens = ( + cache_tokens / total_files if total_files > 0 else 0 + ) + average_input_tokens = ( + total_input_tokens / total_files if total_files > 0 else 0 + ) + average_output_tokens = ( + total_output_tokens / total_files if total_files > 0 else 0 + ) + average_total_tokens = ( + total_tokens / total_files if total_files > 0 else 0 + ) + average_cost = ( + total_cost / total_files if total_files > 0 else 0 + ) + + # Determine region mode + region_mode = "cross-region" if config.ENABLE_RUNTIME_ROTATION else "single-region" + + summary_row = { + "run_id": run_id, + "timestamp": timestamp, + # Averages + "average_fresh_input_tokens": round(average_fresh_input_tokens, 2), + "average_cache_tokens": round(average_cache_tokens, 2), + "average_input_tokens": round(average_input_tokens, 2), + "average_output_tokens": round(average_output_tokens, 2), + "average_total_tokens": round(average_total_tokens, 2), + "average_cost": round(average_cost, 6), + # Totals (split for tractability) + "total_fresh_input_tokens": fresh_input_tokens, + "total_cache_tokens": cache_tokens, + "total_input_tokens": total_input_tokens, # = fresh_input_tokens + cache_tokens + "total_output_tokens": total_output_tokens, # Always fresh (no cache for output) + "total_tokens": total_tokens, + "total_cost": round(total_cost, 6), + "region_mode": region_mode, + } + + return pd.DataFrame([summary_row]) + + +def export_usage_cost_csvs_local(run_id: str, run_timestamp: str) -> None: + """Export usage and cost data to CSV files in local directory. + + Saves two CSV files: + 1. Per-file, per-model usage data + 2. Batch-level summary + + Files are saved to the same directory as final results: + {CONSOLIDATED_OUTPUT_DIRECTORY}/{run_timestamp}/ + + Args: + run_id: Unique identifier for this run (typically BATCH_ID) + run_timestamp: Timestamp string matching the run (e.g., "run_20240101_12-00_test_batch") + """ + # Get DataFrames using the public interface + per_file_df, summary_df = get_usage_dataframes(run_id, run_timestamp) + + # Determine output directory (same as final results) + output_dir = os.path.join(config.CONSOLIDATED_OUTPUT_DIRECTORY, run_timestamp) + os.makedirs(output_dir, exist_ok=True) + + # Save per-file, per-model CSV + if not per_file_df.empty: + per_file_path = os.path.join(output_dir, f"{run_id}-USAGE.csv") + per_file_df.to_csv(per_file_path, index=False, quoting=1) + logging.info(f"Saved usage data to: {per_file_path}") + else: + logging.warning("No usage data to export (per-file CSV empty)") + + # Save batch summary CSV + if not summary_df.empty: + summary_path = os.path.join(output_dir, f"{run_id}-USAGE-SUMMARY.csv") + summary_df.to_csv(summary_path, index=False, quoting=1) + logging.info(f"Saved usage summary to: {summary_path}") + else: + logging.warning("No summary data to export (summary CSV empty)") + diff --git a/fieldExtraction/tests/test_io_utils.py b/fieldExtraction/tests/test_io_utils.py index c01ac5f..7d77e72 100644 --- a/fieldExtraction/tests/test_io_utils.py +++ b/fieldExtraction/tests/test_io_utils.py @@ -1,14 +1,44 @@ from io import StringIO +import os import pandas as pd import pytest +from botocore.exceptions import ClientError from src.utils.io_utils import (filter_already_processed, list_s3_files, read_input, read_input_csv, read_local, read_s3, read_s3_csv, read_xlsb, - remove_txt_extension) + remove_txt_extension, write_s3) class TestIOUtils: + @pytest.fixture(autouse=True) + def cleanup_config_patches(self, mocker): + """Ensure config patches are cleaned up after each test to prevent leakage.""" + # Store original config values before test + import src.config as config_module + original_values = {} + config_attrs = [ + "S3_CLIENT", + "BATCH_ID", + "S3_OUTPUT_BUCKET", + "READ_MODE", + "DF_READ_MODE", + "AC_DF", + "B_DF", + "WRITE_TO_S3", + "OUTPUT_DIRECTORY", + "AC_RESULTS_NAME", + "B_RESULTS_NAME", + ] + for attr in config_attrs: + if hasattr(config_module, attr): + original_values[attr] = getattr(config_module, attr) + + yield # Run the test + + # Restore original values after test + for attr, value in original_values.items(): + setattr(config_module, attr, value) # Fixtures for reusable test data @pytest.fixture def mock_csv_data(self): @@ -296,9 +326,83 @@ class TestIOUtils: # Verify that the mock was called correctly mock_open_workbook.assert_called_once_with("dummy_path.xlsb") mock_workbook.get_sheet.assert_called_once() - # Assert that the result matches the expected DataFrame - pd.testing.assert_frame_equal(result, expected_df) - # Verify that the mock was called correctly - mock_open_workbook.assert_called_once_with("dummy_path.xlsb") - mock_workbook.get_sheet.assert_called_once() + # Tests for write_s3 + def test_write_s3_final(self, mocker): + """Test write_s3 with final output_type as used in main.py.""" + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "test_batch_123") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket") + mocker.patch("src.utils.io_utils.logging") + + df = pd.DataFrame({"col1": [1, 2], "col2": [3, 4]}) + write_s3(df, "", "20240101120000", "final") + + mock_s3_client.put_object.assert_called_once() + call_kwargs = mock_s3_client.put_object.call_args[1] + assert call_kwargs["Bucket"] == "test-bucket" + assert call_kwargs["Key"] == "test_batch_123/20240101120000/test_batch_123-RESULTS.csv" + assert isinstance(call_kwargs["Body"], bytes) + assert b"col1,col2" in call_kwargs["Body"] + + def test_write_s3_error(self, mocker): + """Test write_s3 with error output_type as used in main.py.""" + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "test_batch_123") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket") + mocker.patch("src.utils.io_utils.logging") + + df = pd.DataFrame({"error": ["test error"]}) + write_s3(df, "", "20240101120000", "error") + + call_kwargs = mock_s3_client.put_object.call_args[1] + assert call_kwargs["Key"] == "test_batch_123/20240101120000/test_batch_123-ERRORS.csv" + + def test_write_s3_usage(self, mocker): + """Test write_s3 with usage output_type as used in main.py.""" + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "test_batch_123") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket") + mocker.patch("src.utils.io_utils.logging") + + df = pd.DataFrame({"filename": ["file1.txt"], "tokens": [100]}) + write_s3(df, "", "20240101120000", "usage") + + call_kwargs = mock_s3_client.put_object.call_args[1] + assert call_kwargs["Key"] == "test_batch_123/20240101120000/tracking/test_batch_123-USAGE.csv" + + def test_write_s3_usage_summary(self, mocker): + """Test write_s3 with usage_summary output_type as used in main.py.""" + mock_s3_client = mocker.MagicMock() + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "test_batch_123") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket") + mocker.patch("src.utils.io_utils.logging") + + df = pd.DataFrame({"total_tokens": [1000], "total_cost": [0.50]}) + write_s3(df, "", "20240101120000", "usage_summary") + + call_kwargs = mock_s3_client.put_object.call_args[1] + assert call_kwargs["Key"] == "test_batch_123/20240101120000/tracking/test_batch_123-USAGE-SUMMARY.csv" + + def test_write_s3_client_error(self, mocker): + """Test write_s3 error handling when S3 upload fails.""" + mock_s3_client = mocker.MagicMock() + mock_s3_client.put_object.side_effect = ClientError( + {"Error": {"Code": "AccessDenied", "Message": "Access Denied"}}, + "PutObject", + ) + mocker.patch("src.config.S3_CLIENT", mock_s3_client) + mocker.patch("src.config.BATCH_ID", "test_batch_123") + mocker.patch("src.config.S3_OUTPUT_BUCKET", "test-bucket") + mock_logger = mocker.patch("src.utils.io_utils.logging") + + df = pd.DataFrame({"col1": [1]}) + result = write_s3(df, "", "20240101120000", "final") + + assert result is None + mock_logger.error.assert_called_once() + assert "Error uploading to S3" in str(mock_logger.error.call_args) diff --git a/fieldExtraction/tests/test_llm_utils.py b/fieldExtraction/tests/test_llm_utils.py index 4aac239..465ac14 100644 --- a/fieldExtraction/tests/test_llm_utils.py +++ b/fieldExtraction/tests/test_llm_utils.py @@ -50,12 +50,16 @@ class TestLLMUtils(unittest.TestCase): @patch("builtins.open", new_callable=mock_open) def test_invoke_claude_local_mode(self, mock_open_file): - config.RUN_MODE = "local" - model_id = config.MODEL_ID_CLAUDE2 + original_run_mode = config.RUN_MODE + try: + config.RUN_MODE = "local" + model_id = config.MODEL_ID_CLAUDE2 - with patch("src.utils.llm_utils.local_claude_2", return_value=self.response): - response = llm_utils.invoke_claude(self.prompt, model_id, self.filename) - self.assertEqual(response, self.response) + with patch("src.utils.llm_utils.local_claude_2", return_value=self.response): + response = llm_utils.invoke_claude(self.prompt, model_id, self.filename) + self.assertEqual(response, self.response) + finally: + config.RUN_MODE = original_run_mode if __name__ == "__main__": diff --git a/fieldExtraction/tests/test_usage_tracking.py b/fieldExtraction/tests/test_usage_tracking.py new file mode 100644 index 0000000..041646a --- /dev/null +++ b/fieldExtraction/tests/test_usage_tracking.py @@ -0,0 +1,1600 @@ +"""Unit tests for usage_tracking module. + +Phase 1: Core Functions +- extract_usage_from_response() +- get_cost_per_token() +- track_usage() + +Phase 2: Data Access Functions +- get_usage_data() +- reset_usage_data() + +Phase 3: Export Functions +- get_usage_dataframes() +- _generate_per_file_per_model_df() +- _generate_batch_summary_df() +- export_usage_cost_csvs_local() + +Phase 4: Utility Functions +- get_model_name() +""" + +import logging +import threading +import time + +import pandas as pd +import pytest + +from src.utils import usage_tracking + + +class TestExtractUsageFromResponse: + """Tests for extract_usage_from_response() function.""" + + @pytest.mark.parametrize( + "response_body, expected", + [ + # Valid response with all token types + ( + { + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "cache_creation_input_tokens": 20, + "cache_read_input_tokens": 30, + } + }, + (100, 50, 20, 30), + ), + # Response with only input/output (no cache tokens) + ( + { + "usage": { + "input_tokens": 100, + "output_tokens": 50, + } + }, + (100, 50, 0, 0), + ), + # Response with cache tokens only + ( + { + "usage": { + "cache_creation_input_tokens": 20, + "cache_read_input_tokens": 30, + } + }, + (0, 0, 20, 30), + ), + # Response with missing optional fields + ( + { + "usage": { + "input_tokens": 100, + } + }, + (100, 0, 0, 0), + ), + # Response with zero values + ( + { + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 0, + } + }, + (0, 0, 0, 0), + ), + # Response with string numbers (should convert to int) + ( + { + "usage": { + "input_tokens": "100", + "output_tokens": "50", + "cache_creation_input_tokens": "20", + "cache_read_input_tokens": "30", + } + }, + (100, 50, 20, 30), + ), + ], + ) + def test_extract_usage_from_response_valid(self, response_body, expected): + """Test extraction with valid response bodies.""" + result = usage_tracking.extract_usage_from_response(response_body) + assert result == expected + # Verify all values are integers + assert all(isinstance(x, int) for x in result) + + @pytest.mark.parametrize( + "response_body, expected", + [ + # Missing "usage" key + ({}, (0, 0, 0, 0)), + # Empty "usage" dict + ({"usage": {}}, (0, 0, 0, 0)), + # None response + (None, (0, 0, 0, 0)), + ], + ) + def test_extract_usage_from_response_missing_usage(self, response_body, expected, caplog): + """Test extraction when usage data is missing.""" + with caplog.at_level(logging.DEBUG): + result = usage_tracking.extract_usage_from_response(response_body) + assert result == expected + if response_body is None or "usage" not in (response_body or {}): + assert "No usage data in API response" in caplog.text + + def test_extract_usage_from_response_nested_structure(self): + """Test that function handles nested structure correctly.""" + response_body = { + "other_data": "value", + "usage": { + "input_tokens": 100, + "output_tokens": 50, + "cache_creation_input_tokens": 20, + "cache_read_input_tokens": 30, + }, + "more_data": {"nested": "value"}, + } + result = usage_tracking.extract_usage_from_response(response_body) + assert result == (100, 50, 20, 30) + + +class TestGetCostPerToken: + """Tests for get_cost_per_token() function.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + @pytest.mark.parametrize( + "model_id, expected_input, expected_output, expected_cache_read", + [ + # Claude 2 + ( + "anthropic.claude-instant-v1", + 0.008 / 1000.0, # $0.008 per 1k tokens + 0.024 / 1000.0, # $0.024 per 1k tokens + 0.0008 / 1000.0, # $0.0008 per 1k tokens + ), + # Claude 3 Haiku + ( + "anthropic.claude-3-haiku-20240307-v1:0", + 0.00025 / 1000.0, + 0.00125 / 1000.0, + 0.000025 / 1000.0, + ), + # Claude 3.5 Sonnet + ( + "anthropic.claude-3-5-sonnet-20240620-v1:0", + 0.003 / 1000.0, + 0.015 / 1000.0, + 0.0003 / 1000.0, + ), + ( + "anthropic.claude-3-5-sonnet-20241022-v2:0", + 0.003 / 1000.0, + 0.015 / 1000.0, + 0.0003 / 1000.0, + ), + # Claude 3.7 Sonnet + ( + "us.anthropic.claude-3-7-sonnet-20250219-v1:0", + 0.003 / 1000.0, + 0.015 / 1000.0, + 0.0003 / 1000.0, + ), + # Claude 4 Sonnet + ( + "us.anthropic.claude-sonnet-4-20250514-v1:0", + 0.003 / 1000.0, + 0.015 / 1000.0, + 0.0003 / 1000.0, + ), + # Claude 4.5 Sonnet + ( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + 0.003 / 1000.0, + 0.015 / 1000.0, + 0.0003 / 1000.0, + ), + ], + ) + def test_get_cost_per_token_valid_models( + self, model_id, expected_input, expected_output, expected_cache_read, mocker + ): + """Test cost calculation for valid model IDs.""" + # Mock resolve_model_id to return the model_id as-is + mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id) + + input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id) + + assert abs(input_cost - expected_input) < 1e-10 + assert abs(output_cost - expected_output) < 1e-10 + assert abs(cache_read_cost - expected_cache_read) < 1e-10 + + def test_get_cost_per_token_unknown_model(self, mocker, caplog): + """Test cost calculation for unknown model (should use default).""" + unknown_model = "unknown.model.id" + # Mock resolve_model_id to return the unknown model + mocker.patch( + "src.utils.usage_tracking.config.resolve_model_id", return_value=unknown_model + ) + + with caplog.at_level(logging.WARNING): + input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token( + unknown_model + ) + + # Should use default pricing (Claude 3.5 Sonnet) + expected_input = 0.003 / 1000.0 + expected_output = 0.015 / 1000.0 + expected_cache_read = 0.0003 / 1000.0 + + assert abs(input_cost - expected_input) < 1e-10 + assert abs(output_cost - expected_output) < 1e-10 + assert abs(cache_read_cost - expected_cache_read) < 1e-10 + assert "Unknown model" in caplog.text + assert "using default pricing" in caplog.text + + def test_get_cost_per_token_model_alias(self, mocker): + """Test that model aliases are resolved correctly.""" + alias = "sonnet_37" + resolved_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0" + # Mock resolve_model_id to resolve the alias + mocker.patch( + "src.utils.usage_tracking.config.resolve_model_id", return_value=resolved_id + ) + + input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(alias) + + # Should use Claude 3.7 Sonnet pricing + expected_input = 0.003 / 1000.0 + expected_output = 0.015 / 1000.0 + expected_cache_read = 0.0003 / 1000.0 + + assert abs(input_cost - expected_input) < 1e-10 + assert abs(output_cost - expected_output) < 1e-10 + assert abs(cache_read_cost - expected_cache_read) < 1e-10 + + def test_get_cost_per_token_conversion(self, mocker): + """Test that per-1k costs are correctly converted to per-token.""" + model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0" + mocker.patch("src.utils.usage_tracking.config.resolve_model_id", return_value=model_id) + + input_cost, output_cost, cache_read_cost = usage_tracking.get_cost_per_token(model_id) + + # Verify conversion: per-1k / 1000 = per-token + assert input_cost == 0.003 / 1000.0 + assert output_cost == 0.015 / 1000.0 + assert cache_read_cost == 0.0003 / 1000.0 + + # Verify all costs are floats + assert isinstance(input_cost, float) + assert isinstance(output_cost, float) + assert isinstance(cache_read_cost, float) + + +class TestTrackUsage: + """Tests for track_usage() function.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + @pytest.mark.parametrize( + "filename, model_id, input_tokens, output_tokens, cache_creation, cache_read", + [ + # All token types present + ("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 20, 30), + # No cache tokens + ("file2.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 0, 0), + # Only cache tokens + ("file3.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 20, 30), + # Zero tokens + ("file4.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 0, 0, 0, 0), + ], + ) + def test_track_usage_single_call( + self, + filename, + model_id, + input_tokens, + output_tokens, + cache_creation, + cache_read, + mocker, + ): + """Test tracking a single API call.""" + # Mock get_cost_per_token to return known costs + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage( + filename, model_id, input_tokens, output_tokens, cache_creation, cache_read + ) + + # Verify per-file, per-model data + usage_data = usage_tracking.get_usage_data() + file_data = usage_data["per_file_per_model"][filename][model_id] + + assert file_data["input_tokens"] == input_tokens + assert file_data["output_tokens"] == output_tokens + assert file_data["cache_creation_tokens"] == cache_creation + assert file_data["cache_read_tokens"] == cache_read + assert file_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read + assert file_data["request_count"] == 1 + + # Verify cost calculation + expected_cost = ( + (input_tokens * input_cost) + + (cache_creation * input_cost) # Cache creation uses input rate + + (cache_read * cache_read_cost) + + (output_tokens * output_cost) + ) + assert abs(file_data["total_cost"] - expected_cost) < 1e-10 + + # Verify global totals + global_data = usage_data["global"] + assert global_data["total_fresh_input_tokens"] == input_tokens + assert global_data["total_cache_tokens"] == cache_creation + cache_read + assert global_data["total_input_tokens"] == input_tokens + cache_creation + cache_read + assert global_data["total_output_tokens"] == output_tokens + assert global_data["total_tokens"] == input_tokens + output_tokens + cache_creation + cache_read + assert abs(global_data["total_cost"] - expected_cost) < 1e-10 + assert global_data["total_requests"] == 1 + + def test_track_usage_multiple_calls_same_file_model(self, mocker): + """Test tracking multiple calls for same file/model (accumulation).""" + filename = "file1.txt" + model_id = "anthropic.claude-3-5-sonnet-20240620-v1:0" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # First call + usage_tracking.track_usage(filename, model_id, 100, 50, 10, 20) + # Second call + usage_tracking.track_usage(filename, model_id, 200, 100, 20, 30) + + usage_data = usage_tracking.get_usage_data() + file_data = usage_data["per_file_per_model"][filename][model_id] + + # Verify accumulation + assert file_data["input_tokens"] == 300 # 100 + 200 + assert file_data["output_tokens"] == 150 # 50 + 100 + assert file_data["cache_creation_tokens"] == 30 # 10 + 20 + assert file_data["cache_read_tokens"] == 50 # 20 + 30 + assert file_data["total_tokens"] == 530 # 300 + 150 + 30 + 50 + assert file_data["request_count"] == 2 + + # Verify global totals + global_data = usage_data["global"] + assert global_data["total_fresh_input_tokens"] == 300 + assert global_data["total_cache_tokens"] == 80 # 30 + 50 + assert global_data["total_input_tokens"] == 380 # 300 + 80 + assert global_data["total_output_tokens"] == 150 + assert global_data["total_tokens"] == 530 + assert global_data["total_requests"] == 2 + + def test_track_usage_multiple_files_models(self, mocker): + """Test tracking calls for different files/models (isolation).""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # File 1, Model 1 + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + # File 1, Model 2 + usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30) + # File 2, Model 1 + usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25) + + usage_data = usage_tracking.get_usage_data() + + # Verify file1, model1 + file1_model1 = usage_data["per_file_per_model"]["file1.txt"]["model1"] + assert file1_model1["input_tokens"] == 100 + assert file1_model1["request_count"] == 1 + + # Verify file1, model2 + file1_model2 = usage_data["per_file_per_model"]["file1.txt"]["model2"] + assert file1_model2["input_tokens"] == 200 + assert file1_model2["request_count"] == 1 + + # Verify file2, model1 + file2_model1 = usage_data["per_file_per_model"]["file2.txt"]["model1"] + assert file2_model1["input_tokens"] == 150 + assert file2_model1["request_count"] == 1 + + # Verify global totals (sum of all) + global_data = usage_data["global"] + assert global_data["total_fresh_input_tokens"] == 450 # 100 + 200 + 150 + assert global_data["total_cache_tokens"] == 120 # (10+20) + (20+30) + (15+25) = 30 + 50 + 40 + assert global_data["total_requests"] == 3 + + @pytest.mark.parametrize( + "filename, input_tokens, output_tokens, cache_creation, cache_read, should_reject", + [ + # Empty filename + ("", 100, 50, 0, 0, True), + (None, 100, 50, 0, 0, True), + # Negative input tokens + ("file.txt", -1, 50, 0, 0, True), + # Negative output tokens + ("file.txt", 100, -1, 0, 0, True), + # Negative cache creation tokens + ("file.txt", 100, 50, -1, 0, True), + # Negative cache read tokens + ("file.txt", 100, 50, 0, -1, True), + # Valid (zero values are OK) + ("file.txt", 0, 0, 0, 0, False), + ], + ) + def test_track_usage_input_validation( + self, + filename, + input_tokens, + output_tokens, + cache_creation, + cache_read, + should_reject, + mocker, + caplog, + ): + """Test input validation (should reject invalid inputs).""" + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(0.001, 0.002, 0.0001), + ) + + initial_data = usage_tracking.get_usage_data() + initial_requests = initial_data["global"]["total_requests"] + + with caplog.at_level(logging.WARNING): + usage_tracking.track_usage( + filename, "model1", input_tokens, output_tokens, cache_creation, cache_read + ) + + final_data = usage_tracking.get_usage_data() + final_requests = final_data["global"]["total_requests"] + + if should_reject: + # Should not have updated usage data + assert final_requests == initial_requests + assert "Invalid usage data" in caplog.text + else: + # Should have updated usage data + assert final_requests == initial_requests + 1 + + def test_track_usage_thread_safety(self, mocker): + """Test thread safety with concurrent calls.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + num_threads = 10 + calls_per_thread = 50 + filename = "file1.txt" + model_id = "model1" + + def track_multiple_calls(): + for _ in range(calls_per_thread): + usage_tracking.track_usage(filename, model_id, 10, 5, 2, 3) + time.sleep(0.001) # Small delay to increase chance of race condition + + threads = [] + for _ in range(num_threads): + thread = threading.Thread(target=track_multiple_calls) + threads.append(thread) + thread.start() + + # Wait for all threads to complete + for thread in threads: + thread.join() + + usage_data = usage_tracking.get_usage_data() + file_data = usage_data["per_file_per_model"][filename][model_id] + + # Verify all calls were tracked (no data loss) + expected_requests = num_threads * calls_per_thread + assert file_data["request_count"] == expected_requests + + # Verify token totals are correct + expected_input = expected_requests * 10 + expected_output = expected_requests * 5 + expected_cache_creation = expected_requests * 2 + expected_cache_read = expected_requests * 3 + + assert file_data["input_tokens"] == expected_input + assert file_data["output_tokens"] == expected_output + assert file_data["cache_creation_tokens"] == expected_cache_creation + assert file_data["cache_read_tokens"] == expected_cache_read + + # Verify global totals + global_data = usage_data["global"] + assert global_data["total_requests"] == expected_requests + assert global_data["total_fresh_input_tokens"] == expected_input + assert global_data["total_cache_tokens"] == expected_cache_creation + expected_cache_read + + def test_track_usage_cost_calculation(self, mocker): + """Test that cost calculation is correct for all token types.""" + filename = "file1.txt" + model_id = "model1" + input_cost = 0.003 / 1000.0 # $0.003 per 1k = $0.000003 per token + output_cost = 0.015 / 1000.0 # $0.015 per 1k = $0.000015 per token + cache_read_cost = 0.0003 / 1000.0 # $0.0003 per 1k = $0.0000003 per token + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # Track usage with specific token counts + input_tokens = 1000 + output_tokens = 500 + cache_creation_tokens = 200 # Charged at input rate + cache_read_tokens = 300 # Charged at cache read rate + + usage_tracking.track_usage( + filename, model_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens + ) + + usage_data = usage_tracking.get_usage_data() + file_data = usage_data["per_file_per_model"][filename][model_id] + + # Calculate expected cost manually + expected_cost = ( + (input_tokens * input_cost) + + (cache_creation_tokens * input_cost) # Same as input rate + + (cache_read_tokens * cache_read_cost) # Lower rate + + (output_tokens * output_cost) + ) + + assert abs(file_data["total_cost"] - expected_cost) < 1e-10 + + # Verify cache creation uses input rate, not cache read rate + cache_creation_cost = cache_creation_tokens * input_cost + cache_read_cost_calc = cache_read_tokens * cache_read_cost + assert cache_creation_cost > cache_read_cost_calc # Input rate > cache read rate + + +class TestGetUsageData: + """Tests for get_usage_data() function - Phase 2.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_get_usage_data_empty_state(self): + """Test getting usage data when no usage has been tracked.""" + usage_data = usage_tracking.get_usage_data() + + assert "per_file_per_model" in usage_data + assert "global" in usage_data + assert usage_data["per_file_per_model"] == {} + assert usage_data["global"]["total_tokens"] == 0 + assert usage_data["global"]["total_cost"] == 0.0 + assert usage_data["global"]["total_requests"] == 0 + + def test_get_usage_data_single_file_single_model(self, mocker): + """Test getting usage data with single file and model.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_data = usage_tracking.get_usage_data() + + assert "per_file_per_model" in usage_data + assert "global" in usage_data + assert "file1.txt" in usage_data["per_file_per_model"] + assert "model1" in usage_data["per_file_per_model"]["file1.txt"] + assert usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100 + assert usage_data["global"]["total_fresh_input_tokens"] == 100 + + def test_get_usage_data_multiple_files_single_model(self, mocker): + """Test getting usage data with multiple files and single model.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0) + usage_data = usage_tracking.get_usage_data() + + assert len(usage_data["per_file_per_model"]) == 2 + assert "file1.txt" in usage_data["per_file_per_model"] + assert "file2.txt" in usage_data["per_file_per_model"] + assert usage_data["global"]["total_fresh_input_tokens"] == 300 + + def test_get_usage_data_single_file_multiple_models(self, mocker): + """Test getting usage data with single file and multiple models.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0) + usage_data = usage_tracking.get_usage_data() + + assert len(usage_data["per_file_per_model"]) == 1 + assert "file1.txt" in usage_data["per_file_per_model"] + assert len(usage_data["per_file_per_model"]["file1.txt"]) == 2 + assert "model1" in usage_data["per_file_per_model"]["file1.txt"] + assert "model2" in usage_data["per_file_per_model"]["file1.txt"] + + def test_get_usage_data_multiple_files_multiple_models(self, mocker): + """Test getting usage data with multiple files and models.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30) + usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25) + usage_data = usage_tracking.get_usage_data() + + assert len(usage_data["per_file_per_model"]) == 2 + assert usage_data["global"]["total_fresh_input_tokens"] == 450 + assert usage_data["global"]["total_cache_tokens"] == 120 + assert usage_data["global"]["total_requests"] == 3 + + def test_get_usage_data_deep_copy(self, mocker): + """Test that modifying returned dict doesn't affect internal state.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_data = usage_tracking.get_usage_data() + + # Modify the returned data + usage_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] = 999 + usage_data["global"]["total_tokens"] = 999 + + # Get fresh data and verify original values are unchanged + fresh_data = usage_tracking.get_usage_data() + assert fresh_data["per_file_per_model"]["file1.txt"]["model1"]["input_tokens"] == 100 + assert fresh_data["global"]["total_tokens"] == 150 # 100 + 50 + + def test_get_usage_data_thread_safety(self, mocker): + """Test thread safety when getting data while other thread is updating.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + results = [] + errors = [] + + def track_and_get(): + try: + for i in range(10): + usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0) + data = usage_tracking.get_usage_data() + results.append(data["global"]["total_requests"]) + except Exception as e: + errors.append(e) + + threads = [] + for _ in range(5): + thread = threading.Thread(target=track_and_get) + threads.append(thread) + thread.start() + + for thread in threads: + thread.join() + + # Verify no errors occurred + assert len(errors) == 0 + + # Verify final state is correct + final_data = usage_tracking.get_usage_data() + assert final_data["global"]["total_requests"] == 50 # 5 threads * 10 calls + + +class TestResetUsageData: + """Tests for reset_usage_data() function - Phase 2.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_reset_usage_data_after_tracking(self, mocker): + """Test reset after tracking some usage.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # Track some usage + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_tracking.track_usage("file2.txt", "model2", 200, 100, 20, 30) + + # Verify data exists + usage_data = usage_tracking.get_usage_data() + assert usage_data["global"]["total_requests"] == 2 + assert len(usage_data["per_file_per_model"]) == 2 + + # Reset + usage_tracking.reset_usage_data() + + # Verify all data is cleared + usage_data = usage_tracking.get_usage_data() + assert usage_data["global"]["total_tokens"] == 0 + assert usage_data["global"]["total_cost"] == 0.0 + assert usage_data["global"]["total_requests"] == 0 + assert usage_data["global"]["total_fresh_input_tokens"] == 0 + assert usage_data["global"]["total_cache_tokens"] == 0 + assert usage_data["per_file_per_model"] == {} + + def test_reset_usage_data_when_empty(self): + """Test reset when already empty (should not error).""" + # Reset when already empty + usage_tracking.reset_usage_data() + + # Verify still empty + usage_data = usage_tracking.get_usage_data() + assert usage_data["global"]["total_requests"] == 0 + assert usage_data["per_file_per_model"] == {} + + # Reset again (should not error) + usage_tracking.reset_usage_data() + usage_data = usage_tracking.get_usage_data() + assert usage_data["global"]["total_requests"] == 0 + + def test_reset_usage_data_clears_both_structures(self, mocker): + """Test that reset clears both USAGE_DATA and GLOBAL_USAGE.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_data = usage_tracking.get_usage_data() + + # Verify both structures have data + assert len(usage_data["per_file_per_model"]) > 0 + assert usage_data["global"]["total_requests"] > 0 + + # Reset + usage_tracking.reset_usage_data() + usage_data = usage_tracking.get_usage_data() + + # Verify both structures are cleared + assert usage_data["per_file_per_model"] == {} + assert usage_data["global"]["total_requests"] == 0 + assert usage_data["global"]["total_tokens"] == 0 + assert usage_data["global"]["total_cost"] == 0.0 + assert usage_data["global"]["total_fresh_input_tokens"] == 0 + assert usage_data["global"]["total_cache_tokens"] == 0 + + def test_reset_usage_data_thread_safety(self, mocker): + """Test thread safety of reset operation.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + errors = [] + + def track_and_reset(): + try: + for i in range(10): + usage_tracking.track_usage(f"file{i}.txt", "model1", 10, 5, 0, 0) + usage_tracking.reset_usage_data() + except Exception as e: + errors.append(e) + + threads = [] + for _ in range(5): + thread = threading.Thread(target=track_and_reset) + threads.append(thread) + thread.start() + + for thread in threads: + thread.join() + + # Verify no errors occurred + assert len(errors) == 0 + + # Verify final state is reset + usage_data = usage_tracking.get_usage_data() + assert usage_data["global"]["total_requests"] == 0 + assert usage_data["per_file_per_model"] == {} + + +class TestGetUsageDataframes: + """Tests for get_usage_dataframes() function - Phase 3.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_get_usage_dataframes_empty(self): + """Test getting DataFrames when no usage data exists.""" + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + + per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp) + + assert isinstance(per_file_df, pd.DataFrame) + assert isinstance(summary_df, pd.DataFrame) + assert per_file_df.empty + assert not summary_df.empty # Summary should have one row with zeros + + def test_get_usage_dataframes_single_file_single_model(self, mocker): + """Test getting DataFrames with single file and model.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp) + + assert isinstance(per_file_df, pd.DataFrame) + assert isinstance(summary_df, pd.DataFrame) + assert len(per_file_df) == 1 + assert len(summary_df) == 1 + + # Verify run_id and timestamp are included + assert per_file_df["run_id"].iloc[0] == run_id + assert per_file_df["timestamp"].iloc[0] == timestamp + assert summary_df["run_id"].iloc[0] == run_id + assert summary_df["timestamp"].iloc[0] == timestamp + + def test_get_usage_dataframes_multiple_files_models(self, mocker): + """Test getting DataFrames with multiple files and models.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_tracking.track_usage("file1.txt", "model2", 200, 100, 20, 30) + usage_tracking.track_usage("file2.txt", "model1", 150, 75, 15, 25) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp) + + assert len(per_file_df) == 3 # 3 file/model combinations + assert len(summary_df) == 1 # Single summary row + + def test_get_usage_dataframes_independent(self, mocker): + """Test that modifying one DataFrame doesn't affect the other.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + per_file_df, summary_df = usage_tracking.get_usage_dataframes(run_id, timestamp) + + # Modify per_file_df + per_file_df.loc[0, "input_tokens"] = 999 + + # Verify summary_df is unchanged + assert summary_df["total_fresh_input_tokens"].iloc[0] == 100 + + +class TestGeneratePerFilePerModelDf: + """Tests for _generate_per_file_per_model_df() function - Phase 3.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_generate_per_file_per_model_df_empty(self): + """Test generating DataFrame when no usage data exists.""" + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + assert isinstance(df, pd.DataFrame) + assert df.empty + + def test_generate_per_file_per_model_df_single_file_single_model(self, mocker): + """Test generating DataFrame with single file and model.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "anthropic.claude-3-5-sonnet-20240620-v1:0", 100, 50, 10, 20) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + assert len(df) == 1 + assert list(df.columns) == [ + "run_id", + "timestamp", + "file_name", + "model_name", + "input_tokens", + "output_tokens", + "cache_creation_tokens", + "cache_read_tokens", + "total_tokens", + "cost_per_token", + "total_cost", + ] + + # Verify values + assert df["run_id"].iloc[0] == run_id + assert df["timestamp"].iloc[0] == timestamp + assert df["file_name"].iloc[0] == "file1.txt" + assert df["model_name"].iloc[0] == "Claude 3.5 Sonnet" # Human-readable + assert df["input_tokens"].iloc[0] == 100 + assert df["output_tokens"].iloc[0] == 50 + assert df["cache_creation_tokens"].iloc[0] == 10 + assert df["cache_read_tokens"].iloc[0] == 20 + assert df["total_tokens"].iloc[0] == 180 # 100 + 50 + 10 + 20 + + def test_generate_per_file_per_model_df_multiple_files_single_model(self, mocker): + """Test generating DataFrame with multiple files and single model.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_tracking.track_usage("file2.txt", "model1", 200, 100, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + assert len(df) == 2 + assert set(df["file_name"].values) == {"file1.txt", "file2.txt"} + + def test_generate_per_file_per_model_df_single_file_multiple_models(self, mocker): + """Test generating DataFrame with single file and multiple models.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_tracking.track_usage("file1.txt", "model2", 200, 100, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + assert len(df) == 2 + assert set(df["model_name"].values) == {"Claude Unknown", "Claude Unknown"} # Fallback for unknown models + + def test_generate_per_file_per_model_df_cost_calculation(self, mocker): + """Test cost_per_token calculation and rounding.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 1000, 500, 200, 300) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + # Verify cost_per_token is calculated correctly (total_cost / total_tokens) + total_tokens = 1000 + 500 + 200 + 300 # 2000 + total_cost = ( + (1000 * input_cost) + + (200 * input_cost) # Cache creation uses input rate + + (300 * cache_read_cost) + + (500 * output_cost) + ) + expected_cost_per_token = total_cost / total_tokens + + # Allow for floating point precision (rounding happens in the function) + assert abs(df["cost_per_token"].iloc[0] - round(expected_cost_per_token, 8)) < 1e-10 + assert df["cost_per_token"].iloc[0] == round(expected_cost_per_token, 8) + assert df["total_cost"].iloc[0] == round(total_cost, 6) + + def test_generate_per_file_per_model_df_skips_zero_requests(self, mocker): + """Test that files/models with request_count=0 are skipped.""" + # This is tricky - we need to create a defaultdict entry but with 0 requests + # In practice, this shouldn't happen, but we test the guard clause + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # Track some usage + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_per_file_per_model_df(run_id, timestamp) + + # Should only have the one entry we tracked + assert len(df) == 1 + assert df["file_name"].iloc[0] == "file1.txt" + + +class TestGenerateBatchSummaryDf: + """Tests for _generate_batch_summary_df() function - Phase 3.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_generate_batch_summary_df_empty(self): + """Test generating summary DataFrame when no usage data exists.""" + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + assert isinstance(df, pd.DataFrame) + assert len(df) == 1 # Single row + assert df["total_tokens"].iloc[0] == 0 + assert df["total_cost"].iloc[0] == 0.0 + + def test_generate_batch_summary_df_single_file(self, mocker): + """Test generating summary DataFrame with single file.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + assert len(df) == 1 + assert df["total_fresh_input_tokens"].iloc[0] == 100 + assert df["total_cache_tokens"].iloc[0] == 30 # 10 + 20 + assert df["total_input_tokens"].iloc[0] == 130 # 100 + 30 + assert df["total_output_tokens"].iloc[0] == 50 + assert df["total_tokens"].iloc[0] == 180 # 100 + 50 + 10 + 20 + + # Verify averages (divide by 1 file) + assert df["average_fresh_input_tokens"].iloc[0] == 100.0 + assert df["average_cache_tokens"].iloc[0] == 30.0 + + def test_generate_batch_summary_df_multiple_files(self, mocker): + """Test generating summary DataFrame with multiple files.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + usage_tracking.track_usage("file2.txt", "model1", 200, 100, 20, 30) + usage_tracking.track_usage("file3.txt", "model2", 150, 75, 15, 25) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + assert len(df) == 1 + # Verify totals + assert df["total_fresh_input_tokens"].iloc[0] == 450 # 100 + 200 + 150 + assert df["total_cache_tokens"].iloc[0] == 120 # (10+20) + (20+30) + (15+25) + assert df["total_input_tokens"].iloc[0] == 570 # 450 + 120 + + # Verify averages (divide by 3 files) + assert df["average_fresh_input_tokens"].iloc[0] == 150.0 # 450 / 3 + assert df["average_cache_tokens"].iloc[0] == 40.0 # 120 / 3 + + def test_generate_batch_summary_df_columns(self, mocker): + """Test that summary DataFrame has all required columns.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + expected_columns = [ + "run_id", + "timestamp", + "average_fresh_input_tokens", + "average_cache_tokens", + "average_input_tokens", + "average_output_tokens", + "average_total_tokens", + "average_cost", + "total_fresh_input_tokens", + "total_cache_tokens", + "total_input_tokens", + "total_output_tokens", + "total_tokens", + "total_cost", + "region_mode", + ] + + assert list(df.columns) == expected_columns + + def test_generate_batch_summary_df_region_mode_single(self, mocker): + """Test region_mode is set correctly for single-region mode.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", False) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + assert df["region_mode"].iloc[0] == "single-region" + + def test_generate_batch_summary_df_region_mode_cross(self, mocker): + """Test region_mode is set correctly for cross-region mode.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + mocker.patch("src.utils.usage_tracking.config.ENABLE_RUNTIME_ROTATION", True) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + assert df["region_mode"].iloc[0] == "cross-region" + + def test_generate_batch_summary_df_rounding(self, mocker): + """Test that averages are properly rounded.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + # Create scenario where division results in many decimal places + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + usage_tracking.track_usage("file2.txt", "model1", 333, 167, 0, 0) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + # Verify rounding (to 2 decimal places for averages) + assert isinstance(df["average_fresh_input_tokens"].iloc[0], float) + # 433 / 2 = 216.5, should be rounded to 2 decimals + assert df["average_fresh_input_tokens"].iloc[0] == 216.5 + + def test_generate_batch_summary_df_token_relationship(self, mocker): + """Test that total_fresh_input_tokens + total_cache_tokens = total_input_tokens.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + + run_id = "test_batch" + timestamp = "run_20240101_12-00_test_batch" + df = usage_tracking._generate_batch_summary_df(run_id, timestamp) + + fresh = df["total_fresh_input_tokens"].iloc[0] + cache = df["total_cache_tokens"].iloc[0] + total_input = df["total_input_tokens"].iloc[0] + + assert fresh + cache == total_input + + +class TestExportUsageCostCsvsLocal: + """Tests for export_usage_cost_csvs_local() function - Phase 3.""" + + @pytest.fixture(autouse=True) + def reset_before_test(self): + """Reset usage data before each test.""" + usage_tracking.reset_usage_data() + yield + usage_tracking.reset_usage_data() + + def test_export_usage_cost_csvs_local_with_data(self, mocker, tmp_path, caplog): + """Test exporting CSV files with valid usage data.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + mocker.patch( + "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path) + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 10, 20) + + run_id = "test_batch" + run_timestamp = "run_20240101_12-00_test_batch" + + with caplog.at_level(logging.INFO): + usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp) + + # Verify directory was created + output_dir = tmp_path / run_timestamp + assert output_dir.exists() + + # Verify both files were created + usage_file = output_dir / f"{run_id}-USAGE.csv" + summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv" + + assert usage_file.exists() + assert summary_file.exists() + + # Verify file contents can be read back + usage_df = pd.read_csv(usage_file) + summary_df = pd.read_csv(summary_file) + + assert len(usage_df) == 1 + assert len(summary_df) == 1 + assert usage_df["file_name"].iloc[0] == "file1.txt" + + # Verify logging + assert "Saved usage data to:" in caplog.text + assert "Saved usage summary to:" in caplog.text + + def test_export_usage_cost_csvs_local_empty_data(self, mocker, tmp_path, caplog): + """Test exporting CSV files with empty usage data.""" + mocker.patch( + "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path) + ) + + run_id = "test_batch" + run_timestamp = "run_20240101_12-00_test_batch" + + with caplog.at_level(logging.WARNING): + usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp) + + # Verify directory was created + output_dir = tmp_path / run_timestamp + assert output_dir.exists() + + # Verify files were NOT created (empty data) + usage_file = output_dir / f"{run_id}-USAGE.csv" + summary_file = output_dir / f"{run_id}-USAGE-SUMMARY.csv" + + assert not usage_file.exists() + assert summary_file.exists() # Summary should still exist (with zeros) + + # Verify logging + assert "No usage data to export" in caplog.text + + def test_export_usage_cost_csvs_local_directory_creation(self, mocker, tmp_path): + """Test that directory is created if it doesn't exist.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + mocker.patch( + "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path) + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + run_timestamp = "run_20240101_12-00_test_batch" + + # Directory shouldn't exist yet + output_dir = tmp_path / run_timestamp + assert not output_dir.exists() + + usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp) + + # Directory should now exist + assert output_dir.exists() + + def test_export_usage_cost_csvs_local_file_paths(self, mocker, tmp_path): + """Test that file paths are correct.""" + input_cost = 0.003 / 1000.0 + output_cost = 0.015 / 1000.0 + cache_read_cost = 0.0003 / 1000.0 + mocker.patch( + "src.utils.usage_tracking.get_cost_per_token", + return_value=(input_cost, output_cost, cache_read_cost), + ) + mocker.patch( + "src.utils.usage_tracking.config.CONSOLIDATED_OUTPUT_DIRECTORY", str(tmp_path) + ) + + usage_tracking.track_usage("file1.txt", "model1", 100, 50, 0, 0) + + run_id = "test_batch" + run_timestamp = "run_20240101_12-00_test_batch" + + usage_tracking.export_usage_cost_csvs_local(run_id, run_timestamp) + + # Verify exact file paths + expected_usage_path = tmp_path / run_timestamp / f"{run_id}-USAGE.csv" + expected_summary_path = tmp_path / run_timestamp / f"{run_id}-USAGE-SUMMARY.csv" + + assert expected_usage_path.exists() + assert expected_summary_path.exists() + + +class TestGetModelName: + """Tests for get_model_name() function - Phase 4.""" + + @pytest.mark.parametrize( + "model_id, expected_name", + [ + # Claude 2 + ("anthropic.claude-instant-v1", "Claude 2"), + ("anthropic.claude-2", "Claude 2"), # Note: claude-v2 doesn't match, needs claude-2 + ("ANTHROPIC.CLAUDE-INSTANT-V1", "Claude 2"), # Uppercase + ("Anthropic.Claude-2", "Claude 2"), # Mixed case + # Claude 3 Haiku + ("anthropic.claude-3-haiku-20240307-v1:0", "Claude 3 Haiku"), + ("anthropic.claude-3.5-haiku-20240307-v1:0", "Claude 3 Haiku"), + ("ANTHROPIC.CLAUDE-3-HAIKU-20240307-V1:0", "Claude 3 Haiku"), # Uppercase + # Claude 3.5 Sonnet + ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"), + ("anthropic.claude-3-5-sonnet-20241022-v2:0", "Claude 3.5 Sonnet"), + ("anthropic.claude-3.5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"), + ("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"), # Mixed case + # Claude 3.7 Sonnet + ("us.anthropic.claude-3-7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"), + ("us.anthropic.claude-3.7-sonnet-20250219-v1:0", "Claude 3.7 Sonnet"), + ("US.ANTHROPIC.CLAUDE-3-7-SONNET-20250219-V1:0", "Claude 3.7 Sonnet"), # Uppercase + # Claude 4 Sonnet + ("us.anthropic.claude-sonnet-4-20250514-v1:0", "Claude 4 Sonnet"), + ("us.anthropic.claude-4-sonnet-20250514-v1:0", "Claude 4 Sonnet"), + ("US.ANTHROPIC.CLAUDE-SONNET-4-20250514-V1:0", "Claude 4 Sonnet"), # Uppercase + # Claude 4.5 Sonnet + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"), + ("us.anthropic.claude-sonnet-4.5-20250929-v1:0", "Claude 4.5 Sonnet"), + # Note: claude-4-5-sonnet and claude-4.5-sonnet don't match the pattern + # They need claude-sonnet-4 or claude-4-sonnet structure + ("US.ANTHROPIC.CLAUDE-SONNET-4-5-20250929-V1:0", "Claude 4.5 Sonnet"), # Uppercase + ], + ) + def test_get_model_name_valid_models(self, model_id, expected_name): + """Test model name conversion for all supported models.""" + result = usage_tracking.get_model_name(model_id) + assert result == expected_name + + @pytest.mark.parametrize( + "model_id, expected_pattern", + [ + # Unknown models should use fallback + ("unknown.model.id", "Claude "), # Should extract something or "Unknown" + ("anthropic.claude-999", "Claude "), # Unknown version + ("some.random.model", "Claude "), # Completely unknown + ("", "Claude Unknown"), # Empty string + ], + ) + def test_get_model_name_unknown_models(self, model_id, expected_pattern): + """Test fallback behavior for unknown model IDs.""" + result = usage_tracking.get_model_name(model_id) + assert result.startswith("Claude ") + # Should either extract a version or say "Unknown" + assert "Unknown" in result or result != "Claude " + + def test_get_model_name_case_insensitivity(self): + """Test that function handles case variations correctly.""" + # Test various case combinations + test_cases = [ + ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"), + ("ANTHROPIC.CLAUDE-3-5-SONNET-20240620-V1:0", "Claude 3.5 Sonnet"), + ("Anthropic.Claude-3-5-Sonnet-20240620-V1:0", "Claude 3.5 Sonnet"), + ("AnThRoPiC.cLaUdE-3-5-sOnNeT-20240620-v1:0", "Claude 3.5 Sonnet"), + ] + + for model_id, expected in test_cases: + result = usage_tracking.get_model_name(model_id) + assert result == expected + + def test_get_model_name_malformed_ids(self): + """Test handling of malformed model IDs.""" + # Test various malformed inputs + malformed_ids = [ + "claude", # Too short + ".", # Just a dot + "anthropic.", # Ends with dot + ".claude", # Starts with dot + "anthropic.claude-3-5", # Missing version info + ] + + for model_id in malformed_ids: + result = usage_tracking.get_model_name(model_id) + # Should not crash and should return something starting with "Claude" + assert isinstance(result, str) + assert result.startswith("Claude") + + def test_get_model_name_special_characters(self): + """Test handling of special characters in model IDs.""" + # Test with special characters that might appear + special_cases = [ + ("anthropic.claude-3-5-sonnet-20240620-v1:0", "Claude 3.5 Sonnet"), + ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", "Claude 4.5 Sonnet"), + ] + + for model_id, expected in special_cases: + result = usage_tracking.get_model_name(model_id) + assert result == expected + + def test_get_model_name_fallback_extraction(self): + """Test fallback extraction logic for unknown models.""" + # Test models that should trigger fallback extraction + # Fallback extracts: model_id.split('.')[-1].split('-')[0] + test_cases = [ + ("anthropic.claude-5-sonnet-v1:0", "Claude claude"), # Extracts "claude" from "claude-5-sonnet-v1:0" + ("anthropic.claude-10-v1:0", "Claude claude"), # Extracts "claude" from "claude-10-v1:0" + ("some.prefix.model-7-v1:0", "Claude model"), # Extracts "model" from "model-7-v1:0" + ] + + for model_id, expected in test_cases: + result = usage_tracking.get_model_name(model_id) + assert result == expected + + def test_get_model_name_no_dot_fallback(self): + """Test fallback when model ID has no dot.""" + # Model IDs without dots should return "Claude Unknown" + # BUT: if they match known patterns (like "claude-3-5-sonnet"), they still work + test_cases = [ + ("claude-3-5-sonnet", "Claude 3.5 Sonnet"), # Matches pattern even without dot + ("claude-instant-v1", "Claude 2"), # Matches pattern + ("anthropic-claude-3", "Claude Unknown"), # No dot, doesn't match pattern + ("just-text", "Claude Unknown"), # No dot, doesn't match pattern + ] + + for model_id, expected in test_cases: + result = usage_tracking.get_model_name(model_id) + assert result == expected + + def test_get_model_name_all_supported_models(self): + """Test all officially supported model IDs.""" + supported_models = { + "anthropic.claude-instant-v1": "Claude 2", + "anthropic.claude-2": "Claude 2", # Note: claude-v2 doesn't match pattern + "anthropic.claude-3-haiku-20240307-v1:0": "Claude 3 Haiku", + "anthropic.claude-3-5-sonnet-20240620-v1:0": "Claude 3.5 Sonnet", + "anthropic.claude-3-5-sonnet-20241022-v2:0": "Claude 3.5 Sonnet", + "us.anthropic.claude-3-7-sonnet-20250219-v1:0": "Claude 3.7 Sonnet", + "us.anthropic.claude-sonnet-4-20250514-v1:0": "Claude 4 Sonnet", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0": "Claude 4.5 Sonnet", + } + + for model_id, expected_name in supported_models.items(): + result = usage_tracking.get_model_name(model_id) + assert result == expected_name, f"Failed for {model_id}: got {result}, expected {expected_name}" + diff --git a/poetry.lock b/poetry.lock index 7307aa1..fdccb2d 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "annotated-types" @@ -6,6 +6,7 @@ version = "0.7.0" description = "Reusable constraint types to use with typing.Annotated" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53"}, {file = "annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89"}, @@ -17,6 +18,7 @@ version = "0.36.0" description = "The official Python library for the anthropic API" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "anthropic-0.36.0-py3-none-any.whl", hash = "sha256:9183b9eaa0f409f2047244d7ef02c9c3eb916959c0b2960f7605dcb6cabbf548"}, {file = "anthropic-0.36.0.tar.gz", hash = "sha256:7b0b1457096605572a29559d9a8ce224b9389d379b410e7d1bf5e0c1379f9ee2"}, @@ -42,6 +44,7 @@ version = "4.6.2.post1" description = "High level compatibility layer for multiple asynchronous event loop implementations" optional = false python-versions = ">=3.9" +groups = ["main"] files = [ {file = "anyio-4.6.2.post1-py3-none-any.whl", hash = "sha256:6d170c36fba3bdd840c73d3868c1e777e33676a69c3a72cf0a0d5d6d8009b61d"}, {file = "anyio-4.6.2.post1.tar.gz", hash = "sha256:4c8bc31ccdb51c7f7bd251f51c609e038d63e34219b44aa86e47576389880b4c"}, @@ -53,7 +56,7 @@ sniffio = ">=1.1" [package.extras] doc = ["Sphinx (>=7.4,<8.0)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] -test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1)", "uvloop (>=0.21.0b1)"] +test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "truststore (>=0.9.1) ; python_version >= \"3.10\"", "uvloop (>=0.21.0b1) ; platform_python_implementation == \"CPython\" and platform_system != \"Windows\""] trio = ["trio (>=0.26.1)"] [[package]] @@ -62,6 +65,7 @@ version = "24.10.0" description = "The uncompromising code formatter." optional = false python-versions = ">=3.9" +groups = ["dev"] files = [ {file = "black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812"}, {file = "black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea"}, @@ -106,6 +110,7 @@ version = "1.35.40" description = "The AWS SDK for Python" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "boto3-1.35.40-py3-none-any.whl", hash = "sha256:9352f6d61f15c789231a5d608613f03425059072ed862c32e1ed102b17206abf"}, {file = "boto3-1.35.40.tar.gz", hash = "sha256:33c6a7aeab316f7e0b3ad8552afe95a4a10bfd58519d00741c4d4f3047da8382"}, @@ -125,6 +130,7 @@ version = "1.35.40" description = "Low-level, data-driven core of boto 3." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "botocore-1.35.40-py3-none-any.whl", hash = "sha256:072cc47f29cb1de4fa77ce6632e4f0480af29b70816973ff415fbaa3f50bd1db"}, {file = "botocore-1.35.40.tar.gz", hash = "sha256:547e0a983856c7d7aeaa30fca2a283873c57c07366cd806d2d639856341b3c31"}, @@ -144,6 +150,7 @@ version = "2024.8.30" description = "Python package for providing Mozilla's CA Bundle." optional = false python-versions = ">=3.6" +groups = ["main"] files = [ {file = "certifi-2024.8.30-py3-none-any.whl", hash = "sha256:922820b53db7a7257ffbda3f597266d435245903d80737e34f8a45ff3e3230d8"}, {file = "certifi-2024.8.30.tar.gz", hash = "sha256:bec941d2aa8195e248a60b31ff9f0558284cf01a52591ceda73ea9afffd69fd9"}, @@ -155,6 +162,7 @@ version = "3.4.0" description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." optional = false python-versions = ">=3.7.0" +groups = ["main"] files = [ {file = "charset_normalizer-3.4.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:4f9fc98dad6c2eaa32fc3af1417d95b5e3d08aff968df0cd320066def971f9a6"}, {file = "charset_normalizer-3.4.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0de7b687289d3c1b3e8660d0741874abe7888100efe14bd0f9fd7141bcbda92b"}, @@ -269,6 +277,7 @@ version = "8.1.7" description = "Composable command line interface toolkit" optional = false python-versions = ">=3.7" +groups = ["dev"] files = [ {file = "click-8.1.7-py3-none-any.whl", hash = "sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28"}, {file = "click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de"}, @@ -283,10 +292,12 @@ version = "0.4.6" description = "Cross-platform colored terminal text." optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" +groups = ["main", "dev"] files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] +markers = {main = "platform_system == \"Windows\"", dev = "sys_platform == \"win32\" or platform_system == \"Windows\""} [[package]] name = "distro" @@ -294,6 +305,7 @@ version = "1.9.0" description = "Distro - an OS platform information API" optional = false python-versions = ">=3.6" +groups = ["main"] files = [ {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, @@ -305,6 +317,7 @@ version = "3.16.1" description = "A platform independent file lock." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "filelock-3.16.1-py3-none-any.whl", hash = "sha256:2082e5703d51fbf98ea75855d9d5527e33d8ff23099bec374a134febee6946b0"}, {file = "filelock-3.16.1.tar.gz", hash = "sha256:c249fbfcd5db47e5e2d6d62198e565475ee65e4831e2561c8e313fa7eb961435"}, @@ -313,7 +326,7 @@ files = [ [package.extras] docs = ["furo (>=2024.8.6)", "sphinx (>=8.0.2)", "sphinx-autodoc-typehints (>=2.4.1)"] testing = ["covdefaults (>=2.3)", "coverage (>=7.6.1)", "diff-cover (>=9.2)", "pytest (>=8.3.3)", "pytest-asyncio (>=0.24)", "pytest-cov (>=5)", "pytest-mock (>=3.14)", "pytest-timeout (>=2.3.1)", "virtualenv (>=20.26.4)"] -typing = ["typing-extensions (>=4.12.2)"] +typing = ["typing-extensions (>=4.12.2) ; python_version < \"3.11\""] [[package]] name = "fsspec" @@ -321,6 +334,7 @@ version = "2024.9.0" description = "File-system specification" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "fsspec-2024.9.0-py3-none-any.whl", hash = "sha256:a0947d552d8a6efa72cc2c730b12c41d043509156966cca4fb157b0f2a0c574b"}, {file = "fsspec-2024.9.0.tar.gz", hash = "sha256:4b0afb90c2f21832df142f292649035d80b421f60a9e1c027802e5a0da2b04e8"}, @@ -360,6 +374,7 @@ version = "0.14.0" description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"}, {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"}, @@ -371,6 +386,7 @@ version = "1.0.6" description = "A minimal low-level HTTP client." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "httpcore-1.0.6-py3-none-any.whl", hash = "sha256:27b59625743b85577a8c0e10e55b50b5368a4f2cfe8cc7bcfa9cf00829c2682f"}, {file = "httpcore-1.0.6.tar.gz", hash = "sha256:73f6dbd6eb8c21bbf7ef8efad555481853f5f6acdeaff1edb0694289269ee17f"}, @@ -392,6 +408,7 @@ version = "0.27.2" description = "The next generation HTTP client." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0"}, {file = "httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2"}, @@ -405,7 +422,7 @@ idna = "*" sniffio = "*" [package.extras] -brotli = ["brotli", "brotlicffi"] +brotli = ["brotli ; platform_python_implementation == \"CPython\"", "brotlicffi ; platform_python_implementation != \"CPython\""] cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] @@ -417,6 +434,7 @@ version = "0.25.2" description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub" optional = false python-versions = ">=3.8.0" +groups = ["main"] files = [ {file = "huggingface_hub-0.25.2-py3-none-any.whl", hash = "sha256:1897caf88ce7f97fe0110603d8f66ac264e3ba6accdf30cd66cc0fed5282ad25"}, {file = "huggingface_hub-0.25.2.tar.gz", hash = "sha256:a1014ea111a5f40ccd23f7f7ba8ac46e20fa3b658ced1f86a00c75c06ec6423c"}, @@ -451,6 +469,7 @@ version = "3.10" description = "Internationalized Domain Names in Applications (IDNA)" optional = false python-versions = ">=3.6" +groups = ["main"] files = [ {file = "idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3"}, {file = "idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9"}, @@ -465,6 +484,7 @@ version = "2.0.0" description = "brain-dead simple config-ini parsing" optional = false python-versions = ">=3.7" +groups = ["dev"] files = [ {file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"}, {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"}, @@ -476,6 +496,7 @@ version = "0.6.1" description = "Fast iterable JSON parser." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "jiter-0.6.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:d08510593cb57296851080018006dfc394070178d238b767b1879dc1013b106c"}, {file = "jiter-0.6.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:adef59d5e2394ebbad13b7ed5e0306cceb1df92e2de688824232a91588e77aa7"}, @@ -558,6 +579,7 @@ version = "1.0.1" description = "JSON Matching Expressions" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "jmespath-1.0.1-py3-none-any.whl", hash = "sha256:02e2e4cc71b5bcab88332eebf907519190dd9e6e82107fa7f83b1003a6252980"}, {file = "jmespath-1.0.1.tar.gz", hash = "sha256:90261b206d6defd58fdd5e85f478bf633a2901798906be2ad389150c5c60edbe"}, @@ -569,6 +591,7 @@ version = "1.0.0" description = "Type system extensions for programs checked with the mypy type checker." optional = false python-versions = ">=3.5" +groups = ["dev"] files = [ {file = "mypy_extensions-1.0.0-py3-none-any.whl", hash = "sha256:4392f6c0eb8a5668a69e23d168ffa70f0be9ccfd32b5cc2d26a34ae5b844552d"}, {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"}, @@ -580,6 +603,7 @@ version = "2.1.2" description = "Fundamental package for array computing in Python" optional = false python-versions = ">=3.10" +groups = ["main"] files = [ {file = "numpy-2.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:30d53720b726ec36a7f88dc873f0eec8447fbc93d93a8f079dfac2629598d6ee"}, {file = "numpy-2.1.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8d3ca0a72dd8846eb6f7dfe8f19088060fcb76931ed592d29128e0219652884"}, @@ -642,6 +666,7 @@ version = "24.1" description = "Core utilities for Python packages" optional = false python-versions = ">=3.8" +groups = ["main", "dev"] files = [ {file = "packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124"}, {file = "packaging-24.1.tar.gz", hash = "sha256:026ed72c8ed3fcce5bf8950572258698927fd1dbda10a5e981cdf0ac37f4f002"}, @@ -653,6 +678,7 @@ version = "2.2.3" description = "Powerful data structures for data analysis, time series, and statistics" optional = false python-versions = ">=3.9" +groups = ["main"] files = [ {file = "pandas-2.2.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:1948ddde24197a0f7add2bdc4ca83bf2b1ef84a1bc8ccffd95eda17fd836ecb5"}, {file = "pandas-2.2.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:381175499d3802cde0eabbaf6324cce0c4f5d52ca6f8c377c29ad442f50f6348"}, @@ -735,6 +761,7 @@ version = "0.12.1" description = "Utility library for gitignore style pattern matching of file paths." optional = false python-versions = ">=3.8" +groups = ["dev"] files = [ {file = "pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08"}, {file = "pathspec-0.12.1.tar.gz", hash = "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712"}, @@ -746,6 +773,7 @@ version = "11.2.1" description = "Python Imaging Library (Fork)" optional = false python-versions = ">=3.9" +groups = ["main"] files = [ {file = "pillow-11.2.1-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:d57a75d53922fc20c165016a20d9c44f73305e67c351bbc60d1adaf662e74047"}, {file = "pillow-11.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:127bf6ac4a5b58b3d32fc8289656f77f80567d65660bc46f72c0d77e6600cc95"}, @@ -836,7 +864,7 @@ fpx = ["olefile"] mic = ["olefile"] test-arrow = ["pyarrow"] tests = ["check-manifest", "coverage (>=7.4.2)", "defusedxml", "markdown2", "olefile", "packaging", "pyroma", "pytest", "pytest-cov", "pytest-timeout", "trove-classifiers (>=2024.10.12)"] -typing = ["typing-extensions"] +typing = ["typing-extensions ; python_version < \"3.10\""] xmp = ["defusedxml"] [[package]] @@ -845,6 +873,7 @@ version = "4.3.6" description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`." optional = false python-versions = ">=3.8" +groups = ["dev"] files = [ {file = "platformdirs-4.3.6-py3-none-any.whl", hash = "sha256:73e575e1408ab8103900836b97580d5307456908a03e92031bab39e4554cc3fb"}, {file = "platformdirs-4.3.6.tar.gz", hash = "sha256:357fb2acbc885b0419afd3ce3ed34564c13c9b95c89360cd9563f73aa5e2b907"}, @@ -861,6 +890,7 @@ version = "1.5.0" description = "plugin and hook calling mechanisms for python" optional = false python-versions = ">=3.8" +groups = ["dev"] files = [ {file = "pluggy-1.5.0-py3-none-any.whl", hash = "sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669"}, {file = "pluggy-1.5.0.tar.gz", hash = "sha256:2cffa88e94fdc978c4c574f15f9e59b7f4201d439195c3715ca9e2486f1d0cf1"}, @@ -876,6 +906,7 @@ version = "2.9.2" description = "Data validation using Python type hints" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "pydantic-2.9.2-py3-none-any.whl", hash = "sha256:f048cec7b26778210e28a0459867920654d48e5e62db0958433636cde4254f12"}, {file = "pydantic-2.9.2.tar.gz", hash = "sha256:d155cef71265d1e9807ed1c32b4c8deec042a44a50a4188b25ac67ecd81a9c0f"}, @@ -891,7 +922,7 @@ typing-extensions = [ [package.extras] email = ["email-validator (>=2.0.0)"] -timezone = ["tzdata"] +timezone = ["tzdata ; python_version >= \"3.9\" and sys_platform == \"win32\""] [[package]] name = "pydantic-core" @@ -899,6 +930,7 @@ version = "2.23.4" description = "Core functionality for Pydantic validation and serialization" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "pydantic_core-2.23.4-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:b10bd51f823d891193d4717448fab065733958bdb6a6b351967bd349d48d5c9b"}, {file = "pydantic_core-2.23.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4fc714bdbfb534f94034efaa6eadd74e5b93c8fa6315565a222f7b6f42ca1166"}, @@ -1000,6 +1032,7 @@ version = "1.25.5" description = "A high performance Python library for data extraction, analysis, conversion & manipulation of PDF (and other) documents." optional = false python-versions = ">=3.9" +groups = ["main"] files = [ {file = "pymupdf-1.25.5-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:cde4e1c9cfb09c0e1e9c2b7f4b787dd6bb34a32cfe141a4675e24af7c0c25dd3"}, {file = "pymupdf-1.25.5-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:5a35e2725fae0ab57f058dff77615c15eb5961eac50ba04f41ebc792cd8facad"}, @@ -1017,6 +1050,7 @@ version = "8.3.3" description = "pytest: simple powerful testing with Python" optional = false python-versions = ">=3.8" +groups = ["dev"] files = [ {file = "pytest-8.3.3-py3-none-any.whl", hash = "sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2"}, {file = "pytest-8.3.3.tar.gz", hash = "sha256:70b98107bd648308a7952b06e6ca9a50bc660be218d53c257cc1fc94fda10181"}, @@ -1037,6 +1071,7 @@ version = "2.9.0.post0" description = "Extensions to the standard Python datetime module" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" +groups = ["main"] files = [ {file = "python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3"}, {file = "python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427"}, @@ -1045,12 +1080,28 @@ files = [ [package.dependencies] six = ">=1.5" +[[package]] +name = "python-dotenv" +version = "1.2.1" +description = "Read key-value pairs from a .env file and set them as environment variables" +optional = false +python-versions = ">=3.9" +groups = ["main"] +files = [ + {file = "python_dotenv-1.2.1-py3-none-any.whl", hash = "sha256:b81ee9561e9ca4004139c6cbba3a238c32b03e4894671e181b671e8cb8425d61"}, + {file = "python_dotenv-1.2.1.tar.gz", hash = "sha256:42667e897e16ab0d66954af0e60a9caa94f0fd4ecf3aaf6d2d260eec1aa36ad6"}, +] + +[package.extras] +cli = ["click (>=5.0)"] + [[package]] name = "pytz" version = "2024.2" description = "World timezone definitions, modern and historical" optional = false python-versions = "*" +groups = ["main"] files = [ {file = "pytz-2024.2-py2.py3-none-any.whl", hash = "sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725"}, {file = "pytz-2024.2.tar.gz", hash = "sha256:2aa355083c50a0f93fa581709deac0c9ad65cca8a9e9beac660adcbd493c798a"}, @@ -1062,6 +1113,7 @@ version = "6.0.2" description = "YAML parser and emitter for Python" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "PyYAML-6.0.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0a9a2848a5b7feac301353437eb7d5957887edbf81d56e903999a75a3d743086"}, {file = "PyYAML-6.0.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:29717114e51c84ddfba879543fb232a6ed60086602313ca38cce623c1d62cfbf"}, @@ -1124,6 +1176,7 @@ version = "2.32.3" description = "Python HTTP for Humans." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "requests-2.32.3-py3-none-any.whl", hash = "sha256:70761cfe03c773ceb22aa2f671b4757976145175cdfca038c02654d061d6dcc6"}, {file = "requests-2.32.3.tar.gz", hash = "sha256:55365417734eb18255590a9ff9eb97e9e1da868d4ccd6402399eaf68af20a760"}, @@ -1145,6 +1198,7 @@ version = "0.10.3" description = "An Amazon S3 Transfer Manager" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "s3transfer-0.10.3-py3-none-any.whl", hash = "sha256:263ed587a5803c6c708d3ce44dc4dfedaab4c1a32e8329bab818933d79ddcf5d"}, {file = "s3transfer-0.10.3.tar.gz", hash = "sha256:4f50ed74ab84d474ce614475e0b8d5047ff080810aac5d01ea25231cfc944b0c"}, @@ -1162,6 +1216,7 @@ version = "1.16.0" description = "Python 2 and 3 compatibility utilities" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" +groups = ["main"] files = [ {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"}, {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, @@ -1173,6 +1228,7 @@ version = "1.3.1" description = "Sniff out which async library your code is running under" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2"}, {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, @@ -1184,6 +1240,7 @@ version = "0.20.1" description = "" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "tokenizers-0.20.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:439261da7c0a5c88bda97acb284d49fbdaf67e9d3b623c0bfd107512d22787a9"}, {file = "tokenizers-0.20.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:03dae629d99068b1ea5416d50de0fea13008f04129cc79af77a2a6392792d93c"}, @@ -1301,6 +1358,7 @@ version = "4.66.5" description = "Fast, Extensible Progress Meter" optional = false python-versions = ">=3.7" +groups = ["main"] files = [ {file = "tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd"}, {file = "tqdm-4.66.5.tar.gz", hash = "sha256:e1020aef2e5096702d8a025ac7d16b1577279c9d63f8375b63083e9a5f0fcbad"}, @@ -1321,6 +1379,7 @@ version = "4.12.2" description = "Backported and Experimental Type Hints for Python 3.8+" optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d"}, {file = "typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8"}, @@ -1332,6 +1391,7 @@ version = "2024.2" description = "Provider of IANA time zone data" optional = false python-versions = ">=2" +groups = ["main"] files = [ {file = "tzdata-2024.2-py2.py3-none-any.whl", hash = "sha256:a48093786cdcde33cad18c2555e8532f34422074448fbc874186f0abd79565cd"}, {file = "tzdata-2024.2.tar.gz", hash = "sha256:7d85cc416e9382e69095b7bdf4afd9e3880418a2413feec7069d533d6b4e31cc"}, @@ -1343,18 +1403,19 @@ version = "2.2.3" description = "HTTP library with thread-safe connection pooling, file post, and more." optional = false python-versions = ">=3.8" +groups = ["main"] files = [ {file = "urllib3-2.2.3-py3-none-any.whl", hash = "sha256:ca899ca043dcb1bafa3e262d73aa25c465bfb49e0bd9dd5d59f1d0acba2f8fac"}, {file = "urllib3-2.2.3.tar.gz", hash = "sha256:e7d814a81dad81e6caf2ec9fdedb284ecc9c73076b62654547cc64ccdcae26e9"}, ] [package.extras] -brotli = ["brotli (>=1.0.9)", "brotlicffi (>=0.8.0)"] +brotli = ["brotli (>=1.0.9) ; platform_python_implementation == \"CPython\"", "brotlicffi (>=0.8.0) ; platform_python_implementation != \"CPython\""] h2 = ["h2 (>=4,<5)"] socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] zstd = ["zstandard (>=0.18.0)"] [metadata] -lock-version = "2.0" +lock-version = "2.1" python-versions = "^3.12" -content-hash = "11f240679ad504f8cfc3a866067918339f2dd77f3764b1bbf77f992dac4e7bb3" +content-hash = "c7560454f4aab51b844fa489cfa6844f7800703a9ac87a55c02d05226fb1cefc" diff --git a/pyproject.toml b/pyproject.toml index 31b9c7e..0b887ae 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ boto3 = "^1.35.40" anthropic = "^0.36.0" pymupdf = "^1.25.5" pillow = "^11.2.1" +python-dotenv = "^1.0.0" [tool.poetry.group.dev.dependencies] pytest = "^8.3.3"