Added docstrings and moved to total token count (#116)

* Added docstrings and moved to total token count.:

* Bumping version

* Format.

Author: EmbeddedDevops1

cover_agent/AICaller.py

@@ -75,12 +75,19 @@ def call_model(self, prompt: dict, max_tokens=4096):
 
         model_response = litellm.stream_chunk_builder(chunks, messages=messages)
 
-        if 'WANDB_API_KEY' in os.environ:
+        if "WANDB_API_KEY" in os.environ:
             root_span = Trace(
-                name="inference_"+datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
+                name="inference_"
+                + datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"),
                 kind="llm",  # kind can be "llm", "chain", "agent" or "tool
-                inputs={"user_prompt": prompt["user"], "system_prompt": prompt["system"]},
-                outputs={"model_response": model_response["choices"][0]["message"]["content"]})
+                inputs={
+                    "user_prompt": prompt["user"],
+                    "system_prompt": prompt["system"],
+                },
+                outputs={
+                    "model_response": model_response["choices"][0]["message"]["content"]
+                },
+            )
             root_span.log(name="inference")
 
         # Returns: Response, Prompt token count, and Response token count

cover_agent/CoverAgent.py

@@ -48,8 +48,8 @@ def _duplicate_test_file(self):
             self.args.test_file_output_path = self.args.test_file_path
 
     def run(self):
-        if 'WANDB_API_KEY' in os.environ:
-            wandb.login(key=os.environ['WANDB_API_KEY'])
+        if "WANDB_API_KEY" in os.environ:
+            wandb.login(key=os.environ["WANDB_API_KEY"])
             time_and_date = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
             run_name = f"{self.args.model}_" + time_and_date
             wandb.init(project="cover-agent", name=run_name)
@@ -78,9 +78,7 @@ def run(self):
 
             iteration_count += 1
 
-            if self.test_gen.current_coverage < (
-                self.test_gen.desired_coverage / 100
-            ):
+            if self.test_gen.current_coverage < (self.test_gen.desired_coverage / 100):
                 self.test_gen.run_coverage()
 
         if self.test_gen.current_coverage >= (self.test_gen.desired_coverage / 100):
@@ -96,9 +94,15 @@ def run(self):
             else:
                 self.logger.info(failure_message)
 
-        ReportGenerator.generate_report(
-            test_results_list, self.args.report_filepath
+        # Provide metric on total token usage
+        self.logger.info(
+            f"Total number of input tokens used for LLM model {self.test_gen.ai_caller.model}: {self.test_gen.total_input_token_count}"
         )
+        self.logger.info(
+            f"Total number of output tokens used for LLM model {self.test_gen.ai_caller.model}: {self.test_gen.total_output_token_count}"
+        )
+
+        ReportGenerator.generate_report(test_results_list, self.args.report_filepath)
 
-        if 'WANDB_API_KEY' in os.environ:
+        if "WANDB_API_KEY" in os.environ:
             wandb.finish()

cover_agent/CoverageProcessor.py

@@ -9,7 +9,10 @@
 
 class CoverageProcessor:
     def __init__(
-        self, file_path: str, src_file_path: str, coverage_type: Literal["cobertura", "lcov", "jacoco"]
+        self,
+        file_path: str,
+        src_file_path: str,
+        coverage_type: Literal["cobertura", "lcov", "jacoco"],
     ):
         """
         Initializes a CoverageProcessor object.
@@ -136,22 +139,26 @@ def parse_coverage_report_jacoco(self) -> Tuple[list, list, float]:
         lines_covered, lines_missed = [], []
 
         package_name, class_name = self.extract_package_and_class_java()
-        missed, covered = self.parse_missed_covered_lines_jacoco(package_name, class_name)
+        missed, covered = self.parse_missed_covered_lines_jacoco(
+            package_name, class_name
+        )
 
         total_lines = missed + covered
         coverage_percentage = (float(covered) / total_lines) if total_lines > 0 else 0
 
         return lines_covered, lines_missed, coverage_percentage
 
-    def parse_missed_covered_lines_jacoco(self, package_name: str, class_name: str) -> tuple[int, int]:
-        with open(self.file_path, 'r') as file:
+    def parse_missed_covered_lines_jacoco(
+        self, package_name: str, class_name: str
+    ) -> tuple[int, int]:
+        with open(self.file_path, "r") as file:
             reader = csv.DictReader(file)
             missed, covered = 0, 0
             for row in reader:
-                if row['PACKAGE'] == package_name and row['CLASS'] == class_name:
+                if row["PACKAGE"] == package_name and row["CLASS"] == class_name:
                     try:
-                        missed = int(row['LINE_MISSED'])
-                        covered = int(row['LINE_COVERED'])
+                        missed = int(row["LINE_MISSED"])
+                        covered = int(row["LINE_COVERED"])
                         break
                     except KeyError as e:
                         self.logger.error("Missing expected column in CSV: {e}")
@@ -160,13 +167,13 @@ def parse_missed_covered_lines_jacoco(self, package_name: str, class_name: str)
         return missed, covered
 
     def extract_package_and_class_java(self):
-        package_pattern = re.compile(r'^\s*package\s+([\w\.]+)\s*;.*$')
-        class_pattern = re.compile(r'^\s*public\s+class\s+(\w+).*')
+        package_pattern = re.compile(r"^\s*package\s+([\w\.]+)\s*;.*$")
+        class_pattern = re.compile(r"^\s*public\s+class\s+(\w+).*")
 
         package_name = ""
         class_name = ""
         try:
-            with open(self.src_file_path, 'r') as file:
+            with open(self.src_file_path, "r") as file:
                 for line in file:
                     if not package_name:  # Only match package if not already found
                         package_match = package_pattern.match(line)
@@ -184,4 +191,4 @@ def extract_package_and_class_java(self):
             self.logger.error(f"Error reading file {self.src_file_path}: {e}")
             raise
 
-        return package_name, class_name
+        return package_name, class_name