diff --git a/README.md b/README.md
index f3c8204a..81a343c5 100644
--- a/README.md
+++ b/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
diff --git a/docs/ADD_DOCUMENTATION.md b/docs/ADD_DOCUMENTATION.md
new file mode 100644
index 00000000..b7204f80
--- /dev/null
+++ b/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:
+
+
+
+
+### 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 ...".
\ No newline at end of file
diff --git a/docs/TOOLS_GUIDE.md b/docs/TOOLS_GUIDE.md
index 352195bc..b3831961 100644
--- a/docs/TOOLS_GUIDE.md
+++ b/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.
\ No newline at end of file
diff --git a/pics/add_docs.png b/pics/add_docs.png
new file mode 100644
index 00000000..143ae49b
Binary files /dev/null and b/pics/add_docs.png differ
diff --git a/pics/add_docs_comment.png b/pics/add_docs_comment.png
new file mode 100644
index 00000000..cd40ff12
Binary files /dev/null and b/pics/add_docs_comment.png differ
diff --git a/pr_agent/agent/pr_agent.py b/pr_agent/agent/pr_agent.py
index 11db4dec..3d819af5 100644
--- a/pr_agent/agent/pr_agent.py
+++ b/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())
diff --git a/pr_agent/config_loader.py b/pr_agent/config_loader.py
index 47edfd97..184adb82 100644
--- a/pr_agent/config_loader.py
+++ b/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"
]]
)
diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml
index aa28564c..187ded56 100644
--- a/pr_agent/settings/configuration.toml
+++ b/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 = ""
diff --git a/pr_agent/settings/pr_add_docs.toml b/pr_agent/settings/pr_add_docs.toml
new file mode 100644
index 00000000..b552ec86
--- /dev/null
+++ b/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
+"""
diff --git a/pr_agent/tools/pr_add_docs.py b/pr_agent/tools/pr_add_docs.py
new file mode 100644
index 00000000..4cc9102a
--- /dev/null
+++ b/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"