Added Templates

src/codergpt/cli.py

@@ -44,6 +44,12 @@
     default="gpt-4",
     help="Model to use for code generation.",
 )
+prompt_option = click.option(
+    "-p",
+    "--prompt",
+    type=Union[str, Path],
+    help="Prompt to use for code generation.",
+)
 function_option = click.option("-f", "--function", help="Function name to explain or optimize.")
 class_option = click.option("-c", "--classname", help="Class name to explain or optimize.")
 overwrite_option = click.option(

src/codergpt/commenter/commenter.py

@@ -5,6 +5,8 @@
 
 from langchain_core.runnables.base import RunnableSerializable
 
+from codergpt.constants import TEMPLATES
+
 
 class CodeCommenter:
     """Code Explainer class that extracts and explains code from a given file."""
@@ -25,13 +27,30 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
         :param filename: The original filename of the code file.
         :param overwrite: A boolean indicating whether to overwrite the original file. Default is False.
         """
-        response = self.chain.invoke(
-            {
+        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 comment_template:
+            invoke_params = {
+                "input": f"Rewrite and return this {language} code with\
+                comments: \n{code}\n"
+                f"Use template {comment_template} as reference to render the code comments."
+                "Return just the code block since all this will be a file."
+            }
+        else:
+            invoke_params = {
                 "input": f"Rewrite and return this {language} code with\
-                comments and sphinx docstrings in :param: format: \n{code}\n"
+                    comments: \n{code}\n"
                 "Return just the code block since all this will be a file."
             }
-        )
+
+        response = self.chain.invoke(invoke_params)
 
         # Extract the commented code from the response if necessary
         commented_code = response.content

src/codergpt/constants.py

@@ -7,6 +7,7 @@
 DOCS_DIR = PROJECT_DIR / "docs"
 SRC = Path(__file__).resolve().parents[1]
 PACKAGE_DIR = SRC / "codergpt"
+TEMPLATES_DIR = PACKAGE_DIR / "templates"
 EXTENSION_MAP_FILE = PACKAGE_DIR / "extensions.yaml"
 LANGUAGE_MAP_KEY = "language-map"
 INSPECTION_HEADERS = ["File", "Language"]
@@ -23,3 +24,24 @@
     CLAUDE,
     GEMINI,
 ]
+
+SPHINX_DOCUMENTATION_TEMPLATE = TEMPLATES_DIR / "sphinx_style_document.md"
+PYTHON_CODE_COMMENT_TEMPLATE = TEMPLATES_DIR / "python_code_comment.md"
+
+"""
+Templates for different languages and commands.
+
+Follows format:
+    {
+        "language": {
+            "command": "path/to/comment/template",
+        }
+    }
+"""
+
+TEMPLATES = {
+    "Python": {
+        "comment": PYTHON_CODE_COMMENT_TEMPLATE,
+        "document": SPHINX_DOCUMENTATION_TEMPLATE,
+    }
+}

src/codergpt/documenter/documenter.py

@@ -1,11 +1,10 @@
 """Documenter module for codergpt."""
 
-from pathlib import Path
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Optional
 
 from langchain_core.runnables.base import RunnableSerializable
 
-from codergpt.constants import DOCS_DIR
+from codergpt.constants import DOCS_DIR, TEMPLATES
 
 
 class CodeDocumenter:
@@ -21,7 +20,9 @@ def __init__(self, chain: RunnableSerializable[Dict, Any]):
 
     def document(
         self,
-        filename: Union[str, Path],
+        filename: str,
+        code: str,
+        language: Optional[str] = None,
         outfile: Optional[str] = None,
     ):
         """
@@ -31,24 +32,33 @@ def document(
         :param function: The name of the function to document. Default is None.
         :param classname: The name of the class to document
         """
-        if isinstance(filename, str):
-            filename = Path(filename)
-        with open(filename, "r") as source_file:
-            source_code = source_file.read()
-        response = self.chain.invoke(
-            {
-                "input": f"Document the following code: \n\n```\n{source_code}\n```"
-                "This should be in reStructuredText (RST, ReST, or reST)"
-                " format that is Sphinx compatible. Return only the documentation content."
+        document_template = None
+        if language and language in TEMPLATES.keys():
+            # Check if "document" key exists in the language template
+            if "document" in TEMPLATES[language]:
+                # Get the path to the document template
+                document_template_path = TEMPLATES[language]["document"]
+                with open(document_template_path, "r") as document_template_file:
+                    document_template = document_template_file.read()
+        if document_template:
+            invoke_params = {
+                "input": f"Document the {language} code with the following: \n{code}\n"
+                f"Use template {document_template} as reference to render the code documentation."
+                "Everything in the template are placeholders. Return only the relevant documentation content."
             }
-        )
+        else:
+            invoke_params = {
+                "input": f"Document the {language} code with the following: \n{code}\n"
+                "Return only the documentation content."
+            }
+        response = self.chain.invoke(invoke_params)
 
         # Extract the commented code from the response if necessary
         documentation = response.content
         if outfile:
             destination_path = outfile
         else:
-            destination_path = DOCS_DIR / f"{filename.stem}.rst"
+            destination_path = DOCS_DIR / f"{filename}.rst"
         # Write the documentation to the new file
         with open(destination_path, "w") as updated_file:
             updated_file.write(documentation)

src/codergpt/main.py

@@ -24,7 +24,7 @@ class CoderGPT:
     def __init__(self, model: str = GPT_4_TURBO):
         """Initialize the CoderGPT class."""
         if model is None or model.startswith("gpt-"):
-            self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0, model=model)
+            self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0.3, model=model)
         # elif model == CLAUDE:
         #     self.llm = ChatAnthropic()
         #     print("Coming Soon!")
@@ -153,8 +153,12 @@ def documenter(self, path: Union[str, Path], outfile: str = None):
 
         :param path: The path to the code file.
         """
+        if isinstance(path, str):
+            path = Path(path)
         code_documenter = CodeDocumenter(self.chain)
-        code_documenter.document(filename=path, outfile=outfile)
+        filename = path.stem
+        code, language = self.get_code(filename=path)
+        code_documenter.document(filename=filename, code=code, language=language, outfile=outfile)
 
 
 if __name__ == "__main__":

src/codergpt/templates/python_code_comments.md

@@ -0,0 +1,146 @@
+```python
+#!/usr/bin/env python3
+"""
+Module docstring for the example Python script.
+
+This script provides an example of a Python module including various structures such as classes, functions,
+decorators, and more, with appropriate docstrings and comments.
+"""
+
+# Import statements
+import os
+import sys
+
+# Global variables
+GLOBAL_CONSTANT = 42
+
+def function_with_types_in_docstring(param1, param2):
+    """
+    Example of a function with types specified in the :param: format.
+
+    :param param1: The first parameter.
+    :type param1: int
+    :param param2: The second parameter.
+    :type param2: str
+    :return: The return value. True for success, False otherwise.
+    :rtype: bool
+    """
+    return True
+
+def function_with_pep484_type_annotations(param1: int, param2: str) -> bool:
+    """
+    Example of a function with PEP 484 type annotations.
+
+    :param param1: The first parameter.
+    :param param2: The second parameter.
+    :return: The return value. True for success, False otherwise.
+    """
+    return True
+
+class ExampleClass:
+    """
+    This is an example class with a simple docstring.
+
+    :ivar attr1: Description of `attr1`.
+    :vartype attr1: str
+    :ivar attr2: Description of `attr2`, defaults to 0.
+    :vartype attr2: int, optional
+    """
+
+    def __init__(self, attr1: str, attr2: int = 0):
+        """
+        The constructor for ExampleClass.
+
+        :param attr1: The first attribute.
+        :type attr1: str
+        :param attr2: The second attribute. Defaults to 0.
+        :type attr2: int, optional
+        """
+        self.attr1 = attr1
+        self.attr2 = attr2
+
+    def example_method(self, param1):
+        """
+        An example method.
+
+        :param param1: The first parameter of the method.
+        :type param1: str
+        :return: The return value. True for success, False otherwise.
+        :rtype: bool
+        """
+        return True
+
+    @staticmethod
+    def static_example_method(param1):
+        """
+        An example of a static method.
+
+        :param param1: The first parameter of the method.
+        :type param1: str
+        :return: The return value. True for success, False otherwise.
+        :rtype: bool
+        """
+        return True
+
+    @classmethod
+    def class_example_method(cls, param1):
+        """
+        An example of a class method.
+
+        :param param1: The first parameter of the method.
+        :type param1: str
+        :return: The return value. True for success, False otherwise.
+        :rtype: bool
+        """
+        return True
+
+    @property
+    def example_property(self):
+        """
+        An example of a property.
+
+        :return: The current value of `attr1`.
+        :rtype: str
+        """
+        return self.attr1
+
+    @example_property.setter
+    def example_property(self, value):
+        """
+        The setter for `example_property`.
+
+        :param value: The new value for `attr1`.
+        :type value: str
+        """
+        self.attr1 = value
+
+# Decorators
+def decorator_function(original_function):
+    """
+    A simple decorator function that wraps the input function.
+
+    :param original_function: The function to decorate.
+    :type original_function: callable
+    :return: The wrapped function.
+    :rtype: callable
+    """
+    def wrapper_function(*args, **kwargs):
+        # Do something before the original function is called.
+        result = original_function(*args, **kwargs)
+        # Do something after the original function is called.
+        return result
+    return wrapper_function
+
+@decorator_function
+def function_to_decorate():
+    """
+    Example function that will be decorated.
+    """
+    pass
+
+# Conditional execution
+if __name__ == "__main__":
+    # Code to execute when the script is run directly
+    example_instance = ExampleClass("Hello", 123)
+    print(example_instance.example_method("World"))
+```