test: add tests for token usage extraction in blue and red worker agents#152
Merged
test: add tests for token usage extraction in blue and red worker agents#152
Conversation
…er agents **Added:** - Added extraction of token usage metrics (input_tokens, output_tokens, total_tokens) from agent results in BlueRedisWorkerAgent, BlueWorkerAgent, RedisWorkerAgent, and WorkerAgent. These metrics are included in the result payload for downstream aggregation and tracing. - Instrumented OpenTelemetry/Tempo spans with GenAI semantic conventions for token usage in worker agents. - Introduced unit tests in both `tests/core/blue_worker/test_redis_worker.py` and `tests/core/worker/test_worker.py` to verify token usage extraction and result payload inclusion for both blue and red worker agents. **Changed:** - Updated blue and red worker agent task processing logic to check for agent result usage metrics and include them in the result payload, whether the task completes via callback or partial completion. - Improved docstrings for helper and internal methods throughout the codebase to clarify argument and return types, increasing maintainability and clarity. - Refactored the evaluation workflow to aggregate token usage statistics from completed task results when available, falling back to estimation only if no usage data is present. - Enhanced logging and metrics reporting for token usage in orchestrators and workers. **Removed:** - Removed redundant or outdated comments where code and docstrings are now self-explanatory. - Eliminated placeholder or non-functional code for attack chain reconstruction and other "future work" comments, replacing with concise markers for unimplemented features.
**Changed:** - Removed explanatory comments that repeated obvious code logic in several modules, including agent initialization, prompt building, and data extraction routines - Improved code readability and maintainability by reducing noise and focusing on non-obvious operations
**Changed:** - Removed explanatory and redundant inline comments that simply restate code logic across multiple modules, including blue agent tools, investigation orchestration, CLI commands, core models, recovery routines, task queue, tracing, workflows, evaluation logic, and blue team tools - Improved code readability and reduced noise by eliminating comments that do not provide additional context or rationale beyond what the code already expresses
**Changed:** - Removed redundant or obsolete commented-out code throughout multiple modules to improve code readability and maintainability. This includes eliminating comments that merely restate the following code or describe obvious steps, focusing on reducing visual clutter in the codebase.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Key Changes:
Added:
TestBlueRedisWorkerAgentUsageExtractiontest class totests/core/blue_worker/test_redis_worker.pywith tests for:TestRedisWorkerAgentUsageExtractiontest class totests/core/worker/test_worker.pycovering:Changed:
usagefields in agent results