Merge pull request #337 from Codium-ai/hl/add_docs

Add code documentations in PR

README.md

@@ -28,7 +28,8 @@ CodiumAI `PR-Agent` is an open-source tool aiming to help developers review pull
 ‣ **Update Changelog ([`/update_changelog`](./docs/UPDATE_CHANGELOG.md))**: Automatically updating the CHANGELOG.md file with the PR changes.
 \
 ‣ **Find similar issue ([`/similar_issue`](./docs/SIMILAR_ISSUE.md))**: Automatically retrieves and presents similar issues
-
+\
+‣ **Add Documentation (`/add_docs`)**: Automatically adds [documentation](./docs/ADD_DOCUMENTATION.md) to un-documented functions/classes in the PR.
 
 See the [Usage Guide](./Usage.md) for instructions how to run the different tools from _CLI_, _online usage_, Or by _automatically triggering_ them when a new PR is opened.
 
@@ -111,6 +112,7 @@ See the [Release notes](./RELEASE_NOTES.md) for updates on the latest changes.
 |       | Reflect and Review                          |   :white_check_mark:    |   :white_check_mark:    |   :white_check_mark:        |          |   :white_check_mark:    |    :white_check_mark:    |
 |       | Update CHANGELOG.md                         |   :white_check_mark:    |   :white_check_mark:    |   :white_check_mark:        |          |          |       |
 |       | Find similar issue                          |   :white_check_mark:    |                         |                             |          |          |       |
+|       | Add Documentation                           |   :white_check_mark:    |   :white_check_mark:    |   :white_check_mark:        |   :white_check_mark:    |          |    :white_check_mark:    |
 |       |                                             |        |        |      |      |      |
 | USAGE | CLI                                         |   :white_check_mark:    |   :white_check_mark:    |   :white_check_mark:       |   :white_check_mark:    |   :white_check_mark:    |
 |       | App / webhook                               |   :white_check_mark:    |   :white_check_mark:    |           |          |          |
@@ -197,7 +199,7 @@ Here are some advantages of PR-Agent:
   - [x] Rank the PR (see [here](https://github.com/Codium-ai/pr-agent/pull/89))   
   - [ ] Enforcing CONTRIBUTING.md guidelines
   - [ ] Performance (are there any performance issues)
-  - [ ] Documentation (is the PR properly documented)
+  - [x] Documentation (is the PR properly documented)
   - [ ] ...
 
 ## Similar Projects

docs/ADD_DOCUMENTATION.md

@@ -0,0 +1,15 @@
+# Add Documentation Tool
+The `add_docs` tool scans the PR code changes, and automatically suggests documentation for the undocumented code components (functions, classes, etc.).
+
+It can be invoked manually by commenting on any PR:
+```
+/add_docs
+```
+For example:
+
+<kbd><img src=./../pics/add_docs_comment.png width="768"></kbd>
+<kbd><img src=./../pics/add_docs.png width="768"></kbd>
+
+### Configuration options
+ - `docs_style`: The exact style of the documentation (for python docstring). you can choose between: `google`, `numpy`, `sphinx`, `restructuredtext`, `plain`. Default is `sphinx`.
+ - `extra_instructions`: Optional extra instructions to the tool. For example: "focus on the changes in the file X. Ignore change in ...".

docs/TOOLS_GUIDE.md

@@ -5,5 +5,6 @@
 - [ASK](./ASK.md)
 - [SIMILAR_ISSUE](./SIMILAR_ISSUE.md)
 - [UPDATE CHANGELOG](./UPDATE_CHANGELOG.md)
+- [ADD DOCUMENTATION](./ADD_DOCUMENTATION.md)
 
 See the **[installation guide](/INSTALL.md)** for instructions on how to setup PR-Agent.

pr_agent/agent/pr_agent.py

@@ -6,6 +6,7 @@ import tempfile
 from pr_agent.algo.utils import update_settings_from_args
 from pr_agent.config_loader import get_settings
 from pr_agent.git_providers import get_git_provider
+from pr_agent.tools.pr_add_docs import PRAddDocs
 from pr_agent.tools.pr_code_suggestions import PRCodeSuggestions
 from pr_agent.tools.pr_description import PRDescription
 from pr_agent.tools.pr_information_from_user import PRInformationFromUser
@@ -32,6 +33,7 @@ command2class = {
     "config": PRConfig,
     "settings": PRConfig,
     "similar_issue": PRSimilarIssue,
+    "add_docs": PRAddDocs,
 }
 
 commands = list(command2class.keys())

pr_agent/config_loader.py

@@ -22,6 +22,7 @@ global_settings = Dynaconf(
         "settings/pr_sort_code_suggestions_prompts.toml",
         "settings/pr_information_from_user_prompts.toml",
         "settings/pr_update_changelog_prompts.toml",
+        "settings/pr_add_docs.toml",
         "settings_prod/.secrets.toml"
     ]]
 )

pr_agent/settings/configuration.toml

@@ -47,6 +47,10 @@ rank_extended_suggestions = true
 max_number_of_calls = 5
 final_clip_factor = 0.9
 
+[pr_add_docs] # /add_docs #
+extra_instructions = ""
+docs_style = "Sphinx Style" # "Google Style with Args, Returns, Attributes...etc", "Numpy Style", "Sphinx Style", "PEP257", "reStructuredText"
+
 [pr_update_changelog] # /update_changelog #
 push_changelog_changes=false
 extra_instructions = ""

pr_agent/settings/pr_add_docs.toml

@@ -0,0 +1,115 @@
+[pr_add_docs_prompt]
+system="""You are a language model called PR-Code-Documentation Agent, that specializes in generating documentation for code.
+Your task is to generate meaningfull {{ docs_for_language }} to a PR (the '+' lines).
+
+Example for a PR Diff input:
+'
+## src/file1.py
+
+@@ -12,3 +12,5 @@ def func1():
+__new hunk__
+12  code line that already existed in the file...
+13  code line that already existed in the file....
+14 +new code line1 added in the PR
+15 +new code line2 added in the PR
+16  code line that already existed in the file...
+__old hunk__
+ code line that already existed in the file...
+-code line that was removed in the PR
+ code line that already existed in the file...
+
+
+@@ ... @@ def func2():
+__new hunk__
+...
+__old hunk__
+...
+
+
+## src/file2.py
+...
+'
+
+Specific instructions:
+- Try to identify edited/added code components (classes/functions/methods...) that are undocumented. and generate {{ docs_for_language }} for each one.
+- If there are documented (any type of {{ language }} documentation) code components in the PR, Don't generate {{ docs_for_language }} for them.
+- Ignore code components that don't appear fully in the '__new hunk__' section. For example. you must see the component header and body,
+- Make sure the {{ docs_for_language }} starts and ends with standart {{ language }} {{ docs_for_language }} signs.
+- The {{ docs_for_language }} should be in standard format.
+- Provide the exact line number (inclusive) where the {{ docs_for_language }} should be added.
+
+
+{%- if extra_instructions %}
+
+Extra instructions from the user:
+{{ extra_instructions }}
+{%- endif %}
+
+You must use the following YAML schema to format your answer:
+```yaml
+Code Documentation:
+  type: array
+  uniqueItems: true
+  items:
+    relevant file:
+      type: string
+      description: the relevant file full path
+    relevant line:
+      type: integer
+      description: |-
+        The relevant line number from a '__new hunk__' section where the {{ docs_for_language }} should be added.
+    doc placement:
+      type: string
+      enum:
+        - before
+        - after
+      description: |-
+        The {{ docs_for_language }} placement relative to the relevant line (code component).
+    documentation:
+      type: string
+      description: |-
+        The {{ docs_for_language }} content. It should be complete, correctly formatted and indented, and without line numbers.
+```
+
+Example output:
+```yaml
+Code Documentation:
+-   relevant file: |-
+        src/file1.py
+    relevant lines: 12
+    doc placement: after
+    documentation: |-
+        \"\"\"
+        This is a python docstring for func1.
+        \"\"\"
+- ...
+...
+```
+
+
+Each YAML output MUST be after a newline, indented, with block scalar indicator ('|-').
+Don't repeat the prompt in the answer, and avoid outputting the 'type' and 'description' fields.
+"""
+
+user="""PR Info:
+
+Title: '{{ title }}'
+
+Branch: '{{ branch }}'
+
+Description: '{{description}}'
+
+{%- if language %}
+
+Main language: {{language}}
+{%- endif %}
+
+
+The PR Diff:
+```
+{{- diff|trim }}
+```
+
+Response (should be a valid YAML, and nothing else):
+```yaml
+"""

pr_agent/tools/pr_add_docs.py

@@ -0,0 +1,173 @@
+import copy
+import logging
+import textwrap
+from typing import List, Dict
+from jinja2 import Environment, StrictUndefined
+
+from pr_agent.algo.ai_handler import AiHandler
+from pr_agent.algo.pr_processing import get_pr_diff, retry_with_fallback_models, get_pr_multi_diffs
+from pr_agent.algo.token_handler import TokenHandler
+from pr_agent.algo.utils import load_yaml
+from pr_agent.config_loader import get_settings
+from pr_agent.git_providers import BitbucketProvider, get_git_provider
+from pr_agent.git_providers.git_provider import get_main_pr_language
+
+
+class PRAddDocs:
+    def __init__(self, pr_url: str, cli_mode=False, args: list = None):
+
+        self.git_provider = get_git_provider()(pr_url)
+        self.main_language = get_main_pr_language(
+            self.git_provider.get_languages(), self.git_provider.get_files()
+        )
+
+        self.ai_handler = AiHandler()
+        self.patches_diff = None
+        self.prediction = None
+        self.cli_mode = cli_mode
+        self.vars = {
+            "title": self.git_provider.pr.title,
+            "branch": self.git_provider.get_pr_branch(),
+            "description": self.git_provider.get_pr_description(),
+            "language": self.main_language,
+            "diff": "",  # empty diff for initial calculation
+            "extra_instructions": get_settings().pr_add_docs.extra_instructions,
+            "commit_messages_str": self.git_provider.get_commit_messages(),
+            'docs_for_language': get_docs_for_language(self.main_language,
+                                                       get_settings().pr_add_docs.docs_style),
+        }
+        self.token_handler = TokenHandler(self.git_provider.pr,
+                                          self.vars,
+                                          get_settings().pr_add_docs_prompt.system,
+                                          get_settings().pr_add_docs_prompt.user)
+
+    async def run(self):
+        try:
+            logging.info('Generating code Docs for PR...')
+            if get_settings().config.publish_output:
+                self.git_provider.publish_comment("Generating Documentation...", is_temporary=True)
+
+            logging.info('Preparing PR documentation...')
+            await retry_with_fallback_models(self._prepare_prediction)
+            data = self._prepare_pr_code_docs()
+            if (not data) or (not 'Code Documentation' in data):
+                logging.info('No code documentation found for PR.')
+                return
+
+            if get_settings().config.publish_output:
+                logging.info('Pushing PR documentation...')
+                self.git_provider.remove_initial_comment()
+                logging.info('Pushing inline code documentation...')
+                self.push_inline_docs(data)
+        except Exception as e:
+            logging.error(f"Failed to generate code documentation for PR, error: {e}")
+
+    async def _prepare_prediction(self, model: str):
+        logging.info('Getting PR diff...')
+        self.patches_diff = get_pr_diff(self.git_provider,
+                                        self.token_handler,
+                                        model,
+                                        add_line_numbers_to_hunks=True,
+                                        disable_extra_lines=True)
+
+        logging.info('Getting AI prediction...')
+        self.prediction = await self._get_prediction(model)
+
+    async def _get_prediction(self, model: str):
+        variables = copy.deepcopy(self.vars)
+        variables["diff"] = self.patches_diff  # update diff
+        environment = Environment(undefined=StrictUndefined)
+        system_prompt = environment.from_string(get_settings().pr_add_docs_prompt.system).render(variables)
+        user_prompt = environment.from_string(get_settings().pr_add_docs_prompt.user).render(variables)
+        if get_settings().config.verbosity_level >= 2:
+            logging.info(f"\nSystem prompt:\n{system_prompt}")
+            logging.info(f"\nUser prompt:\n{user_prompt}")
+        response, finish_reason = await self.ai_handler.chat_completion(model=model, temperature=0.2,
+                                                                        system=system_prompt, user=user_prompt)
+
+        return response
+
+    def _prepare_pr_code_docs(self) -> Dict:
+        docs = self.prediction.strip()
+        data = load_yaml(docs)
+        if isinstance(data, list):
+            data = {'Code Documentation': data}
+        return data
+
+    def push_inline_docs(self, data):
+        docs = []
+
+        if not data['Code Documentation']:
+            return self.git_provider.publish_comment('No code documentation found to improve this PR.')
+
+        for d in data['Code Documentation']:
+            try:
+                if get_settings().config.verbosity_level >= 2:
+                    logging.info(f"add_docs: {d}")
+                relevant_file = d['relevant file'].strip()
+                relevant_line = int(d['relevant line'])  # absolute position
+                documentation = d['documentation']
+                doc_placement = d['doc placement'].strip()
+                if documentation:
+                    new_code_snippet = self.dedent_code(relevant_file, relevant_line, documentation, doc_placement,
+                                                        add_original_line=True)
+
+                    body = f"**Suggestion:** Proposed documentation\n```suggestion\n" + new_code_snippet + "\n```"
+                    docs.append({'body': body, 'relevant_file': relevant_file,
+                                             'relevant_lines_start': relevant_line,
+                                             'relevant_lines_end': relevant_line})
+            except Exception:
+                if get_settings().config.verbosity_level >= 2:
+                    logging.info(f"Could not parse code docs: {d}")
+
+        is_successful = self.git_provider.publish_code_suggestions(docs)
+        if not is_successful:
+            logging.info("Failed to publish code docs, trying to publish each docs separately")
+            for doc_suggestion in docs:
+                self.git_provider.publish_code_suggestions([doc_suggestion])
+
+    def dedent_code(self, relevant_file, relevant_lines_start, new_code_snippet, doc_placement='after',
+                    add_original_line=False):
+        try:  # dedent code snippet
+            self.diff_files = self.git_provider.diff_files if self.git_provider.diff_files \
+                else self.git_provider.get_diff_files()
+            original_initial_line = None
+            for file in self.diff_files:
+                if file.filename.strip() == relevant_file:
+                    original_initial_line = file.head_file.splitlines()[relevant_lines_start - 1]
+                    break
+            if original_initial_line:
+                if doc_placement == 'after':
+                    line = file.head_file.splitlines()[relevant_lines_start]
+                else:
+                    line = original_initial_line
+                suggested_initial_line = new_code_snippet.splitlines()[0]
+                original_initial_spaces = len(line) - len(line.lstrip())
+                suggested_initial_spaces = len(suggested_initial_line) - len(suggested_initial_line.lstrip())
+                delta_spaces = original_initial_spaces - suggested_initial_spaces
+                if delta_spaces > 0:
+                    new_code_snippet = textwrap.indent(new_code_snippet, delta_spaces * " ").rstrip('\n')
+                if add_original_line:
+                    if doc_placement == 'after':
+                        new_code_snippet = original_initial_line + "\n" + new_code_snippet
+                    else:
+                        new_code_snippet = new_code_snippet.rstrip() + "\n" + original_initial_line
+        except Exception as e:
+            if get_settings().config.verbosity_level >= 2:
+                logging.info(f"Could not dedent code snippet for file {relevant_file}, error: {e}")
+
+        return new_code_snippet
+
+
+def get_docs_for_language(language, style):
+    language = language.lower()
+    if language == 'java':
+        return "Javadocs"
+    elif language in ['python', 'lisp', 'clojure']:
+        return f"Docstring ({style})"
+    elif language in ['javascript', 'typescript']:
+        return "JSdocs"
+    elif language == 'c++':
+        return "Doxygen"
+    else:
+        return "Docs"