feat: add token cost to mutation coverage

Author: jungs1

README.md

@@ -46,6 +46,8 @@ Mutahunter uses LLM models to inject context-aware faults into your codebase. Th
 
 ## Getting Started
 
+⚠️ We highly suggest using `--modified-files-only` flag to run mutation testing on only on modified files and `--only-mutate-file-paths` flag to focus on specific files. This will make the mutation testing **faster** and **cost effective.** ⚠️
+
 ```bash
 # Install Mutahunter package via GitHub. Python 3.11+ is required.
 $ pip install git+https://github.com/codeintegrity-ai/mutahunter.git
@@ -73,6 +75,7 @@ $ mutahunter run --test-command "pytest tests/unit" --code-coverage-report-path
 2024-07-05 00:26:13,421 INFO: 🕒 Timeout Mutants: 0 🕒
 2024-07-05 00:26:13,421 INFO: 🔥 Compile Error Mutants: 0 🔥
 2024-07-05 00:26:13,421 INFO: 🎯 Mutation Coverage: 61.54% 🎯
+2024-07-05 00:26:13,421 INFO: 💰 Total Cost: $0.00583 USD 💰
 2024-07-05 00:26:13,421 INFO: Report saved to logs/_latest/mutation_coverage.json
 2024-07-05 00:26:13,421 INFO: Report saved to logs/_latest/mutation_coverage_detail.json
 2024-07-05 00:26:13,421 INFO: Mutation Testing Ended. Took 43s

src/mutahunter/core/hunter.py

@@ -10,6 +10,7 @@
 from mutahunter.core.logger import logger
 from mutahunter.core.mutator import MutantGenerator
 from mutahunter.core.report import MutantReport
+from mutahunter.core.router import LLMRouter
 from mutahunter.core.runner import TestRunner
 
 
@@ -23,6 +24,9 @@ def __init__(self, config: Dict[str, Any]) -> None:
         self.mutant_report = MutantReport(config=self.config)
         self.analyzer = Analyzer(self.config)
         self.test_runner = TestRunner()
+        self.router = LLMRouter(
+            model=self.config["model"], api_base=self.config["api_base"]
+        )
 
     def run(self) -> None:
         """
@@ -40,7 +44,7 @@ def run(self) -> None:
                 logger.info("🦠 Running mutation testing on entire codebase... 🦠")
                 self.run_mutation_testing()
             logger.info("🎯 Generating Mutation Report... 🎯")
-            self.mutant_report.generate_report(self.mutants)
+            self.mutant_report.generate_report(self.mutants, self.router.total_cost)
             logger.info(f"Mutation Testing Ended. Took {round(time.time() - start)}s")
         except Exception as e:
             logger.error(
@@ -151,6 +155,7 @@ def generate_mutations(self, file_path: str, executed_lines: List[int]) -> None:
                 function_name=function_name,
                 start_byte=start_byte,
                 end_byte=end_byte,
+                router=self.router,
             )
             for mutant_data in mutant_generator.generate():
                 self.process_mutant(mutant_data, file_path, start_byte, end_byte)

src/mutahunter/core/mutator.py

@@ -6,7 +6,6 @@
 from mutahunter.core.logger import logger
 from mutahunter.core.pilot.aider.repomap import RepoMap
 from mutahunter.core.pilot.prompts.factory import PromptFactory
-from mutahunter.core.router import LLMRouter
 
 
 class MutantGenerator:
@@ -19,6 +18,7 @@ def __init__(
         function_name,
         start_byte,
         end_byte,
+        router,
     ):
         self.config = config
         self.executed_lines = executed_lines
@@ -27,10 +27,7 @@ def __init__(
         self.function_name = function_name
         self.start_byte = start_byte
         self.end_byte = end_byte
-
-        self.router = LLMRouter(
-            model=self.config["model"], api_base=self.config["api_base"]
-        )
+        self.router = router
         self.repo_map = RepoMap(model=self.config["model"])
         self.language = filename_to_lang(self.source_file_path)
         self.prompt = PromptFactory.get_prompt(language=self.language)

src/mutahunter/core/pilot/prompts/factory.py

@@ -2,12 +2,10 @@
 Module for generating prompts based on the programming language.
 """
 
-from mutahunter.core.pilot.prompts.examples import (
-    GO_EXAMPLE_OUTPUT,
-    JAVA_EXAMPLE_OUTPUT,
-    JAVASCRIPT_EXAMPLE_OUTPUT,
-    PYTHON_EXAMPLE_OUTPUT,
-)
+from mutahunter.core.pilot.prompts.examples import (GO_EXAMPLE_OUTPUT,
+                                                    JAVA_EXAMPLE_OUTPUT,
+                                                    JAVASCRIPT_EXAMPLE_OUTPUT,
+                                                    PYTHON_EXAMPLE_OUTPUT)
 from mutahunter.core.pilot.prompts.system import SYSTEM_PROMPT
 from mutahunter.core.pilot.prompts.user import USER_PROMPT
 

src/mutahunter/core/report.py

@@ -22,7 +22,7 @@ class MutantReport:
     def __init__(self, config) -> None:
         self.config = config
 
-    def generate_report(self, mutants: List[Mutant]) -> None:
+    def generate_report(self, mutants: List[Mutant], total_cost) -> None:
         """
         Generates a comprehensive mutation testing report.
 
@@ -32,10 +32,10 @@ def generate_report(self, mutants: List[Mutant]) -> None:
         mutants = [asdict(mutant) for mutant in mutants]
         self.save_report("logs/_latest/mutants.json", mutants)
         print(MUTAHUNTER_ASCII)
-        self.generate_mutant_report(mutants)
+        self.generate_mutant_report(mutants, total_cost)
         self.generate_mutant_report_detail(mutants)
 
-    def generate_mutant_report(self, mutants: List[Mutant]) -> None:
+    def generate_mutant_report(self, mutants: List[Mutant], total_cost) -> None:
         killed_mutants = [mutant for mutant in mutants if mutant["status"] == "KILLED"]
         survived_mutants = [
             mutant for mutant in mutants if mutant["status"] == "SURVIVED"
@@ -62,6 +62,7 @@ def generate_mutant_report(self, mutants: List[Mutant]) -> None:
         logger.info("🕒 Timeout Mutants: %d 🕒", len(timeout_mutants))
         logger.info("🔥 Compile Error Mutants: %d 🔥", len(compile_error_mutants))
         logger.info("🎯 Mutation Coverage: %s 🎯", total_mutation_coverage)
+        logger.info("💰 Expected Cost: $%.5f USD 💰", total_cost)
 
         mutation_coverage = {
             "total_mutants": len(mutants),
@@ -70,6 +71,7 @@ def generate_mutant_report(self, mutants: List[Mutant]) -> None:
             "timeout_mutants": len(timeout_mutants),
             "compile_error_mutants": len(compile_error_mutants),
             "mutation_coverage": total_mutation_coverage,
+            "expected_cost": total_cost,
         }
 
         self.save_report("logs/_latest/mutation_coverage.json", mutation_coverage)

src/mutahunter/core/router.py

@@ -1,6 +1,6 @@
 import time
 
-import litellm
+from litellm import completion, completion_cost, litellm, success_callback
 
 
 class LLMRouter:
@@ -10,6 +10,24 @@ def __init__(self, model: str, api_base: str = ""):
         """
         self.model = model
         self.api_base = api_base
+        self.total_cost = 0
+        litellm.success_callback = [self.track_cost_callback]
+
+        # track_cost_callback
+
+    def track_cost_callback(
+        self,
+        kwargs,  # kwargs to completion
+        completion_response,  # response from completion
+        start_time,
+        end_time,  # start/end time
+    ):
+        try:
+            response_cost = kwargs.get("response_cost", 0)
+            self.total_cost += response_cost
+            print("streaming response_cost", response_cost)
+        except:
+            pass
 
     def generate_response(
         self, prompt: dict, max_tokens: int = 4096, streaming: bool = False
@@ -25,6 +43,7 @@ def generate_response(
         Returns:
             tuple: Generated response, prompt tokens used, and completion tokens used.
         """
+        streaming = False
         self._validate_prompt(prompt)
         messages = self._build_messages(prompt)
         completion_params = self._build_completion_params(
@@ -86,7 +105,7 @@ def _stream_response(self, completion_params: dict) -> list:
         """
         response_chunks = []
         print("\nStreaming results from LLM model...")
-        response = litellm.completion(**completion_params)
+        response = completion(**completion_params)
         for chunk in response:
             print(chunk.choices[0].delta.content or "", end="", flush=True)
             response_chunks.append(chunk)
@@ -100,7 +119,7 @@ def _non_stream_response(self, completion_params: dict) -> tuple:
         """
         Get the non-streamed response from the LLM model.
         """
-        response = litellm.completion(**completion_params)
+        response = completion(**completion_params)
         content = response["choices"][0]["message"]["content"]
         prompt_tokens = int(response["usage"]["prompt_tokens"])
         completion_tokens = int(response["usage"]["completion_tokens"])

tests/test_hunter.py

@@ -326,7 +326,6 @@ def test_should_skip_file_exclude_files(mutant_hunter):
     assert mutant_hunter.should_skip_file("excluded_file.py") is True
 
 
-
 def test_process_test_result_compile_error(mutant_hunter):
     mutant = Mutant(
         id="1",
@@ -356,6 +355,7 @@ def test_process_test_result_timeout(mutant_hunter):
     assert mutant.status == "TIMEOUT"
     assert mutant.error_msg == "timeout error"
 
+
 @patch.object(MutantHunter, "prepare_mutant_file", return_value="")
 def test_process_mutant_compile_error(mock_prepare_mutant_file, mutant_hunter):
     mutant_data = {"mutant_code": "code", "type": "type", "description": "description"}

tests/test_report.py

@@ -1,13 +1,12 @@
 import json
+from dataclasses import asdict
 from unittest.mock import mock_open, patch
 
 import pytest
-from dataclasses import asdict
+
 from mutahunter.core.entities.mutant import Mutant
 from mutahunter.core.report import MutantReport
 
-from unittest.mock import patch
-
 
 @pytest.fixture
 def config():
@@ -135,14 +134,15 @@ def test_generate_mutant_report(mutants, config):
         patch.object(report, "save_report") as mock_save_report,
         patch("mutahunter.core.logger.logger.info") as mock_logger_info,
     ):
-        report.generate_mutant_report(mutants)
+        report.generate_mutant_report(mutants, 0.0)
 
         mock_logger_info.assert_any_call("🦠 Total Mutants: %d 🦠", len(mutants))
         mock_logger_info.assert_any_call("🛡️ Survived Mutants: %d 🛡️", 2)
         mock_logger_info.assert_any_call("🗡️ Killed Mutants: %d 🗡️", 2)
         mock_logger_info.assert_any_call("🕒 Timeout Mutants: %d 🕒", 0)
         mock_logger_info.assert_any_call("🔥 Compile Error Mutants: %d 🔥", 0)
         mock_logger_info.assert_any_call("🎯 Mutation Coverage: %s 🎯", "50.00%")
+        mock_logger_info.assert_any_call("💰 Expected Cost: $%.5f USD 💰", 0.0)
 
         mock_save_report.assert_called_once_with(
             "logs/_latest/mutation_coverage.json",
@@ -153,6 +153,7 @@ def test_generate_mutant_report(mutants, config):
                 "timeout_mutants": 0,
                 "compile_error_mutants": 0,
                 "mutation_coverage": "50.00%",
+                "expected_cost": 0.0,
             },
         )