Added code for bug finding + fixing

README.md

@@ -5,7 +5,7 @@ CoderGPT is a versatile command-line interface (CLI) designed to enhance coding
 # Model Providers Implemented
  - [x] OpenAI [`gpt-3.5-turbo`, `gpt-4`, `gpt-4-turbo-preview`(default)]
  - [x] Google [`gemini-pro`]
- - [x] Anthropic [`claude-2`]
+ - [x] Anthropic [`claude-2.1`]
 
 ## Prerequisites
 
@@ -46,7 +46,7 @@ codergpt [OPTIONS] COMMAND [ARGS]...
   - Available models:
     - OpenAI: [`gpt-3.5-turbo`, `gpt-4`, `gpt-4-turbo-preview`(default)]
     - Google: [`gemini-pro`]
-    - Anthropic[`claude-2`]
+    - Anthropic[`claude-2.1`]
 
 #### Commands
 

src/codergpt/bug_finder/__init__.py

@@ -0,0 +1,5 @@
+"""Bug-finder module for the package."""
+
+from .bug_finder import BugFinder
+
+__all__ = ["BugFinder"]

src/codergpt/bug_finder/bug_finder.py

@@ -0,0 +1,100 @@
+"""Bug-finder class for the package."""
+
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+from langchain_core.runnables.base import RunnableSerializable
+
+from codergpt.utils import extract_code_from_response
+
+
+class BugFinder:
+    """Bug-finder class for the package."""
+
+    def __init__(self, chain: RunnableSerializable[Dict, Any]):
+        """Initialize the BugFinder class."""
+        self.chain = chain
+
+    def find_bugs(
+        self, code: str, function: Optional[str] = None, classname: Optional[str] = None, language: Optional[str] = None
+    ):
+        """
+        Find bugs in the given code.
+
+        :param code: The code to find bugs in.
+        :param function: The name of the function to find bugs in. Default is None.
+        :param classname: The name of the class to find bugs in. Default is None.
+        :param language: The language of the code. Default is None.
+        """
+        if function:
+            response = self.chain.invoke(
+                {
+                    "input": f"Find and list all the bugs in the function {function}"
+                    f" in the following {language} code: \n\n```\n{code}\n```"
+                }
+            )
+            # Pretty print the response
+            print(f"Bugs found in '{function}':\n{response.content}")
+        elif classname:
+            response = self.chain.invoke(
+                {
+                    "input": f"Find and list all the bugs in the class {classname}"
+                    f" in the following {language} code: \n\n```\n{code}\n```"
+                }
+            )
+            # Pretty print the response
+            print(f"Bugs found in '{classname}':\n{response.content}")
+        else:
+            # Find bugs in full code
+            response = self.chain.invoke(
+                {"input": f"Find and list all the bugs in the following {language} code: \n\n```\n{code}\n```"}
+            )
+            # Pretty print the response
+            print(f"Bugs found in the code:\n{response.content}")
+
+    def fix_bugs(
+        self,
+        filename: Union[str, Path],
+        code: str,
+        function: Optional[str] = None,
+        classname: Optional[str] = None,
+        language: Optional[str] = None,
+        outfile: Optional[str] = None,
+    ) -> None:
+        """
+        Fix bugs in the given code.
+
+        :param code: The code to fix bugs in.
+        :param function: The name of the function to fix bugs in. Default is None.
+        :param classname: The name of the class to fix bugs
+        :param outfile:Path for output file with bug-fix code. Default is None.
+        """
+        if function:
+            response = self.chain.invoke(
+                {
+                    "input": f"List all the bug fixes if any and rewrite the function {function}"
+                    f" in the following {language} code: \n\n```\n{code}\n```"
+                }
+            )
+            # Pretty print the response
+            print(f"Fixed code for '{function}':\n{response.content}")
+            return response.content
+        elif classname:
+            response = self.chain.invoke(
+                {
+                    "input": f"List all the bug fixes if any and rewrite the class {classname}"
+                    f" in the following {language} code: \n\n```\n{code}\n```"
+                }
+            )
+            # Pretty print the response
+            print(f"Fixed code for '{classname}':\n{response.content}")
+            return response.content
+        else:
+            # Fix bugs in full code
+            response = self.chain.invoke(
+                {
+                    "input": f"List all the bug fixes if any and rewrite the following {language}"
+                    f" code: \n\n```\n{code}\n```"
+                }
+            )
+            return extract_code_from_response(language, response.content, filename, outfile)

src/codergpt/cli.py

@@ -206,5 +206,52 @@ def write_documentation(path: Union[str, Path], outfile: Union[str, Path] = None
         raise ValueError("The path provided is not a file.")
 
 
+@main.command("find-bugs")
+@path_argument
+@function_option
+@class_option
+def find_bugs_in_code(path: Union[str, Path], function: str, classname: str):
+    """
+    Write tests for the code file.
+
+    :param path: The path to the code file.
+    :param function: The name of the function to test. Default is None.
+    :param classname: The name of the class to test. Default is None.
+    """
+    # Ensure path is a string or Path object for consistency
+    if isinstance(path, str):
+        path = Path(path)
+
+    # Check if path is a file
+    if path.is_file():
+        coder.bug_finder(path=path, function=function, classname=classname)
+    else:
+        raise ValueError("The path provided is not a file.")
+
+
+@main.command("fix-bugs")
+@path_argument
+@function_option
+@class_option
+@output_option
+def fix_bugs_in_code(path: Union[str, Path], function: str, classname: str, outfile: Union[str, Path] = None):
+    """
+    Write tests for the code file.
+
+    :param path: The path to the code file.
+    :param function: The name of the function to test. Default is None.
+    :param classname: The name of the class to test. Default is None.
+    """
+    # Ensure path is a string or Path object for consistency
+    if isinstance(path, str):
+        path = Path(path)
+
+    # Check if path is a file
+    if path.is_file():
+        coder.bug_fixer(path=path, function=function, classname=classname, outfile=outfile)
+    else:
+        raise ValueError("The path provided is not a file.")
+
+
 if __name__ == "__main__":
     main()

src/codergpt/commenter/commenter.py

@@ -1,11 +1,13 @@
 """Commenter Module."""
 
 import os
+from pathlib import Path
 from typing import Any, Dict, Optional
 
 from langchain_core.runnables.base import RunnableSerializable
 
 from codergpt.constants import TEMPLATES
+from codergpt.utils import extract_code_from_response, get_language_from_extension
 
 
 class CodeCommenter:
@@ -29,13 +31,17 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
         :param language: Coding language of the file, defaults to None
         """
         comment_template = None
-        if language and language in TEMPLATES.keys():
-            # Check if "comment" key exists in the language template
-            if "comment" in TEMPLATES[language]:
-                # Get the path to the comment template
-                comment_template_path = TEMPLATES[language]["comment"]
-                with open(comment_template_path, "r") as comment_template_file:
-                    comment_template = comment_template_file.read()
+        if language:
+            if language in TEMPLATES.keys():
+                # Check if "comment" key exists in the language template
+                if "comment" in TEMPLATES[language]:
+                    # Get the path to the comment template
+                    comment_template_path = TEMPLATES[language]["comment"]
+                    with open(comment_template_path, "r") as comment_template_file:
+                        comment_template = comment_template_file.read()
+        else:
+            # Get the language from the file extension
+            language = get_language_from_extension(filename)
 
         if comment_template:
             invoke_params = {
@@ -54,15 +60,15 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
 
         response = self.chain.invoke(invoke_params)
 
-        # Extract the commented code from the response if necessary
-        commented_code = response.content
-
         new_filename = filename
         if not overwrite:
             # Create a new filename with the _updated suffix
             base, ext = os.path.splitext(filename)
-            new_filename = f"{base}_updated{ext}"
+            new_filename = f"{Path(filename).parent / base}_updated{ext}"
 
-        # Write the commented code to the new file
-        with open(new_filename, "w") as updated_file:
-            updated_file.write(commented_code)
+        if language:
+            return extract_code_from_response(language, response.content, filename, new_filename)
+        else:
+            # Write the commented code to the new file
+            with open(new_filename, "w") as updated_file:
+                updated_file.write(response.content)

src/codergpt/constants.py

@@ -14,7 +14,7 @@
 GPT_3_5_TURBO = "gpt-3.5-turbo"
 GPT_4 = "gpt-4"
 GPT_4_TURBO = "gpt-4-turbo-preview"
-CLAUDE = "claude-2"
+CLAUDE = "claude-2.1"
 GEMINI = "gemini-pro"
 
 ALL_MODELS = [

src/codergpt/main.py

@@ -4,19 +4,20 @@
 from pathlib import Path
 from typing import Optional, Union
 
-import yaml
 from langchain_anthropic import ChatAnthropicMessages
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_openai import ChatOpenAI
 from tabulate import tabulate
 
+from codergpt.bug_finder.bug_finder import BugFinder
 from codergpt.commenter.commenter import CodeCommenter
-from codergpt.constants import CLAUDE, EXTENSION_MAP_FILE, GEMINI, GPT_4_TURBO, INSPECTION_HEADERS
+from codergpt.constants import CLAUDE, GEMINI, GPT_4_TURBO, INSPECTION_HEADERS
 from codergpt.documenter.documenter import CodeDocumenter
 from codergpt.explainer.explainer import CodeExplainer
 from codergpt.optimizer.optimizer import CodeOptimizer
 from codergpt.test_writer.test_writer import CodeTester
+from codergpt.utils import get_language_from_extension
 
 
 class CoderGPT:
@@ -52,23 +53,20 @@ def inspect_package(self, path: Union[str, Path]):
         """
         print("Inspecting the code.")
 
-        with open(EXTENSION_MAP_FILE, "r") as file:
-            extension_to_language = yaml.safe_load(file)
-
         path = Path(path)
 
         file_language_list = []
         file_language_dict = {}
 
         if path.is_dir():
             for file in path.rglob("*.*"):
-                language = extension_to_language["language-map"].get(file.suffix)
+                language = get_language_from_extension(filename=file)
                 if language is not None:
                     file_language_list.append((str(file), language))
                     file_language_dict[str(file)] = language
 
         elif path.is_file():
-            language = extension_to_language["language-map"].get(path.suffix)
+            language = get_language_from_extension(filename=path)
             if language is not None:
                 file_language_list.append((str(path), language))
                 file_language_dict[str(path)] = language
@@ -166,6 +164,44 @@ def documenter(self, path: Union[str, Path], outfile: str = None):
         code, language = self.get_code(filename=path)
         code_documenter.document(filename=filename, code=code, language=language, outfile=outfile)
 
+    def bug_finder(self, path: Union[str, Path], function: Optional[str] = None, classname: Optional[str] = None):
+        """
+        Find bugs in the code file.
+
+        :param path: The path to the code file.
+        :param function: The name of the function to find bugs in. Default is None.
+        :param classname: The name of the class to find bugs in. Default is None.
+        """
+        if isinstance(path, str):
+            path = Path(path)
+        bug_finder = BugFinder(self.chain)
+        code, language = self.get_code(filename=path, function_name=function, class_name=classname)
+        bug_finder.find_bugs(code=code, function=function, classname=classname, language=language)
+
+    def bug_fixer(
+        self,
+        path: Union[str, Path],
+        function: Optional[str] = None,
+        classname: Optional[str] = None,
+        outfile: Optional[str] = None,
+    ):
+        """
+        Fix bugs in the code file.
+
+        :param path: The path to the code file.
+        :param function: The name of the function to fix bugs in. Default is None.
+        :param classname: The name of the class to fix bugs in. Default is None.
+        :param outfile: The path to the output file. Default is None.
+        """
+        if isinstance(path, str):
+            path = Path(path)
+        bug_finder = BugFinder(self.chain)
+        code, language = self.get_code(filename=path, function_name=function, class_name=classname)
+        filename = path.stem
+        bug_finder.fix_bugs(
+            filename=filename, code=code, function=function, classname=classname, language=language, outfile=outfile
+        )
+
 
 if __name__ == "__main__":
     coder = CoderGPT()

src/codergpt/utils.py

@@ -0,0 +1,55 @@
+"""Utility functions for the codergpt package."""
+
+import os
+import re
+from pathlib import Path
+from typing import Optional, Union
+
+import yaml
+
+from codergpt.constants import EXTENSION_MAP_FILE
+
+
+def extract_code_from_response(
+    language: str, response: str, filename: Union[str, Path], outfile: Optional[str] = None
+) -> str:
+    """
+    Generate code files based on LLM responses.
+
+    :param language: Code language.
+    :param response: LLM response.
+    :param filename: Source code file.
+    :param outfile: Destination filepath, defaults to None
+    """
+    base, ext = os.path.splitext(filename)
+    file_parent = Path(filename).parent
+
+    if not language:
+        get_language_from_extension(filename)
+
+    code_pattern_block = rf"```{language.lower()}(.*?)(?<=\n)```"
+    matches = re.findall(code_pattern_block, response, re.DOTALL)
+
+    if matches:
+        code_to_save = matches[0].strip()
+        if not outfile:
+            outfile = f"{file_parent/base}_updated{ext}"
+        with open(outfile, "w") as file:
+            file.write(code_to_save)
+        print(f"Fixed code saved in file: {outfile}")
+
+    print(response)
+    return response
+
+
+def get_language_from_extension(filename: Union[str, Path]) -> Optional[str]:
+    """
+    Get the language of a file from its extension.
+
+    :param filename: The filename to get the language for.
+    :return: The language of the file, if found.
+    """
+    with open(EXTENSION_MAP_FILE, "r") as file:
+        extension_to_language = yaml.safe_load(file)
+    language = extension_to_language["language-map"].get(Path(filename).suffix)
+    return language