Added templates for documentation and commenting

pyproject.toml

@@ -31,6 +31,10 @@ sphinx-autodoc-typehints = {version = ">=1.2.0"}
 sphinx-click = {version = ">=4.3.0"}
 myst-parser = {version = ">=0.18.1"}
 
+
+[tool.poetry.group.anthropic.dependencies]
+langchain-anthropic = "^0.0.1.post1"
+
 [tool.poetry.scripts]
 codergpt = "codergpt.cli:main"
 

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."""
@@ -24,14 +26,33 @@ def comment(self, code: str, filename: str, overwrite: bool = False, language: O
         :param code: The string containing the code to be commented.
         :param filename: The original filename of the code file.
         :param overwrite: A boolean indicating whether to overwrite the original file. Default is False.
+        :param language: Coding language of the file, defaults to None
         """
-        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 and sphinx docstrings in :param: format: \n{code}\n"
+                comments: \n{code}\n"
+                f"Use template {comment_template} as a style reference to render the docstrings."
+                "Everything in the template are placeholders."
+                "Return just the code block since all this will be saved in a file."
+            }
+        else:
+            invoke_params = {
+                "input": f"Rewrite and return this {language} code with\
+                    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,13 +7,14 @@
 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"]
 GPT_3_5_TURBO = "gpt-3.5-turbo"
 GPT_4 = "gpt-4"
 GPT_4_TURBO = "gpt-4-turbo-preview"
-CLAUDE = "claude"
+CLAUDE = "claude-2"
 GEMINI = "gemini-pro"
 
 ALL_MODELS = [
@@ -23,3 +24,24 @@
     CLAUDE,
     GEMINI,
 ]
+
+SPHINX_DOCUMENTATION_TEMPLATE = TEMPLATES_DIR / "sphinx_style_document.md"
+PYTHON_CODE_COMMENT_TEMPLATE = TEMPLATES_DIR / "python_code_comments.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,34 +20,47 @@ 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,
     ):
         """
         Document the contents of the code file by invoking the runnable chain.
 
+        :param filename: filename of the code file.
         :param code: The string containing the code to be documented.
-        :param function: The name of the function to document. Default is None.
-        :param classname: The name of the class to document
+        :param language: Coding language of the file, defaults to None
+        :param outfile: Destination filepath, defaults to None
         """
-        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 a style reference for documentation."
+                "Everything in the template are placeholders. Be as detailed as possible."
+                "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/explainer/explainer.py

@@ -27,11 +27,15 @@ def explain(
         :param classname: The name of the class to explain. Default is None.
         """
         if function:
-            response = self.chain.invoke({"input": f"Explain the following {language} code: \n\n```\n{code}\n```"})
+            response = self.chain.invoke(
+                {"input": f"Explain the function {function} in the following {language} code: \n\n```\n{code}\n```"}
+            )
             # Pretty print the response
             print(f"Explanation for '{function}':\n{response.content}")
         elif classname:
-            response = self.chain.invoke({"input": f"Explain the following {language} code: \n\n```\n{code}\n```"})
+            response = self.chain.invoke(
+                {"input": f"Explain the class {classname} in the following {language} code: \n\n```\n{code}\n```"}
+            )
             # Pretty print the response
             print(f"Explanation for '{classname}':\n{response.content}")
         else:

src/codergpt/main.py

@@ -5,13 +5,14 @@
 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.commenter.commenter import CodeCommenter
-from codergpt.constants import EXTENSION_MAP_FILE, GEMINI, GPT_4_TURBO, INSPECTION_HEADERS
+from codergpt.constants import CLAUDE, EXTENSION_MAP_FILE, GEMINI, GPT_4_TURBO, INSPECTION_HEADERS
 from codergpt.documenter.documenter import CodeDocumenter
 from codergpt.explainer.explainer import CodeExplainer
 from codergpt.optimizer.optimizer import CodeOptimizer
@@ -23,11 +24,16 @@ class CoderGPT:
 
     def __init__(self, model: str = GPT_4_TURBO):
         """Initialize the CoderGPT class."""
+        temp = 0.3
         if model is None or model.startswith("gpt-"):
-            self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=0.7, model=model)
-        # elif model == CLAUDE:
-        #     self.llm = ChatAnthropic()
-        #     print("Coming Soon!")
+            self.llm = ChatOpenAI(openai_api_key=os.environ.get("OPENAI_API_KEY"), temperature=temp, model=model)
+        elif model == CLAUDE:
+            self.llm = ChatAnthropicMessages(
+                model_name=model,
+                anthropic_api_key=os.environ.get("ANTHROPIC_API_KEY"),
+                temperature=temp,
+                max_tokens=2048,
+            )
         elif model == GEMINI:
             self.llm = ChatGoogleGenerativeAI(model=model, convert_system_message_to_human=True)
         else:
@@ -115,7 +121,7 @@ def explainer(self, path: Union[str, Path], function: str = None, classname=None
         """
         code_explainer = CodeExplainer(self.chain)
         code, language = self.get_code(filename=path, function_name=function, class_name=classname)
-        code_explainer.explain(code, language)
+        code_explainer.explain(code=code, function=function, classname=classname, language=language)
 
     def commenter(self, path: Union[str, Path], overwrite: bool = False):
         """
@@ -153,8 +159,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__":