T

Praneel Panchigar ff400f4b1c Merged in feature/llm-utils-refactor (pull request #776 )

Feature/llm utils refactor

* - 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

* Phase 1: Remove unused code and align cost constants

- Remove count_tokens() function (unused, replaced by exact API counts)
- Remove unused cost constants from invoke_claude() and invoke_claude_with_image_array()
- Remove unused math import
- Update usage_tracking.py to match correct cost constants from llm_utils.py:
  - Corrected Haiku cache_read: 0.000025 -> 0.00003
  - Implemented correct cache creation cost calculation:
    * Sonnet models: 25% higher than input (0.00375 vs 0.003)
    * Haiku: 20% higher than input (0.0003 vs 0.00025)
- Remove test_count_tokens() test
- Update test_track_usage_single_call() to account for corrected cache creation costs

All 437 tests passing. Token counts verified against previous runs.

* Phase 2: Remove redundant update_stats() tracking

- Remove all update_stats() calls from EC2 mode (3 locations)
- Remove update_stats() function definition
- Remove estimated token counting (tokens_sent, tokens_received)
- Remove elapsed_time calculations
- Remove unused start_time assignments
- Remove test_update_stats() test case
- Clean up setUp() method in test_llm_utils.py
- Update docstrings to reflect usage_tracking.py as source of truth

All token tracking now uses exact API counts via usage_tracking.track_usage().
All 436 tests passing.

* Merge main into feature/llm-utils-refactor

Resolved merge conflicts:
- llm_utils.py: Kept refactored code (removed legacy update_stats, math import)
- usage_tracking.py: Kept corrected cost constants and cache creation pricing logic
- io_utils.py: Merged main's new 'dtc' and 'dtc_summary' output types
- test_usage_tracking.py: Kept corrected test values for cache_read and cache_creation costs

All conflicts resolved while preserving Phase 1 & 2 refactoring work.

* Fix merge conflict artifact in test_usage_tracking.py

Removed leftover merge conflict marker (<<<<<<< HEAD) that was missed
during merge resolution. The test now correctly calculates cache creation
costs for Sonnet (25% higher) and Haiku (20% higher) models.

* Phase 3: Add thread safety and enhance cache key generation

Step 3.1: Add thread safety to cache operations
- Added _cache_lock = threading.Lock() at module level
- Wrapped all cache operations (get, set, move_to_end, popitem) with lock
- All 4 cache operation locations now thread-safe:
  * Cache check/retrieval in invoke_claude()
  * Cache storage in invoke_claude()
  * Cache check/retrieval in invoke_claude_with_image_array()
  * Cache storage in invoke_claude_with_image_array()

Step 3.2: Enhance cache key generation
- Updated get_cache_key() to include instruction, max_tokens, and cache flag
- Updated get_cache_key_with_image_array() with same enhancements
- Cache keys now differentiate based on all relevant parameters
- Backward compatible (all new parameters are optional)

Step 3.3: Update call sites
- Updated invoke_claude() to pass instruction, max_tokens, cache to get_cache_key()
- Updated invoke_claude_with_image_array() to pass max_tokens

All 412 tests passing. Cache operations are now thread-saf…
* Phase 4: Add comprehensive tests for cache operations and token tracking

Step 4.1: Remove legacy test cases
- Already completed in Phase 2 (test_update_stats, test_count_tokens removed)
- setUp() cleaned up (no MODEL_STATS/GLOBAL_STATS setup needed)

Step 4.2: Add tests for thread-safe cache operations
- test_cache_thread_safety_concurrent_reads: 10 concurrent cache reads
- test_cache_thread_safety_concurrent_writes: 10 concurrent cache writes
- test_cache_thread_safety_mixed_operations: mixed read/write operations

Step 4.3: Add tests for enhanced cache key generation
- test_get_cache_key_basic: backward compatibility test
- test_get_cache_key_with_instruction: instruction parameter
- test_get_cache_key_with_max_tokens: max_tokens parameter
- test_get_cache_key_with_cache_flag: cache flag
- test_get_cache_key_all_parameters: all parameters together
- test_get_cache_key_with_image_array: image array variant
- test_invoke_claude_uses_enhanced_cache_key: integration test

Step 4.4: Verify token tracking uses e…
* Phase 5: Replace legacy CSV export with usage_tracking.py

Step 5.1: Replace write_stats_to_csv() implementation
- Updated write_stats_to_csv() to use usage_tracking.get_usage_dataframes() instead of config.MODEL_STATS/GLOBAL_STATS
- Added import for usage_tracking module
- Function now reads from usage_tracking.py as the single source of truth

Step 5.2: Add optional run_id and timestamp parameters
- Added optional run_id parameter (defaults to config.BATCH_ID)
- Added optional timestamp parameter (defaults to current timestamp)
- Updated upload_tracking_logs() to pass batch_id and run_timestamp

Step 5.3: Transform usage_tracking format to legacy CSV format
- Aggregates per-file, per-model data to per-file data
- Maps model names to Claude 2/3/3.5 request counts
- Calculates tokens_sent (fresh input tokens) and tokens_received (output tokens)
- Maintains backward compatibility with legacy CSV structure
- Handles empty data gracefully (writes headers-only CSV)

Step 5.4: Verify call sites work with updated f…
* refactor(llm_utils): extract helper functions and reduce code duplication

- Add private helper functions: _track_usage_from_response(), _build_claude_3_request_body(), _store_in_cache()
- Simplify routing logic using model sets (_CLAUDE_2_MODELS, _CLAUDE_3_AND_UP_MODELS)
- Replace 8 instances of usage tracking pattern with single function call
- Replace 3 instances of request body building with helper function
- Replace 2 instances of cache storage with helper function
- Add Claude 4.5 Sonnet to cache support list
- Reduce file size from 1,004 to 860 lines (14.3% reduction, 144 lines removed)
- All tests passing (422 tests)
- Maintains backward compatibility and functionality

* Re-support Claude 3 Haiku and add comprehensive tests for new helper functions

- Re-added Claude 3 Haiku to supported models (removed from deprecated patterns)
- Updated model validation to include Haiku in supported models list
- Added comprehensive test coverage for new helper functions:
  - _validate_model_support() - 5 tests covering deprecation, support, and context
  - _supports_prompt_cache() - 2 tests for cache support detection
  - _build_claude_3_request_body() - 4 tests for request body building with/without cache
- Updated error messages to reflect Claude 3+ models (Haiku, Sonnet 3.5/3.7/4/4.5)
- All 23 tests passing

* Merged main into feature/llm-utils-refactor

* Merged main into feature/llm-utils-refactor


Approved-by: Siddhant Medar
Approved-by: Katon Minhas

2025-11-17 14:15:46 +00:00