Refactor AgentCompletion to take in function parameter. (#282)

* Updating AgentCompletion to require updated parameters from issue.

* Adding function parameters and cleaning up TOML.

* Refactored UnitTestGenerator with PromptBuilder parameters.

* Working UnitTestValidator/Generator.

* Adding more detailed failure info.

* Working refactor. Moved to Claude for testing.

* Working call to PromptBuilder without init. Deleting PromptBuilder test.

* Removed more of PromptBuilder.

* Adding defaults to ABC.

* Adding defaults to DefaultAgentCompeletion..

* Updated docstrings.

* Fixing path to test file.

* Fixed docstrings.

* Moving to tagged version of GPT-4o

* Adding more logging.

Author: EmbeddedDevops1

cover_agent/AICaller.py

@@ -32,7 +32,7 @@ def retry_wrapper():
 
 
 class AICaller:
-    def __init__(self, model: str, api_base: str = "", enable_retry=True):
+    def __init__(self, model: str, api_base: str = "", enable_retry=True, max_tokens=16384):
         """
         Initializes an instance of the AICaller class.
 
@@ -43,15 +43,15 @@ def __init__(self, model: str, api_base: str = "", enable_retry=True):
         self.model = model
         self.api_base = api_base
         self.enable_retry = enable_retry
+        self.max_tokens = max_tokens
 
     @conditional_retry  # You can access self.enable_retry here
-    def call_model(self, prompt: dict, max_tokens=16384, stream=True):
+    def call_model(self, prompt: dict, stream=True):
         """
         Call the language model with the provided prompt and retrieve the response.
 
         Parameters:
             prompt (dict): The prompt to be sent to the language model.
-            max_tokens (int, optional): The maximum number of tokens to generate in the response. Defaults to 16384.
             stream (bool, optional): Whether to stream the response or not. Defaults to True.
 
         Returns:
@@ -84,15 +84,15 @@ def call_model(self, prompt: dict, max_tokens=16384, stream=True):
             "messages": messages,
             "stream": stream,  # Use the stream parameter passed to the method
             "temperature": 0.2,
-            "max_tokens": max_tokens,
+            "max_tokens": self.max_tokens,
         }
 
         # Model-specific adjustments
         if self.model in ["o1-preview", "o1-mini", "o1", "o3-mini"]:
             stream = False  # o1 doesn't support streaming
             completion_params["temperature"] = 1
             completion_params["stream"] = False  # o1 doesn't support streaming
-            completion_params["max_completion_tokens"] = 2*max_tokens
+            completion_params["max_completion_tokens"] = 2*self.max_tokens
             # completion_params["reasoning_effort"] = "high"
             completion_params.pop("max_tokens", None)  # Remove 'max_tokens' if present
 

cover_agent/AgentCompletionABC.py

@@ -1,68 +1,192 @@
 from abc import ABC, abstractmethod
 from typing import Tuple
 
-
 class AgentCompletionABC(ABC):
-    """Abstract base class for AI-driven prompt handling."""
+    """
+    Abstract base class for AI-driven prompt handling. Each method accepts
+    specific input parameters (e.g. source/test content, logs, coverage data)
+    and returns a tuple containing the AI response, along with additional
+    metadata (e.g. token usage and the generated prompt).
+    """
 
     @abstractmethod
     def generate_tests(
         self,
-        failed_tests: str,
+        source_file_name: str,
+        max_tests: int,
+        source_file_numbered: str,
+        code_coverage_report: str,
         language: str,
-        test_framework: str,
-        coverage_report: str,
+        test_file: str,
+        test_file_name: str,
+        testing_framework: str,
+        additional_instructions_text: str = None,
+        additional_includes_section: str = None,
+        failed_tests_section: str = None,
     ) -> Tuple[str, int, int, str]:
         """
-        Generates additional unit tests to improve test coverage.
+        Generates additional unit tests to improve coverage or handle edge cases.
+
+        Args:
+            source_file_name (str): Name of the source file under test.
+            max_tests (int): Maximum number of test functions to propose.
+            source_file_numbered (str): The source code with line numbers.
+            code_coverage_report (str): Coverage details highlighting untested lines.
+            language (str): The programming language (e.g. "python", "java").
+            test_file (str): Contents of the existing test file.
+            test_file_name (str): The name/path of the test file.
+            testing_framework (str): The test framework in use (e.g. "pytest", "junit").
+            additional_instructions_text (str, optional): Extra instructions or context.
+            additional_includes_section (str, optional): Additional code or includes.
+            failed_tests_section (str, optional): Details of failed tests to consider.
 
         Returns:
-            Tuple[str, int, int, str]: AI-generated test cases, input token count, output token count, and generated prompt.
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated test suggestions (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
         """
         pass
 
     @abstractmethod
     def analyze_test_failure(
-        self, stderr: str, stdout: str, processed_test_file: str
+        self,
+        source_file_name: str,
+        source_file: str,
+        processed_test_file: str,
+        stdout: str,
+        stderr: str,
+        test_file_name: str,
     ) -> Tuple[str, int, int, str]:
         """
-        Analyzes a test failure and returns insights.
+        Analyzes the output of a failed test to determine possible causes and
+        recommended fixes.
+
+        Args:
+            source_file_name (str): Name of the source file being tested.
+            source_file (str): Raw content of the source file.
+            processed_test_file (str): Content of the failing test file (pre-processed).
+            stdout (str): Captured standard output from the test run.
+            stderr (str): Captured standard error from the test run.
+            test_file_name (str): Name/path of the failing test file.
 
         Returns:
-            Tuple[str, int, int, str]: AI-generated analysis, input token count, output token count, and generated prompt.
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated analysis or explanation (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
         """
         pass
 
     @abstractmethod
-    def analyze_test_insert_line(self, test_file: str) -> Tuple[str, int, int, str]:
+    def analyze_test_insert_line(
+        self,
+        language: str,
+        test_file_numbered: str,
+        test_file_name: str,
+        additional_instructions_text: str = None,
+    ) -> Tuple[str, int, int, str]:
         """
-        Determines where to insert new test cases.
+        Determines the correct placement for inserting new test cases into
+        an existing test file.
+
+        Args:
+            language (str): The programming language of the test file.
+            test_file_numbered (str): The test file content, labeled with line numbers.
+            test_file_name (str): Name/path of the test file.
+            additional_instructions_text (str, optional): Any extra instructions or context.
 
         Returns:
-            Tuple[str, int, int, str]: Suggested insertion point, input token count, output token count, and generated prompt.
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated suggestion or instructions (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
         """
         pass
 
     @abstractmethod
     def analyze_test_against_context(
-        self, test_code: str, context: str
+        self,
+        language: str,
+        test_file_content: str,
+        test_file_name_rel: str,
+        context_files_names_rel: str,
     ) -> Tuple[str, int, int, str]:
         """
-        Validates whether a test is appropriate for its corresponding source code.
+        Evaluates a test file against a set of related context files to identify:
+        1. If it is a unit test,
+        2. Which context file the test is primarily targeting.
+
+        Args:
+            language (str): The programming language of the test file.
+            test_file_content (str): Raw content of the test file under review.
+            test_file_name_rel (str): Relative path/name of the test file.
+            context_files_names_rel (str): One or more file names related to the context.
 
         Returns:
-            Tuple[str, int, int, str]: AI validation result, input token count, output token count, and generated prompt.
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated classification or analysis (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
         """
         pass
 
     @abstractmethod
     def analyze_suite_test_headers_indentation(
-        self, test_file: str
+        self,
+        language: str,
+        test_file_name: str,
+        test_file: str,
+    ) -> Tuple[str, int, int, str]:
+        """
+        Analyzes an existing test suite to determine its indentation style,
+        the number of existing tests, and potentially the testing framework.
+
+        Args:
+            language (str): The programming language of the test file.
+            test_file_name (str): Name/path of the test file.
+            test_file (str): Raw content of the test file.
+
+        Returns:
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated suite analysis (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
+        """
+        pass
+
+    @abstractmethod
+    def adapt_test_command_for_a_single_test_via_ai(
+        self,
+        test_file_relative_path: str,
+        test_command: str,
+        project_root_dir: str,
     ) -> Tuple[str, int, int, str]:
         """
-        Determines the indentation style used in test suite headers.
+        Adapts an existing test command line to run only a single test file,
+        preserving other relevant flags and arguments where possible.
+
+        Args:
+            test_file_relative_path (str): Path to the specific test file to be isolated.
+            test_command (str): The original command line used for running multiple tests.
+            project_root_dir (str): Root directory of the project.
 
         Returns:
-            Tuple[str, int, int, str]: Suggested indentation style, input token count, output token count, and generated prompt.
+            Tuple[str, int, int, str]:
+                A 4-element tuple containing:
+                - The AI-generated modified command line (string),
+                - The input token count (int),
+                - The output token count (int),
+                - The final constructed prompt (string).
         """
         pass

cover_agent/CoverAgent.py

@@ -7,14 +7,14 @@
 from typing import List
 
 from cover_agent.CustomLogger import CustomLogger
-from cover_agent.PromptBuilder import adapt_test_command_for_a_single_test_via_ai
+from cover_agent.PromptBuilder import PromptBuilder, adapt_test_command_for_a_single_test_via_ai
 from cover_agent.UnitTestGenerator import UnitTestGenerator
 from cover_agent.UnitTestValidator import UnitTestValidator
 from cover_agent.UnitTestDB import UnitTestDB
 from cover_agent.AICaller import AICaller
-from cover_agent.PromptBuilder import PromptBuilder
 from cover_agent.AgentCompletionABC import AgentCompletionABC
 from cover_agent.DefaultAgentCompletion import DefaultAgentCompletion
+import cover_agent.utils
 
 
 class CoverAgent:
@@ -42,18 +42,8 @@ def __init__(self, args, agent_completion: AgentCompletionABC = None):
             self.agent_completion = agent_completion
         else:
             # Default to using the DefaultAgentCompletion object with the PromptBuilder and AICaller
-            self.ai_caller = AICaller(model=args.model, api_base=args.api_base)
-            self.prompt_builder = PromptBuilder(
-                source_file_path=args.source_file_path,
-                test_file_path=args.test_file_output_path,
-                code_coverage_report="",
-                included_files=UnitTestGenerator.get_included_files(args.included_files, args.project_root),
-                additional_instructions=args.additional_instructions,
-                failed_test_runs="",
-                language="",
-                testing_framework="",
-                project_root=args.project_root,
-            )
+            self.ai_caller = AICaller(model=args.model, api_base=args.api_base, max_tokens=8192)
+            self.prompt_builder = PromptBuilder()
             self.agent_completion = DefaultAgentCompletion(
                 builder=self.prompt_builder, caller=self.ai_caller
             )
@@ -231,13 +221,18 @@ def run_test_gen(
             )
 
             # Loop through each new test and validate it
-            for generated_test in generated_tests_dict.get("new_tests", []):
-                # Validate the test and record the result
-                test_result = self.test_validator.validate_test(generated_test)
-
-                # Insert the test result into the database
-                test_result["prompt"] = self.test_gen.prompt["user"]
-                self.test_db.insert_attempt(test_result)
+            try:
+                for generated_test in generated_tests_dict.get("new_tests", []):
+                    # Validate the test and record the result
+                    test_result = self.test_validator.validate_test(generated_test)
+
+                    # Insert the test result into the database
+                    test_result["prompt"] = self.test_gen.prompt["user"]
+                    self.test_db.insert_attempt(test_result)
+            except AttributeError as e:
+                self.logger.error(
+                    f"Failed to validate the test {generated_test} within {generated_tests_dict}. Error: {e}"
+                )
 
             # Increment the iteration count
             iteration_count += 1