├── .gitignore ├── CHANGELOG.md ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── INSTALL.md ├── LICENSE ├── README.md ├── assets ├── meta_logo_white.png └── system.png ├── config └── example_config.yaml ├── data ├── raw_test_repo │ ├── README.md │ ├── __init__.py │ ├── example.py │ ├── inventory │ │ ├── __init__.py │ │ └── inventory_manager.py │ ├── models │ │ ├── __init__.py │ │ └── product.py │ ├── payment │ │ ├── __init__.py │ │ └── payment_processor.py │ └── vending_machine.py └── raw_test_repo_simple │ ├── helper.py │ ├── inner │ └── inner_functions.py │ ├── main.py │ ├── processor.py │ └── test_file.py ├── eval_completeness.py ├── generate_docstrings.py ├── output └── dependency_graphs │ └── raw_test_repo_dependency_graph.json ├── run_web_ui.py ├── setup.py ├── src ├── DocstringGenerator.egg-info │ ├── PKG-INFO │ ├── SOURCES.txt │ ├── dependency_links.txt │ ├── requires.txt │ └── top_level.txt ├── __init__.py ├── __pycache__ │ ├── __init__.cpython-310.pyc │ ├── __init__.cpython-311.pyc │ ├── __init__.cpython-312.pyc │ ├── __init__.cpython-38.pyc │ └── __init__.cpython-39.pyc ├── agent │ ├── README.md │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-311.pyc │ │ ├── base.cpython-310.pyc │ │ ├── base.cpython-311.pyc │ │ ├── docstring_workflow.cpython-310.pyc │ │ ├── orchestrator.cpython-310.pyc │ │ ├── orchestrator.cpython-311.pyc │ │ ├── reader.cpython-310.pyc │ │ ├── reader.cpython-311.pyc │ │ ├── searcher.cpython-310.pyc │ │ ├── searcher.cpython-311.pyc │ │ ├── verifier.cpython-310.pyc │ │ ├── verifier.cpython-311.pyc │ │ ├── workflow.cpython-310.pyc │ │ ├── writer.cpython-310.pyc │ │ └── writer.cpython-311.pyc │ ├── base.py │ ├── llm │ │ ├── __init__.py │ │ ├── __pycache__ │ │ │ ├── __init__.cpython-310.pyc │ │ │ ├── __init__.cpython-311.pyc │ │ │ ├── base.cpython-310.pyc │ │ │ ├── base.cpython-311.pyc │ │ │ ├── claude_llm.cpython-310.pyc │ │ │ ├── claude_llm.cpython-311.pyc │ │ │ ├── factory.cpython-310.pyc │ │ │ ├── factory.cpython-311.pyc │ │ │ ├── gemini_llm.cpython-310.pyc │ │ │ ├── gemini_llm.cpython-311.pyc │ │ │ ├── huggingface_llm.cpython-310.pyc │ │ │ ├── huggingface_llm.cpython-311.pyc │ │ │ ├── openai_llm.cpython-310.pyc │ │ │ ├── openai_llm.cpython-311.pyc │ │ │ ├── rate_limiter.cpython-310.pyc │ │ │ └── rate_limiter.cpython-311.pyc │ │ ├── base.py │ │ ├── claude_llm.py │ │ ├── factory.py │ │ ├── gemini_llm.py │ │ ├── huggingface_llm.py │ │ ├── openai_llm.py │ │ └── rate_limiter.py │ ├── orchestrator.py │ ├── reader.py │ ├── searcher.py │ ├── tool │ │ ├── README.md │ │ ├── __pycache__ │ │ │ ├── ast.cpython-310.pyc │ │ │ ├── ast.cpython-311.pyc │ │ │ ├── ast_analyzer.cpython-310.pyc │ │ │ ├── internal_traverse.cpython-310.pyc │ │ │ ├── internal_traverse.cpython-311.pyc │ │ │ ├── perplexity_api.cpython-310.pyc │ │ │ └── perplexity_api.cpython-311.pyc │ │ ├── ast.py │ │ ├── internal_traverse.py │ │ └── perplexity_api.py │ ├── verifier.py │ ├── workflow.py │ └── writer.py ├── analyze_helpfulness_significance.py ├── data │ └── parse │ │ ├── data_process.py │ │ ├── downloader.py │ │ └── repo_tree.py ├── dependency_analyzer │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-311.pyc │ │ ├── __init__.cpython-312.pyc │ │ ├── analyzer.cpython-310.pyc │ │ ├── ast_parser.cpython-310.pyc │ │ ├── ast_parser.cpython-311.pyc │ │ ├── ast_parser.cpython-312.pyc │ │ ├── graph.cpython-310.pyc │ │ ├── topo_sort.cpython-310.pyc │ │ ├── topo_sort.cpython-311.pyc │ │ ├── topo_sort.cpython-312.pyc │ │ └── visualizer.cpython-310.pyc │ ├── ast_parser.py │ └── topo_sort.py ├── evaluate_helpfulness.py ├── evaluator │ ├── README.md │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-38.pyc │ │ ├── __init__.cpython-39.pyc │ │ ├── base.cpython-310.pyc │ │ ├── base.cpython-38.pyc │ │ ├── base.cpython-39.pyc │ │ ├── completeness.cpython-310.pyc │ │ ├── completeness.cpython-38.pyc │ │ ├── completeness.cpython-39.pyc │ │ ├── evaluation_common.cpython-310.pyc │ │ ├── helpfulness_arguments.cpython-310.pyc │ │ ├── helpfulness_attributes.cpython-310.pyc │ │ ├── helpfulness_description.cpython-310.pyc │ │ ├── helpfulness_evaluator.cpython-310.pyc │ │ ├── helpfulness_evaluator_ablation.cpython-310.pyc │ │ ├── helpfulness_examples.cpython-310.pyc │ │ ├── helpfulness_parameters.cpython-310.pyc │ │ ├── helpfulness_summary.cpython-310.pyc │ │ └── segment.cpython-39.pyc │ ├── base.py │ ├── completeness.py │ ├── evaluation_common.py │ ├── helper │ │ ├── __pycache__ │ │ │ └── context_finder.cpython-310.pyc │ │ └── context_finder.py │ ├── helpfulness_attributes.py │ ├── helpfulness_description.py │ ├── helpfulness_evaluator.py │ ├── helpfulness_evaluator_ablation.py │ ├── helpfulness_examples.py │ ├── helpfulness_parameters.py │ ├── helpfulness_summary.py │ ├── segment.py │ └── truthfulness.py ├── visualizer │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── __init__.cpython-311.pyc │ │ ├── graph_visualizer.cpython-310.pyc │ │ ├── progress.cpython-310.pyc │ │ ├── progress.cpython-311.pyc │ │ ├── status.cpython-310.pyc │ │ ├── status.cpython-311.pyc │ │ └── web_bridge.cpython-310.pyc │ ├── progress.py │ ├── status.py │ └── web_bridge.py ├── web │ ├── README.md │ ├── __init__.py │ ├── __pycache__ │ │ ├── __init__.cpython-310.pyc │ │ ├── app.cpython-310.pyc │ │ ├── config_handler.cpython-310.pyc │ │ ├── process_handler.cpython-310.pyc │ │ ├── run.cpython-310.pyc │ │ └── visualization_handler.cpython-310.pyc │ ├── app.py │ ├── config_handler.py │ ├── process_handler.py │ ├── run.py │ ├── static │ │ ├── assets │ │ │ └── meta_logo_white.png │ │ ├── css │ │ │ └── style.css │ │ └── js │ │ │ ├── completeness.js │ │ │ ├── config.js │ │ │ ├── log-handler.js │ │ │ ├── main.js │ │ │ ├── repo-structure.js │ │ │ └── status-visualizer.js │ ├── templates │ │ └── index.html │ └── visualization_handler.py └── web_eval │ ├── README.md │ ├── __pycache__ │ ├── app.cpython-310.pyc │ ├── app.cpython-38.pyc │ ├── helpers.cpython-310.pyc │ ├── helpers.cpython-38.pyc │ └── helpers.cpython-39.pyc │ ├── app.py │ ├── helpers.py │ ├── requirements.txt │ ├── start_server.sh │ ├── static │ ├── assets │ │ └── meta_logo_white.png │ └── css │ │ └── style.css │ ├── templates │ ├── index.html │ └── results.html │ └── test_docstring_parser.py └── tool ├── remove_docstrings.py ├── remove_docstrings.sh └── serve_local_llm.sh /.gitignore: -------------------------------------------------------------------------------- 1 | config/agent_config.yaml 2 | tool/add_header.sh -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | 0.0.1 (April 17, 2025) 2 | 3 | ### First Version 4 | 5 | Include web UI, CLI for DocAgent. -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | # Code of Conduct 2 | 3 | ## Our Pledge 4 | 5 | In the interest of fostering an open and welcoming environment, we as 6 | contributors and maintainers pledge to make participation in our project and 7 | our community a harassment-free experience for everyone, regardless of age, body 8 | size, disability, ethnicity, sex characteristics, gender identity and expression, 9 | level of experience, education, socio-economic status, nationality, personal 10 | appearance, race, religion, or sexual identity and orientation. 11 | 12 | ## Our Standards 13 | 14 | Examples of behavior that contributes to creating a positive environment 15 | include: 16 | 17 | * Using welcoming and inclusive language 18 | * Being respectful of differing viewpoints and experiences 19 | * Gracefully accepting constructive criticism 20 | * Focusing on what is best for the community 21 | * Showing empathy towards other community members 22 | 23 | Examples of unacceptable behavior by participants include: 24 | 25 | * The use of sexualized language or imagery and unwelcome sexual attention or 26 | advances 27 | * Trolling, insulting/derogatory comments, and personal or political attacks 28 | * Public or private harassment 29 | * Publishing others' private information, such as a physical or electronic 30 | address, without explicit permission 31 | * Other conduct which could reasonably be considered inappropriate in a 32 | professional setting 33 | 34 | ## Our Responsibilities 35 | 36 | Project maintainers are responsible for clarifying the standards of acceptable 37 | behavior and are expected to take appropriate and fair corrective action in 38 | response to any instances of unacceptable behavior. 39 | 40 | Project maintainers have the right and responsibility to remove, edit, or 41 | reject comments, commits, code, wiki edits, issues, and other contributions 42 | that are not aligned to this Code of Conduct, or to ban temporarily or 43 | permanently any contributor for other behaviors that they deem inappropriate, 44 | threatening, offensive, or harmful. 45 | 46 | ## Scope 47 | 48 | This Code of Conduct applies within all project spaces, and it also applies when 49 | an individual is representing the project or its community in public spaces. 50 | Examples of representing a project or community include using an official 51 | project e-mail address, posting via an official social media account, or acting 52 | as an appointed representative at an online or offline event. Representation of 53 | a project may be further defined and clarified by project maintainers. 54 | 55 | This Code of Conduct also applies outside the project spaces when there is a 56 | reasonable belief that an individual's behavior may have a negative impact on 57 | the project or its community. 58 | 59 | ## Enforcement 60 | 61 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 62 | reported by contacting the project team at . All 63 | complaints will be reviewed and investigated and will result in a response that 64 | is deemed necessary and appropriate to the circumstances. The project team is 65 | obligated to maintain confidentiality with regard to the reporter of an incident. 66 | Further details of specific enforcement policies may be posted separately. 67 | 68 | Project maintainers who do not follow or enforce the Code of Conduct in good 69 | faith may face temporary or permanent repercussions as determined by other 70 | members of the project's leadership. 71 | 72 | ## Attribution 73 | 74 | This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4, 75 | available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html 76 | 77 | [homepage]: https://www.contributor-covenant.org 78 | 79 | For answers to common questions about this code of conduct, see 80 | https://www.contributor-covenant.org/faq -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing to DocAgent 2 | We want to make contributing to this project as easy and transparent as 3 | possible. 4 | 5 | ## Pull Requests 6 | We actively welcome your pull requests. 7 | 8 | 1. Fork the repo and create your branch from `main`. 9 | 2. If you've added code that should be tested, add tests. 10 | 3. If you've changed APIs, update the documentation. 11 | 4. Ensure the test suite passes. 12 | 5. Make sure your code lints. 13 | 6. If you haven't already, complete the Contributor License Agreement ("CLA"). 14 | 15 | ## Contributor License Agreement ("CLA") 16 | In order to accept your pull request, we need you to submit a CLA. You only need 17 | to do this once to work on any of Meta's open source projects. 18 | 19 | Complete your CLA here: 20 | 21 | ## Issues 22 | We use GitHub issues to track public bugs. Please ensure your description is 23 | clear and has sufficient instructions to be able to reproduce the issue. 24 | 25 | Meta has a [bounty program](https://bugbounty.meta.com/) for the safe 26 | disclosure of security bugs. In those cases, please go through the process 27 | outlined on that page and do not file a public issue. 28 | 29 | ## Coding Style 30 | * 2 spaces for indentation rather than tabs 31 | * 80 character line length 32 | * Use [Black](https://github.com/psf/black) for code formatting. 33 | * Use [Flake8](https://flake8.pycqa.org/en/latest/) for linting. 34 | * Follow [PEP 8](https://www.python.org/dev/peps/pep-0008/) style guidelines. 35 | * Use snake_case for variable and function names. 36 | * Use PascalCase for class names. 37 | * Write docstrings for all public modules, classes, functions, and methods using Google style. 38 | * Use type hints for function signatures. 39 | * Keep imports organized: standard library first, then third-party libraries, then local application/library specific imports, each group separated by a blank line. Use [isort](https://pycqa.github.io/isort/) to automate this. 40 | 41 | ## License 42 | By contributing to DocAgent, you agree that your contributions will be licensed 43 | under the LICENSE file in the root directory of this source tree. -------------------------------------------------------------------------------- /INSTALL.md: -------------------------------------------------------------------------------- 1 | # Installation Guide 2 | 3 | This guide details how to set up the environment for DocAgent. 4 | 5 | ## Option 1: Installation with pip (Recommended) 6 | 7 | ### Basic Installation 8 | To install the basic package with core dependencies: 9 | 10 | ```bash 11 | # For all dependencies 12 | pip install -e ".[all]" 13 | ``` 14 | 15 | 16 | 17 | ## Development Setup 18 | 19 | For development, we recommend installing in editable mode with dev dependencies: 20 | 21 | ```bash 22 | # Install the package in editable mode with dev dependencies 23 | pip install -e ".[dev]" 24 | 25 | # Run tests 26 | pytest 27 | ``` 28 | 29 | ## Troubleshooting 30 | 31 | ### GraphViz Dependencies 32 | 33 | For visualization components, you may need to install system-level dependencies for GraphViz: 34 | 35 | ```bash 36 | # Ubuntu/Debian 37 | sudo apt-get install graphviz graphviz-dev 38 | 39 | # CentOS/RHEL 40 | sudo yum install graphviz graphviz-devel 41 | 42 | # macOS 43 | brew install graphviz 44 | ``` 45 | 46 | ### CUDA Support 47 | 48 | If you're using CUDA for accelerated processing, ensure you have the correct CUDA toolkit installed that matches your PyTorch version. -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) Meta Platforms, Inc. and affiliates. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # DocAgent: Agentic Hierarchical Docstring Generation System 2 | 3 |

4 | Meta Logo 5 |

6 | 7 | DocAgent is a system designed to generate high-quality, context-aware docstrings for Python codebases using a multi-agent approach and hierarchical processing. 8 | 9 | ## Citation 10 | 11 | If you use DocAgent in your research, please cite our paper: 12 | 13 | ```bibtex 14 | @misc{yang2025docagent, 15 | title={DocAgent: A Multi-Agent System for Automated Code Documentation Generation}, 16 | author={Dayu Yang and Antoine Simoulin and Xin Qian and Xiaoyi Liu and Yuwei Cao and Zhaopu Teng and Grey Yang}, 17 | year={2025}, 18 | eprint={2504.08725}, 19 | archivePrefix={arXiv}, 20 | primaryClass={cs.SE} 21 | } 22 | ``` 23 | 24 | You can find the paper on arXiv: [https://arxiv.org/abs/2504.08725](https://arxiv.org/abs/2504.08725) 25 | 26 | ## Table of Contents 27 | 28 | - [Motivation](#motivation) 29 | - [Methodology](#methodology) 30 | - [Installation](#installation) 31 | - [Components](#components) 32 | - [Configuration](#configuration) 33 | - [Usage](#usage) 34 | - [Running the Evaluation System](#running-the-evaluation-system) 35 | - [Optional: Using a Local LLM](#optional-using-a-local-llm) 36 | 37 | ## Motivation 38 | 39 | High-quality docstrings are crucial for code readability, usability, and maintainability, especially in large repositories. They should explain the purpose, parameters, returns, exceptions, and usage within the broader context. Current LLMs often struggle with this, producing superficial or redundant comments and failing to capture essential context or rationale. DocAgent aims to address these limitations by generating informative, concise, and contextually aware docstrings. 40 | 41 | ## Methodology 42 | 43 | DocAgent employs two key strategies: 44 | 45 | 1. **Hierarchical Traversal**: Processes code components by analyzing dependencies, starting with files having fewer dependencies. This builds a documented foundation before tackling more complex code, addressing the challenge of documenting context that itself lacks documentation. 46 | 2. **Agentic System**: Utilizes a team of specialized agents (`Reader`, `Searcher`, `Writer`, `Verifier`) coordinated by an `Orchestrator`. This system gathers context (internal and external), drafts docstrings according to standards, and verifies their quality in an iterative process. 47 | 48 | System Overview 49 | 50 | For more details on the agentic framework, see the [Agent Component README](./src/agent/README.md). 51 | 52 | ## Installation 53 | 54 | 1. Clone the repository: 55 | ```bash 56 | git clone 57 | cd DocAgent 58 | ``` 59 | 2. Install the necessary dependencies. It's recommended to use a virtual environment: 60 | ```bash 61 | python -m venv venv 62 | source venv/bin/activate # if you use venv, you can also use conda 63 | pip install -e . 64 | ``` 65 | *Note: For optional features like development tools, web UI components, or specific hardware support (e.g., CUDA), refer to the comments in `setup.py` and install extras as needed (e.g., `pip install -e ".[dev,web]"`).* 66 | 67 | ## Components 68 | 69 | DocAgent is composed of several key parts: 70 | 71 | - **[Core Agent Framework](./src/agent/README.md)**: Implements the multi-agent system (Reader, Searcher, Writer, Verifier, Orchestrator) responsible for the generation logic. 72 | - **[Docstring Evaluator](./src/evaluator/README.md)**: Provides tools for evaluating docstring quality, primarily focusing on completeness based on static code analysis (AST). *Note: Evaluation is run separately, see its README.* 73 | - **[Generation Web UI](./src/web/README.md)**: A web interface for configuring, running, and monitoring the docstring *generation* process in real-time. 74 | 75 | ## Configuration 76 | 77 | Before running DocAgent, you **must** create a configuration file named `config/agent_config.yaml`. This file specifies crucial parameters for the agents, such as the LLM endpoints, API keys (if required), model names, and generation settings. 78 | 79 | 1. **Copy the Example**: An example configuration file is provided at `config/example_config.yaml`. Copy this file to `config/agent_config.yaml`: 80 | ```bash 81 | cp config/example_config.yaml config/agent_config.yaml 82 | ``` 83 | 2. **Edit the Configuration**: Open `config/agent_config.yaml` in a text editor and modify the settings according to your environment and requirements. Pay close attention to the LLM provider, model selection, and any necessary API credentials. 84 | 85 | ## Usage 86 | 87 | You can run the docstring generation process using either the command line or the web UI. 88 | 89 | **1. Command Line Interface (CLI)** 90 | 91 | This is the primary method for running the generation process directly. 92 | 93 | ```bash 94 | # Example: Run on a test repo (remove existing docstrings first if desired) 95 | ./test/tool/remove_docstrings.sh data/raw_test_repo 96 | python generate_docstrings.py --repo-path data/raw_test_repo 97 | ``` 98 | Use `python generate_docstrings.py --help` to see available options, such as specifying different configurations or test modes. 99 | 100 | **2. Generation Web UI** 101 | 102 | The web UI provides a graphical interface to configure, run, and monitor the process. 103 | 104 | - Note that when input repo path, always put complete absolute path. 105 | 106 | ```bash 107 | # Launch the web UI server 108 | python run_web_ui.py --host 0.0.0.0 --port 5000 109 | ``` 110 | 111 | Then, access the UI in your web browser, typically at `http://localhost:5000`. If running the server remotely, you might need to set up SSH tunneling (see instructions below or the [Web UI README](./src/web/README.md)). 112 | 113 | *Basic SSH Tunneling (if running server remotely):* 114 | ```bash 115 | # In your local terminal 116 | ssh -L 5000:localhost:5000 @ 117 | # Then access http://localhost:5000 in your local browser 118 | ``` 119 | 120 | ## Running the Evaluation System 121 | 122 | DocAgent includes a separate web-based interface for evaluating the quality of generated docstrings. 123 | 124 | **1. Running Locally** 125 | 126 | To run the evaluation system on your local machine: 127 | 128 | ```bash 129 | python src/web_eval/app.py 130 | ``` 131 | 132 | Then, access the evaluation UI in your web browser at `http://localhost:5001`. 133 | 134 | **2. Running on a Remote Server** 135 | 136 | To run the evaluation system on a remote server: 137 | 138 | ```bash 139 | python src/web_eval/app.py --host 0.0.0.0 --port 5001 140 | ``` 141 | 142 | Then, set up SSH tunneling to access the remote server from your local machine: 143 | 144 | ```bash 145 | ssh -L 5001:localhost:5001 @ 146 | ``` 147 | 148 | Once the tunnel is established, access the evaluation UI in your local web browser at `http://localhost:5001`. 149 | 150 | ## Optional: Using a Local LLM 151 | 152 | If you prefer to use a local LLM (e.g., one hosted via Hugging Face), you can configure DocAgent to interact with it via an API endpoint. 153 | 154 | 1. **Serve the Local LLM**: Use a tool like `vllm` to serve your model. A convenience script is provided: 155 | ```bash 156 | # Ensure vllm is installed: pip install vllm 157 | bash tool/serve_local_llm.sh 158 | ``` 159 | This script will likely start an OpenAI-compatible API server (check the script details). Note the URL where the model is served (e.g., `http://localhost:8000/v1`). 160 | 161 | 2. **Configure DocAgent**: Update your `config/agent_config.yaml` to point to the local LLM API endpoint. You'll typically need to set: 162 | - The `provider` to `openai` (if using an OpenAI-compatible server like vllm's default). 163 | - The `api_base` or equivalent URL parameter to your local server address (e.g., `http://localhost:8000/v1`). 164 | - The `model_name` to the appropriate identifier for your local model. 165 | - Set the `api_key` to `None` or an empty string if no key is required by your local server. 166 | 167 | 3. **Run DocAgent**: Run the generation process as usual (CLI or Web UI). DocAgent will now send requests to your local LLM. 168 | 169 | ## License 170 | 171 | This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details. 172 | 173 | 174 | -------------------------------------------------------------------------------- /assets/meta_logo_white.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/assets/meta_logo_white.png -------------------------------------------------------------------------------- /assets/system.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/assets/system.png -------------------------------------------------------------------------------- /config/example_config.yaml: -------------------------------------------------------------------------------- 1 | # Example configuration file for DocAgent 2 | # Copy this file to agent_config.yaml and add your own API keys 3 | 4 | # LLM configuration for all agents 5 | llm: 6 | # Choose ONE of the following LLM provider configurations by uncommenting 7 | 8 | # Option 1: Claude (Anthropic) 9 | type: "claude" 10 | api_key: "your-anthropic-api-key-here" 11 | model: "claude-3-5-haiku-latest" # Options: claude-3-5-sonnet, claude-3-opus, etc. 12 | temperature: 0.1 13 | max_output_tokens: 4096 14 | max_input_tokens: 100000 # Maximum number of tokens for input context 15 | 16 | # Option 2: OpenAI 17 | # type: "openai" 18 | # api_key: "your-openai-api-key-here" 19 | # model: "gpt-4o" # Options: gpt-4o, gpt-4-turbo, gpt-3.5-turbo, etc. 20 | # temperature: 0.1 21 | # max_output_tokens: 4096 22 | # max_input_tokens: 100000 23 | 24 | # Option 3: Gemini 25 | # type: "gemini" 26 | # api_key: "your-gemini-api-key-here" 27 | # model: "gemini-1.5-pro" 28 | # temperature: 0.1 29 | # max_output_tokens: 4096 30 | # max_input_tokens: 100000 31 | 32 | # Option 4: HuggingFace (for local models) 33 | # type: "huggingface" 34 | # model: "codellama/CodeLlama-34b-Instruct-hf" 35 | # api_base: "http://localhost:8000/v1" # Local API endpoint 36 | # api_key: "EMPTY" # Can be empty for local models 37 | # device: "cuda" # Options: cuda, cpu 38 | # torch_dtype: "float16" 39 | # temperature: 0.1 40 | # max_output_tokens: 4096 41 | # max_input_tokens: 32000 42 | 43 | # Rate limit settings for different LLM providers 44 | # These are default values - adjust based on your specific API tier 45 | rate_limits: 46 | # Claude rate limits 47 | claude: 48 | requests_per_minute: 50 49 | input_tokens_per_minute: 20000 50 | output_tokens_per_minute: 8000 51 | input_token_price_per_million: 3.0 52 | output_token_price_per_million: 15.0 53 | 54 | # OpenAI rate limits 55 | openai: 56 | requests_per_minute: 500 57 | input_tokens_per_minute: 200000 58 | output_tokens_per_minute: 100000 59 | input_token_price_per_million: 0.15 60 | output_token_price_per_million: 0.60 61 | 62 | # Gemini rate limits 63 | gemini: 64 | requests_per_minute: 60 65 | input_tokens_per_minute: 30000 66 | output_tokens_per_minute: 10000 67 | input_token_price_per_million: 0.125 68 | output_token_price_per_million: 0.375 69 | 70 | # Flow control parameters 71 | flow_control: 72 | max_reader_search_attempts: 2 # Maximum times reader can call searcher 73 | max_verifier_rejections: 1 # Maximum times verifier can reject a docstring 74 | status_sleep_time: 1 # Time to sleep between status updates (seconds) 75 | 76 | # Docstring generation options 77 | docstring_options: 78 | overwrite_docstrings: false # Whether to overwrite existing docstrings (default: false) 79 | 80 | # Perplexity API configuration (for web search capability) 81 | perplexity: 82 | api_key: "your-perplexity-api-key-here" # Replace with your actual Perplexity API key 83 | model: "sonar" # Default model 84 | temperature: 0.1 85 | max_output_tokens: 250 -------------------------------------------------------------------------------- /data/raw_test_repo/README.md: -------------------------------------------------------------------------------- 1 | # Vending Machine Test Repository 2 | 3 | A comprehensive vending machine implementation in Python that demonstrates various programming concepts, design patterns, and documentation styles. This repository serves as a test bed for docstring generation systems and code documentation analysis. 4 | 5 | ## Project Structure 6 | 7 | ``` 8 | test_repo_vm/ 9 | ├── __init__.py # Main package initialization 10 | ├── example.py # Example usage demonstration 11 | ├── vending_machine.py # Main vending machine implementation 12 | ├── models/ # Data models 13 | │ ├── __init__.py 14 | │ └── product.py # Product class definition 15 | ├── payment/ # Payment processing 16 | │ ├── __init__.py 17 | │ └── payment_processor.py # Payment-related classes 18 | └── inventory/ # Inventory management 19 | ├── __init__.py 20 | └── inventory_manager.py # Inventory tracking system 21 | ``` 22 | 23 | ## Components 24 | 25 | ### 1. Product Management (`models/product.py`) 26 | - `Product` class with attributes like ID, name, price, quantity, and expiry date 27 | - Methods for checking availability and managing stock 28 | 29 | ### 2. Payment Processing (`payment/payment_processor.py`) 30 | - Abstract `PaymentMethod` base class for different payment types 31 | - `CashPayment` implementation for handling cash transactions 32 | - `PaymentTransaction` class for tracking payment status 33 | - `PaymentStatus` enum for transaction states 34 | 35 | ### 3. Inventory Management (`inventory/inventory_manager.py`) 36 | - `InventoryManager` class for product storage and retrieval 37 | - Slot-based product organization 38 | - Stock level tracking 39 | - Product availability checking 40 | 41 | ### 4. Main Vending Machine (`vending_machine.py`) 42 | - `VendingMachine` class that coordinates all components 43 | - Product selection and purchase workflow 44 | - Payment processing and change calculation 45 | - Exception handling for error cases 46 | 47 | ## Code Features 48 | 49 | This repository demonstrates various Python programming features: 50 | 51 | 1. **Object-Oriented Design** 52 | - Abstract base classes 53 | - Inheritance 54 | - Encapsulation 55 | - Interface definitions 56 | 57 | 2. **Modern Python Features** 58 | - Type hints 59 | - Dataclasses 60 | - Enums 61 | - Optional types 62 | - Package organization 63 | 64 | 3. **Documentation** 65 | - Comprehensive docstrings 66 | - Type annotations 67 | - Code organization 68 | - Exception documentation 69 | 70 | 4. **Best Practices** 71 | - SOLID principles 72 | - Clean code architecture 73 | - Error handling 74 | - Modular design 75 | 76 | ## Usage Example 77 | 78 | ```python 79 | from decimal import Decimal 80 | from vending_machine import VendingMachine 81 | from models.product import Product 82 | 83 | # Create a vending machine 84 | vm = VendingMachine() 85 | 86 | # Add products 87 | product = Product( 88 | id="COLA001", 89 | name="Cola Classic", 90 | price=1.50, 91 | quantity=10, 92 | category="drinks" 93 | ) 94 | vm.inventory.add_product(product, slot=0) 95 | 96 | # Insert money 97 | vm.insert_money(Decimal('2.00')) 98 | 99 | # Purchase product 100 | product, change = vm.purchase_product(slot=0) 101 | print(f"Purchased: {product.name}") 102 | print(f"Change: ${change:.2f}") 103 | ``` 104 | 105 | ## Running the Example 106 | 107 | To run the example implementation: 108 | 109 | ```bash 110 | python example.py 111 | ``` 112 | 113 | This will demonstrate: 114 | 1. Creating a vending machine 115 | 2. Adding products to inventory 116 | 3. Displaying available products 117 | 4. Making a purchase 118 | 5. Handling change 119 | 6. Updating inventory 120 | 121 | ## Testing Documentation Generation 122 | 123 | This repository is structured to test various aspects of documentation generation: 124 | 125 | 1. **Complex Imports** 126 | - Cross-module dependencies 127 | - Package-level imports 128 | - Relative imports 129 | 130 | 2. **Documentation Styles** 131 | - Function documentation 132 | - Class documentation 133 | - Module documentation 134 | - Package documentation 135 | 136 | 3. **Code Complexity** 137 | - Multiple inheritance 138 | - Abstract classes 139 | - Type annotations 140 | - Exception hierarchies 141 | 142 | ## Requirements 143 | 144 | - Python 3.7+ 145 | - No external dependencies required 146 | 147 | ## License 148 | 149 | This project is open source and available under the MIT License. -------------------------------------------------------------------------------- /data/raw_test_repo/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Vending Machine Package 4 | 5 | A comprehensive vending machine implementation with: 6 | - Product management 7 | - Inventory tracking 8 | - Payment processing 9 | - Transaction handling 10 | """ 11 | -------------------------------------------------------------------------------- /data/raw_test_repo/example.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from decimal import Decimal 3 | from datetime import datetime, timedelta 4 | from models.product import Item 5 | from vending_machine import Sys, SysErr 6 | 7 | 8 | def main(): 9 | s = Sys() 10 | items = [Item(code='D1', label='Drink1', val=1.5, count=10, grp='d', 11 | exp=datetime.now() + timedelta(days=90)), Item(code='S1', label= 12 | 'Snack1', val=1.0, count=15, grp='s', exp=datetime.now() + 13 | timedelta(days=30)), Item(code='S2', label='Snack2', val=2.0, count 14 | =8, grp='s', exp=datetime.now() + timedelta(days=60))] 15 | for i, item in enumerate(items): 16 | s.store.put(item, i) 17 | try: 18 | print('Items:') 19 | for pos, item in s.ls(): 20 | print(f'Pos {pos}: {item.label} - ${item.val:.2f}') 21 | pos = 0 22 | print('\nAdding $2.00...') 23 | s.add_money(Decimal('2.00')) 24 | item, ret = s.buy(pos) 25 | print(f'\nBought: {item.label}') 26 | if ret: 27 | print(f'Return: ${ret:.2f}') 28 | print('\nUpdated Items:') 29 | for pos, item in s.ls(): 30 | print( 31 | f'Pos {pos}: {item.label} - ${item.val:.2f} (Count: {item.count})' 32 | ) 33 | except SysErr as e: 34 | print(f'Err: {str(e)}') 35 | 36 | 37 | if __name__ == '__main__': 38 | main() 39 | -------------------------------------------------------------------------------- /data/raw_test_repo/inventory/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """Inventory management package for product stock tracking.""" 3 | -------------------------------------------------------------------------------- /data/raw_test_repo/inventory/inventory_manager.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import Dict, List, Optional 3 | from ..models.product import Item 4 | 5 | 6 | class Store: 7 | 8 | def __init__(self, cap: int=20): 9 | self.cap = cap 10 | self._data: Dict[str, Item] = {} 11 | self._map: Dict[int, str] = {} 12 | 13 | def put(self, obj: Item, pos: Optional[int]=None) ->bool: 14 | if obj.code in self._data: 15 | curr = self._data[obj.code] 16 | curr.count += obj.count 17 | return True 18 | if pos is not None: 19 | if pos < 0 or pos >= self.cap: 20 | return False 21 | if pos in self._map: 22 | return False 23 | self._map[pos] = obj.code 24 | else: 25 | for i in range(self.cap): 26 | if i not in self._map: 27 | self._map[i] = obj.code 28 | break 29 | else: 30 | return False 31 | self._data[obj.code] = obj 32 | return True 33 | 34 | def rm(self, code: str) ->bool: 35 | if code not in self._data: 36 | return False 37 | for k, v in list(self._map.items()): 38 | if v == code: 39 | del self._map[k] 40 | del self._data[code] 41 | return True 42 | 43 | def get(self, code: str) ->Optional[Item]: 44 | return self._data.get(code) 45 | 46 | def get_at(self, pos: int) ->Optional[Item]: 47 | if pos not in self._map: 48 | return None 49 | code = self._map[pos] 50 | return self._data.get(code) 51 | 52 | def ls(self) ->List[Item]: 53 | return [obj for obj in self._data.values() if obj.check()] 54 | 55 | def find(self, code: str) ->Optional[int]: 56 | for k, v in self._map.items(): 57 | if v == code: 58 | return k 59 | return None 60 | -------------------------------------------------------------------------------- /data/raw_test_repo/models/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """Models package for data structures used in the vending machine.""" 3 | -------------------------------------------------------------------------------- /data/raw_test_repo/models/product.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from dataclasses import dataclass 3 | from typing import Optional 4 | from datetime import datetime 5 | 6 | @dataclass 7 | class Item: 8 | """ 9 | Summary: 10 | Represents an item with associated attributes for tracking and management in various contexts. 11 | 12 | Description: 13 | This class serves as a blueprint for creating items that can be tracked and managed within a system. Each item has attributes such as a unique code, a label, a value, a count, an optional expiration date, and a group classification. The primary motivation behind this class is to facilitate resource management, inventory tracking, or any scenario where items need to be monitored for validity and availability. 14 | 15 | Use this class when you need to represent items that may have a limited lifespan or quantity, such as in inventory systems, gaming resources, or token management. It provides methods to check the validity of an item and to modify its count, ensuring that operations on the item are safe and consistent. 16 | 17 | The class fits into larger systems by allowing for easy integration with resource management workflows, enabling developers to track item states and manage their lifecycle effectively. 18 | 19 | Example: 20 | ```python 21 | from datetime import datetime, timedelta 22 | 23 | # Create an item with a specific expiration date 24 | item = Item(code='A123', label='Sample Item', val=10.0, count=5, exp=datetime.now() + timedelta(days=1)) 25 | 26 | # Check if the item is valid 27 | is_valid = item.check() # Returns True if count > 0 and not expired 28 | 29 | # Modify the count of the item 30 | item.mod(2) # Decreases count by 2, returns True 31 | ``` 32 | 33 | Parameters: 34 | - code (str): A unique identifier for the item. 35 | - label (str): A descriptive name for the item. 36 | - val (float): The value associated with the item, representing its worth. 37 | - count (int): The quantity of the item available. Must be a non-negative integer. 38 | - exp (Optional[datetime]): An optional expiration date for the item. If set, the item will be considered invalid after this date. 39 | - grp (str): A classification group for the item, defaulting to 'misc'. 40 | 41 | Attributes: 42 | - code (str): The unique identifier for the item. 43 | - label (str): The name or description of the item. 44 | - val (float): The monetary or functional value of the item. 45 | - count (int): The current quantity of the item available, must be non-negative. 46 | - exp (Optional[datetime]): The expiration date of the item, if applicable. 47 | - grp (str): The group classification of the item, useful for categorization. 48 | """ 49 | code: str 50 | label: str 51 | val: float 52 | count: int 53 | exp: Optional[datetime] = None 54 | grp: str = 'misc' 55 | 56 | def check(self) -> bool: 57 | """ 58 | Validates the current object's state based on count and expiration. 59 | 60 | Checks whether the object is still valid by verifying two key conditions: 61 | 1. The object's count is greater than zero 62 | 2. The object has not exceeded its expiration timestamp 63 | 64 | This method is typically used to determine if an object is still usable 65 | or has become stale/invalid. It provides a quick state validation check 66 | that can be used in resource management, token validation, or lifecycle 67 | tracking scenarios. 68 | 69 | Returns: 70 | bool: True if the object is valid (count > 0 and not expired), 71 | False otherwise. 72 | """ 73 | if self.count <= 0: 74 | return False 75 | if self.exp and datetime.now() > self.exp: 76 | return False 77 | return True 78 | 79 | def mod(self, n: int=1) -> bool: 80 | """ 81 | Summary: 82 | Determines if the current count can be decremented by a specified value. 83 | 84 | Description: 85 | This method checks if the `count` attribute is greater than or equal to the provided integer `n`. If so, it decrements `count` by `n` and returns `True`. If `count` is less than `n`, it returns `False`, indicating that the operation could not be performed. 86 | 87 | Use this function when managing resources or operations that require a controlled decrement of a count, ensuring that the count does not drop below zero. This is particularly useful in scenarios such as resource allocation, gaming mechanics, or iterative processes. 88 | 89 | The method is integral to classes that require precise control over a count, allowing for safe decrements while maintaining the integrity of the count value. 90 | 91 | Args: 92 | n (int, optional): The value to decrement from `count`. Must be a positive integer that does not exceed the current `count`. Default is 1. 93 | 94 | Returns: 95 | bool: Returns `True` if the decrement was successful (i.e., `count` was greater than or equal to `n`), otherwise returns `False`. 96 | 97 | Raises: 98 | No exceptions are raised by this method. Ensure that `n` is a positive integer and does not exceed the current `count` to avoid logical errors. 99 | 100 | Examples: 101 | ```python 102 | obj = YourClass() 103 | obj.count = 5 104 | result = obj.mod(2) # result will be True, obj.count will be 3 105 | result = obj.mod(4) # result will be False, obj.count remains 3 106 | result = obj.mod(0) # result will be False, as n should be greater than 0 107 | result = obj.mod(-1) # result will be False, as n should be a positive integer 108 | ``` 109 | """ 110 | if self.count >= n: 111 | self.count -= n 112 | return True 113 | return False -------------------------------------------------------------------------------- /data/raw_test_repo/payment/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """Payment processing package for handling different payment methods.""" 3 | -------------------------------------------------------------------------------- /data/raw_test_repo/payment/payment_processor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from abc import ABC, abstractmethod 3 | from dataclasses import dataclass 4 | from enum import Enum 5 | from typing import Optional 6 | from decimal import Decimal 7 | 8 | 9 | class TxStatus(Enum): 10 | WAIT = 'pending' 11 | DONE = 'completed' 12 | ERR = 'failed' 13 | RET = 'refunded' 14 | 15 | 16 | @dataclass 17 | class Tx: 18 | id: str 19 | amt: Decimal 20 | st: TxStatus 21 | mth: str 22 | msg: Optional[str] = None 23 | 24 | 25 | class Handler(ABC): 26 | 27 | @abstractmethod 28 | def proc(self, amt: Decimal) ->Tx: 29 | pass 30 | 31 | @abstractmethod 32 | def rev(self, tx: Tx) ->bool: 33 | pass 34 | 35 | 36 | class Cash(Handler): 37 | 38 | def __init__(self): 39 | self.bal: Decimal = Decimal('0.00') 40 | 41 | def add(self, amt: Decimal) ->None: 42 | self.bal += amt 43 | 44 | def proc(self, amt: Decimal) ->Tx: 45 | if self.bal >= amt: 46 | self.bal -= amt 47 | return Tx(id=f'C_{id(self)}', amt=amt, st=TxStatus.DONE, mth='cash' 48 | ) 49 | return Tx(id=f'C_{id(self)}', amt=amt, st=TxStatus.ERR, mth='cash', 50 | msg='insufficient') 51 | 52 | def rev(self, tx: Tx) ->bool: 53 | if tx.st == TxStatus.DONE: 54 | self.bal += tx.amt 55 | tx.st = TxStatus.RET 56 | return True 57 | return False 58 | 59 | def ret(self) ->Decimal: 60 | tmp = self.bal 61 | self.bal = Decimal('0.00') 62 | return tmp 63 | -------------------------------------------------------------------------------- /data/raw_test_repo/vending_machine.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from decimal import Decimal 3 | from typing import Optional, List, Tuple 4 | from .models.product import Item 5 | from .payment.payment_processor import Handler, Tx, TxStatus, Cash 6 | from .inventory.inventory_manager import Store 7 | 8 | 9 | class SysErr(Exception): 10 | pass 11 | 12 | 13 | class Sys: 14 | 15 | def __init__(self, h: Optional[Handler]=None): 16 | self.store = Store() 17 | self.h = h or Cash() 18 | self._tx: Optional[Tx] = None 19 | 20 | def ls(self) ->List[Tuple[int, Item]]: 21 | items = [] 22 | for item in self.store.ls(): 23 | pos = self.store.find(item.code) 24 | if pos is not None: 25 | items.append((pos, item)) 26 | return sorted(items, key=lambda x: x[0]) 27 | 28 | def pick(self, pos: int) ->Optional[Item]: 29 | item = self.store.get_at(pos) 30 | if not item: 31 | raise SysErr('invalid pos') 32 | if not item.check(): 33 | raise SysErr('unavailable') 34 | return item 35 | 36 | def add_money(self, amt: Decimal) ->None: 37 | if not isinstance(self.h, Cash): 38 | raise SysErr('cash not supported') 39 | self.h.add(amt) 40 | 41 | def buy(self, pos: int) ->Tuple[Item, Optional[Decimal]]: 42 | item = self.pick(pos) 43 | tx = self.h.proc(Decimal(str(item.val))) 44 | self._tx = tx 45 | if tx.st != TxStatus.DONE: 46 | raise SysErr(tx.msg or 'tx failed') 47 | if not item.mod(): 48 | self.h.rev(tx) 49 | raise SysErr('dispense failed') 50 | ret = None 51 | if isinstance(self.h, Cash): 52 | ret = self.h.ret() 53 | return item, ret 54 | 55 | def cancel(self) ->Optional[Decimal]: 56 | if not self._tx: 57 | raise SysErr('no tx') 58 | ok = self.h.rev(self._tx) 59 | if not ok: 60 | raise SysErr('rev failed') 61 | ret = None 62 | if isinstance(self.h, Cash): 63 | ret = self.h.ret() 64 | self._tx = None 65 | return ret 66 | -------------------------------------------------------------------------------- /data/raw_test_repo_simple/helper.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | class HelperClass: 3 | """ 4 | Represents a utility for managing and processing data. 5 | 6 | The `HelperClass` is designed to facilitate data processing tasks by leveraging the `DataProcessor` class. It serves as an intermediary that manages the workflow of data processing, making it easier to handle data updates and retrievals within a system. This class is particularly useful in scenarios where data needs to be processed and accessed in a structured manner. 7 | 8 | The `HelperClass` fits into the larger system architecture as a component that coordinates data processing tasks. It achieves its purpose by using the `DataProcessor` to perform the actual data processing and then managing the processed data internally. 9 | 10 | Example: 11 | # Initialize the HelperClass 12 | helper = HelperClass() 13 | 14 | # Process data using the helper 15 | helper.process_data() 16 | 17 | # Retrieve the processed data result 18 | result = helper.get_result() 19 | print(result) # Output: '[1, 2, 3]' 20 | 21 | Attributes: 22 | data (list): Stores the processed data, initially an empty list. 23 | """ 24 | 25 | def __init__(self): 26 | self.data = [] 27 | 28 | def process_data(self): 29 | """ 30 | Processes and updates the internal data. 31 | 32 | This method orchestrates the data processing workflow by invoking the `DataProcessor.process()` method to perform the main data processing task. It then calls `_internal_process()` to finalize the processing and update the internal `data` attribute. Use this method when you need to refresh or initialize the data within the `HelperClass` instance. 33 | 34 | Returns: 35 | None: This method updates the internal state and does not return a value. 36 | """ 37 | self.data = DataProcessor.process() 38 | self._internal_process() 39 | 40 | def _internal_process(self): 41 | """ 42 | No docstring provided. 43 | """ 44 | return self.data 45 | 46 | def get_result(self): 47 | """ 48 | No docstring provided. 49 | """ 50 | return str(self.data) 51 | 52 | class DataProcessor: 53 | ''' 54 | """Handles basic data processing tasks within a system. 55 | 56 | This class is designed to perform simple data processing operations, providing 57 | utility methods that can be used in various scenarios where basic data manipulation 58 | is required. It is particularly useful in contexts where a straightforward list of 59 | integers is needed for further processing or testing. 60 | 61 | The `DataProcessor` class fits into the larger system architecture as a utility 62 | component, offering static and internal methods to handle specific processing tasks. 63 | It achieves its purpose by providing a static method for general use and an internal 64 | method for class-specific operations. 65 | 66 | Example: 67 | # Initialize the DataProcessor class 68 | processor = DataProcessor() 69 | 70 | # Use the static method to process data 71 | result = DataProcessor.process() 72 | print(result) # Output: [1, 2, 3] 73 | 74 | # Use the internal method for internal processing 75 | internal_result = processor._internal_process() 76 | print(internal_result) # Output: 'processed' 77 | """ 78 | ''' 79 | 80 | @staticmethod 81 | def process(): 82 | ''' 83 | """Processes data and returns a list of integers. 84 | 85 | This static method is designed to perform a basic data processing task 86 | and return a predefined list of integers. It can be used whenever a simple 87 | list of integers is required for further operations or testing purposes. 88 | 89 | Returns: 90 | list of int: A list containing the integers [1, 2, 3]. 91 | """ 92 | ''' 93 | return [1, 2, 3] 94 | 95 | def _internal_process(self): 96 | ''' 97 | """Processes internal data and returns a status message. 98 | 99 | This method is used internally within the `DataProcessor` class to perform 100 | specific data processing tasks that are not exposed publicly. It is typically 101 | called by other methods within the class to handle intermediate processing 102 | steps. 103 | 104 | Returns: 105 | str: A string indicating the processing status, specifically 'processed'. 106 | """ 107 | ''' 108 | return 'processed' -------------------------------------------------------------------------------- /data/raw_test_repo_simple/inner/inner_functions.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | def inner_function(): 3 | """ 4 | Returns a greeting message from an inner function. 5 | 6 | This function is designed to return a simple greeting message, which can be used in nested or internal function calls to verify execution flow or for debugging purposes. It is typically used in development environments where confirming the execution of specific code paths is necessary. 7 | 8 | Returns: 9 | str: A greeting message stating 'Hello from inner function!' 10 | 11 | Example: 12 | >>> message = inner_function() 13 | >>> print(message) 14 | 'Hello from inner function!' 15 | """ 16 | return 'Hello from inner function!' 17 | 18 | def get_random_quote(): 19 | """ 20 | Fetches a predefined inspirational quote. 21 | 22 | This function is designed to provide users with a motivational quote, which can be used in applications that aim to inspire or uplift users. It is particularly useful in scenarios where a quick, positive message is needed to enhance user experience. 23 | 24 | Returns: 25 | str: A quote string stating 'The best way to predict the future is to create it.' 26 | 27 | Example: 28 | >>> quote = get_random_quote() 29 | >>> print(quote) 30 | 'The best way to predict the future is to create it.' 31 | """ 32 | return 'The best way to predict the future is to create it.' 33 | 34 | def generate_timestamp(): 35 | """ 36 | Generates and returns a static timestamp. 37 | 38 | This function provides a hardcoded timestamp string, which can be used in scenarios where a consistent and predictable timestamp is required for testing or logging purposes. It fits into workflows where a fixed date and time representation is needed without relying on the current system time. 39 | 40 | Returns: 41 | str: A string representing the static timestamp '2023-05-15 14:30:22'. 42 | """ 43 | return '2023-05-15 14:30:22' 44 | 45 | def get_system_status(): 46 | """ 47 | Provides a static message indicating the operational status of systems. 48 | 49 | This function is used to retrieve a fixed status message that confirms all systems are functioning correctly. It is useful in monitoring dashboards or status pages where a quick confirmation of system health is required. 50 | 51 | Returns: 52 | str: A status message stating 'All systems operational.' 53 | 54 | Example: 55 | >>> status = get_system_status() 56 | >>> print(status) 57 | 'All systems operational' 58 | """ 59 | return 'All systems operational' 60 | 61 | def fetch_user_message(): 62 | ''' 63 | """Fetches a predefined user message indicating notifications. 64 | 65 | This function is used to retrieve a static message that informs the user about the number of notifications they have. It is typically used in scenarios where a quick status update is needed for user engagement. 66 | 67 | Returns: 68 | str: A message string stating 'Welcome back! You have 3 notifications.' 69 | 70 | Example: 71 | >>> message = fetch_user_message() 72 | >>> print(message) 73 | 'Welcome back! You have 3 notifications.' 74 | """ 75 | ''' 76 | return 'Welcome back! You have 3 notifications.' -------------------------------------------------------------------------------- /data/raw_test_repo_simple/main.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from helper import HelperClass 3 | from inner.inner_functions import inner_function, get_random_quote, generate_timestamp, get_system_status, fetch_user_message 4 | 5 | def main_function(): 6 | """ 7 | Executes data processing and utility operations, returning the processed data as a string. 8 | 9 | This function initializes a `HelperClass` instance to manage and process data, invokes a utility function to provide a placeholder value, and generates a static timestamp for consistency in logging or testing scenarios. The function is useful when a complete data processing sequence is needed, integrating utility operations to produce a final result. 10 | 11 | Returns: 12 | str: The processed data result as a string, derived from the `HelperClass` instance after executing the data processing and utility functions. 13 | 14 | Example: 15 | # Execute the main function to process data and retrieve the result 16 | result = main_function() 17 | print(result) # Output: '[1, 2, 3]' 18 | """ 19 | helper = HelperClass() 20 | helper.process_data() 21 | utility_function() 22 | generate_timestamp() 23 | return helper.get_result() 24 | 25 | def utility_function(): 26 | """ 27 | Returns a utility string. 28 | 29 | This function provides a simple utility string, which can be used in various contexts where a placeholder or a generic return value is needed. It is typically used within workflows that require a consistent return value for testing or demonstration purposes. 30 | 31 | Returns: 32 | str: The string 'utility', serving as a generic utility value. 33 | """ 34 | return 'utility' -------------------------------------------------------------------------------- /data/raw_test_repo_simple/processor.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from helper import HelperClass 3 | from processor import DataProcessor 4 | from main import utility_function 5 | 6 | class AdvancedProcessor: 7 | """ 8 | Facilitates advanced data processing by coordinating multiple processing components. 9 | 10 | The `AdvancedProcessor` class is designed to manage and execute complex data processing workflows by integrating the functionalities of `HelperClass` and `DataProcessor`. It is ideal for scenarios where a comprehensive processing sequence is needed, providing a streamlined approach to handle data operations and produce a final result. 11 | 12 | This class fits into the larger system architecture as a high-level orchestrator of data processing tasks, ensuring that each component's capabilities are effectively utilized to achieve the desired outcome. 13 | 14 | Example: 15 | # Initialize the AdvancedProcessor 16 | processor = AdvancedProcessor() 17 | 18 | # Execute the processing workflow 19 | result = processor.run() 20 | print(result) # Output: 'utility' 21 | 22 | Attributes: 23 | helper (HelperClass): An instance of `HelperClass` used to manage data processing tasks. 24 | data_processor (DataProcessor): An instance of `DataProcessor` used to perform specific data processing operations. 25 | """ 26 | 27 | def __init__(self): 28 | self.helper = HelperClass() 29 | self.data_processor = DataProcessor() 30 | 31 | def run(self): 32 | """ 33 | Executes the complete data processing workflow and returns the result. 34 | 35 | This method coordinates the data processing tasks by utilizing both the `HelperClass` and `DataProcessor` to perform necessary operations. It is designed to be used when a full processing sequence is required, culminating in a final result that indicates the completion of these tasks. 36 | 37 | Returns: 38 | str: The result of the processing workflow, typically a utility string indicating successful completion. 39 | 40 | Example: 41 | # Create an instance of AdvancedProcessor 42 | processor = AdvancedProcessor() 43 | 44 | # Run the processing workflow 45 | result = processor.run() 46 | print(result) # Output: 'utility' 47 | """ 48 | self.helper.process_data() 49 | self.data_processor._internal_process() 50 | return self.process_result() 51 | 52 | def process_result(self): 53 | """ 54 | Returns a utility string as the result of processing. 55 | 56 | This method is part of the `AdvancedProcessor` class workflow, providing a consistent utility value after processing operations. It is typically used when a placeholder or generic result is needed following the execution of data processing tasks within the class. 57 | 58 | Returns: 59 | str: The string 'utility', serving as a generic utility value to indicate the completion of processing tasks. 60 | """ 61 | return utility_function() -------------------------------------------------------------------------------- /data/raw_test_repo_simple/test_file.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | def test_function(): 3 | """ 4 | Returns a boolean value indicating a successful test condition. 5 | 6 | This function is typically used in scenarios where a simple, consistent boolean value is required to represent a successful outcome or condition. It can be integrated into workflows that need a straightforward pass/fail indicator for testing or validation purposes. 7 | 8 | Returns: 9 | bool: The boolean value `True`, indicating a successful or positive condition. 10 | """ 11 | return True -------------------------------------------------------------------------------- /run_web_ui.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import eventlet 3 | eventlet.monkey_patch() 4 | # Copyright (c) Meta Platforms, Inc. and affiliates 5 | """ 6 | Web UI Launcher for DocAgent Docstring Generator 7 | 8 | This script launches the web-based user interface for the docstring generation tool. 9 | The UI provides a more interactive and visual way to use the docstring generator, 10 | with real-time feedback and progress tracking. 11 | 12 | Usage: 13 | python run_web_ui.py [--host HOST] [--port PORT] [--debug] 14 | """ 15 | 16 | import argparse 17 | import os 18 | import sys 19 | import logging 20 | from pathlib import Path 21 | 22 | # Configure logging 23 | logging.basicConfig( 24 | level=logging.INFO, 25 | format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', 26 | handlers=[ 27 | logging.StreamHandler(sys.stdout) 28 | ] 29 | ) 30 | logger = logging.getLogger("docstring_web") 31 | 32 | # Add the current directory to the path 33 | sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) 34 | 35 | def check_dependencies(): 36 | """Check if all required dependencies are installed.""" 37 | try: 38 | import flask 39 | import flask_socketio 40 | import eventlet 41 | import yaml 42 | import tabulate 43 | import colorama 44 | return True 45 | except ImportError as e: 46 | missing_module = str(e).split("'")[1] 47 | logger.error(f"Missing dependency: {missing_module}") 48 | logger.error("Please install all required dependencies with:") 49 | logger.error("pip install -r requirements-web.txt") 50 | return False 51 | 52 | def main(): 53 | """Parse command line arguments and start the web UI.""" 54 | parser = argparse.ArgumentParser(description='Launch the DocAgent Web UI') 55 | parser.add_argument('--host', default='127.0.0.1', help='Host to bind the server to') 56 | parser.add_argument('--port', type=int, default=5000, help='Port to bind the server to') 57 | parser.add_argument('--debug', action='store_true', help='Run in debug mode') 58 | 59 | args = parser.parse_args() 60 | 61 | # Check dependencies 62 | if not check_dependencies(): 63 | return 1 64 | 65 | # Print banner 66 | print("\n" + "=" * 80) 67 | print("DocAgent Web Interface".center(80)) 68 | print("=" * 80) 69 | 70 | # Import and run the web app 71 | try: 72 | # First try to import eventlet to ensure it's properly initialized 73 | import eventlet 74 | eventlet.monkey_patch() 75 | 76 | from src.web.app import create_app 77 | 78 | app, socketio = create_app(debug=args.debug) 79 | 80 | logger.info(f"Starting DocAgent Web UI at: http://{args.host}:{args.port}") 81 | logger.info("Press Ctrl+C to stop the server") 82 | 83 | # Start the server 84 | socketio.run(app, host=args.host, port=args.port, debug=args.debug, allow_unsafe_werkzeug=True) 85 | 86 | return 0 87 | except ImportError as e: 88 | logger.error(f"Error importing web application: {e}") 89 | logger.error("Make sure the src/web directory exists and contains the necessary files.") 90 | return 1 91 | except Exception as e: 92 | logger.error(f"Error running web application: {e}") 93 | return 1 94 | 95 | if __name__ == '__main__': 96 | try: 97 | sys.exit(main()) 98 | except KeyboardInterrupt: 99 | print("\nServer stopped.") 100 | sys.exit(0) -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from setuptools import setup, find_packages 3 | 4 | # Read the contents of README file 5 | from pathlib import Path 6 | this_directory = Path(__file__).parent 7 | long_description = (this_directory / "README.md").read_text() 8 | 9 | # Prepare all extras 10 | dev_requires = [ 11 | "pytest>=8.3.4", 12 | "pytest-cov>=2.0", 13 | "black>=22.0", 14 | "flake8>=3.9", 15 | ] 16 | 17 | web_requires = [ 18 | "flask>=3.1.0", 19 | "flask-socketio>=5.5.1", 20 | "eventlet>=0.39.0", 21 | "python-socketio>=5.12.1", 22 | "python-engineio>=4.11.2", 23 | "bidict>=0.23.0", 24 | "dnspython>=2.7.0", 25 | "six>=1.16.0", 26 | ] 27 | 28 | visualization_requires = [ 29 | "matplotlib>=3.10.0", 30 | "pygraphviz>=1.14", 31 | "networkx>=3.4.2", 32 | ] 33 | 34 | cuda_requires = [ 35 | "torch>=2.0.0", 36 | "accelerate>=1.4.0", 37 | ] 38 | 39 | # Combine all extras for the 'all' option 40 | all_requires = dev_requires + web_requires + visualization_requires + cuda_requires 41 | 42 | setup( 43 | name="DocstringGenerator", 44 | version="0.1.0", 45 | author="Dayu Yang", 46 | author_email="dayuyang@meta.com", 47 | description="DocAgent for High-quality docstring generation in Large-scale Python projects", 48 | long_description=long_description, 49 | long_description_content_type="text/markdown", 50 | packages=find_packages(where="src"), 51 | package_dir={"": "src"}, 52 | classifiers=[ 53 | "Development Status :: 3 - Alpha", 54 | "Intended Audience :: Developers", 55 | "License :: OSI Approved :: MIT License", 56 | "Programming Language :: Python :: 3", 57 | "Programming Language :: Python :: 3.8", 58 | "Programming Language :: Python :: 3.9", 59 | "Programming Language :: Python :: 3.10", 60 | ], 61 | python_requires=">=3.8", 62 | install_requires=[ 63 | # Core dependencies 64 | "numpy>=1.23.5", 65 | "pyyaml>=6.0", 66 | "jinja2>=3.1.5", 67 | "requests>=2.32.0", 68 | "urllib3>=2.3.0", 69 | 70 | # Code analysis tools 71 | "astor>=0.8.1", 72 | "code2flow>=2.5.1", 73 | "pydeps>=3.0.0", 74 | 75 | # AI/LLM related dependencies 76 | "anthropic>=0.45.0", 77 | "openai>=1.60.1", 78 | "langchain-anthropic>=0.3.4", 79 | "langchain-openai>=0.3.2", 80 | "langchain-core>=0.3.31", 81 | "langgraph>=0.2.67", 82 | "tiktoken>=0.8.0", 83 | "transformers>=4.48.0", 84 | "huggingface-hub>=0.28.0", 85 | "google-generativeai>=0.6.0", 86 | 87 | # Utility packages 88 | "tqdm>=4.67.1", 89 | "tabulate>=0.9.0", 90 | "colorama>=0.4.6", 91 | "termcolor>=2.5.0", 92 | "pydantic>=2.10.0", 93 | 94 | # Web requirements 95 | "flask>=3.1.0", 96 | "flask-socketio>=5.5.1", 97 | "eventlet>=0.39.0", 98 | "python-socketio>=5.12.1", 99 | "python-engineio>=4.11.2", 100 | "bidict>=0.23.0", 101 | "dnspython>=2.7.0", 102 | "six>=1.16.0", 103 | 104 | # CUDA requirements 105 | "torch>=2.0.0", 106 | "accelerate>=1.4.0", 107 | ], 108 | extras_require={ 109 | "dev": dev_requires, 110 | "web": web_requires, # Keep for potential compatibility, now included in core 111 | "visualization": visualization_requires, 112 | "cuda": cuda_requires, # Keep for potential compatibility, now included in core 113 | "all": all_requires, 114 | } 115 | ) -------------------------------------------------------------------------------- /src/DocstringGenerator.egg-info/SOURCES.txt: -------------------------------------------------------------------------------- 1 | README.md 2 | setup.py 3 | src/DocstringGenerator.egg-info/PKG-INFO 4 | src/DocstringGenerator.egg-info/SOURCES.txt 5 | src/DocstringGenerator.egg-info/dependency_links.txt 6 | src/DocstringGenerator.egg-info/requires.txt 7 | src/DocstringGenerator.egg-info/top_level.txt 8 | src/agent/__init__.py 9 | src/agent/base.py 10 | src/agent/orchestrator.py 11 | src/agent/reader.py 12 | src/agent/searcher.py 13 | src/agent/verifier.py 14 | src/agent/workflow.py 15 | src/agent/writer.py 16 | src/agent/llm/__init__.py 17 | src/agent/llm/base.py 18 | src/agent/llm/claude_llm.py 19 | src/agent/llm/factory.py 20 | src/agent/llm/gemini_llm.py 21 | src/agent/llm/huggingface_llm.py 22 | src/agent/llm/openai_llm.py 23 | src/agent/llm/rate_limiter.py 24 | src/dependency_analyzer/__init__.py 25 | src/dependency_analyzer/ast_parser.py 26 | src/dependency_analyzer/topo_sort.py 27 | src/evaluator/__init__.py 28 | src/evaluator/base.py 29 | src/evaluator/completeness.py 30 | src/evaluator/evaluation_common.py 31 | src/evaluator/helpfulness_attributes.py 32 | src/evaluator/helpfulness_description.py 33 | src/evaluator/helpfulness_evaluator.py 34 | src/evaluator/helpfulness_evaluator_ablation.py 35 | src/evaluator/helpfulness_examples.py 36 | src/evaluator/helpfulness_parameters.py 37 | src/evaluator/helpfulness_summary.py 38 | src/evaluator/segment.py 39 | src/evaluator/truthfulness.py 40 | src/visualizer/__init__.py 41 | src/visualizer/progress.py 42 | src/visualizer/status.py 43 | src/visualizer/web_bridge.py 44 | src/web/__init__.py 45 | src/web/app.py 46 | src/web/config_handler.py 47 | src/web/process_handler.py 48 | src/web/run.py 49 | src/web/visualization_handler.py -------------------------------------------------------------------------------- /src/DocstringGenerator.egg-info/dependency_links.txt: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/DocstringGenerator.egg-info/requires.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.23.5 2 | pyyaml>=6.0 3 | jinja2>=3.1.5 4 | requests>=2.32.0 5 | urllib3>=2.3.0 6 | astor>=0.8.1 7 | code2flow>=2.5.1 8 | pydeps>=3.0.0 9 | anthropic>=0.45.0 10 | openai>=1.60.1 11 | langchain-anthropic>=0.3.4 12 | langchain-openai>=0.3.2 13 | langchain-core>=0.3.31 14 | langgraph>=0.2.67 15 | tiktoken>=0.8.0 16 | transformers>=4.48.0 17 | huggingface-hub>=0.28.0 18 | google-generativeai>=0.6.0 19 | tqdm>=4.67.1 20 | tabulate>=0.9.0 21 | colorama>=0.4.6 22 | termcolor>=2.5.0 23 | pydantic>=2.10.0 24 | flask>=3.1.0 25 | flask-socketio>=5.5.1 26 | eventlet>=0.39.0 27 | python-socketio>=5.12.1 28 | python-engineio>=4.11.2 29 | bidict>=0.23.0 30 | dnspython>=2.7.0 31 | six>=1.16.0 32 | torch>=2.0.0 33 | accelerate>=1.4.0 34 | 35 | [all] 36 | pytest>=8.3.4 37 | pytest-cov>=2.0 38 | black>=22.0 39 | flake8>=3.9 40 | flask>=3.1.0 41 | flask-socketio>=5.5.1 42 | eventlet>=0.39.0 43 | python-socketio>=5.12.1 44 | python-engineio>=4.11.2 45 | bidict>=0.23.0 46 | dnspython>=2.7.0 47 | six>=1.16.0 48 | matplotlib>=3.10.0 49 | pygraphviz>=1.14 50 | networkx>=3.4.2 51 | torch>=2.0.0 52 | accelerate>=1.4.0 53 | 54 | [cuda] 55 | torch>=2.0.0 56 | accelerate>=1.4.0 57 | 58 | [dev] 59 | pytest>=8.3.4 60 | pytest-cov>=2.0 61 | black>=22.0 62 | flake8>=3.9 63 | 64 | [visualization] 65 | matplotlib>=3.10.0 66 | pygraphviz>=1.14 67 | networkx>=3.4.2 68 | 69 | [web] 70 | flask>=3.1.0 71 | flask-socketio>=5.5.1 72 | eventlet>=0.39.0 73 | python-socketio>=5.12.1 74 | python-engineio>=4.11.2 75 | bidict>=0.23.0 76 | dnspython>=2.7.0 77 | six>=1.16.0 78 | -------------------------------------------------------------------------------- /src/DocstringGenerator.egg-info/top_level.txt: -------------------------------------------------------------------------------- 1 | agent 2 | dependency_analyzer 3 | evaluator 4 | visualizer 5 | web 6 | -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-312.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/__pycache__/__init__.cpython-312.pyc -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /src/agent/README.md: -------------------------------------------------------------------------------- 1 | # Agent Framework for Docstring Generation 2 | 3 | This directory contains the core components of the multi-agent system responsible for generating high-quality docstrings for code components. 4 | 5 | ## Overview 6 | 7 | The system employs a collaborative workflow involving several specialized agents, managed by an Orchestrator. The goal is to analyze code, gather necessary context (both internal and external), generate a docstring, and verify its quality before finalizing. 8 | 9 | The main workflow is initiated via the `generate_docstring` function in `workflow.py`. 10 | 11 | ## Agents 12 | 13 | 1. **`BaseAgent` (`base.py`)** 14 | * **Role:** Abstract base class for all agents. 15 | * **Functionality:** Provides common infrastructure including LLM initialization (using `LLMFactory`), configuration loading, memory management (storing conversation history), and basic LLM interaction (`generate_response`). Ensures consistency across agents. 16 | 17 | 2. **`Reader` (`reader.py`)** 18 | * **Role:** Contextual Analysis and Information Needs Assessment. 19 | * **Functionality:** Analyzes the input code component (`focal_component`) and any existing context. Determines if additional information is required to write a comprehensive docstring. If more information is needed, it generates a structured request specifying whether internal codebase details (e.g., callers, callees) or external web search results are required. 20 | 21 | 3. **`Searcher` (`searcher.py`)** 22 | * **Role:** Information Retrieval. 23 | * **Functionality:** Acts upon the requests generated by the `Reader`. It retrieves the specified information by: 24 | * Querying the internal codebase using AST analysis (`ASTNodeAnalyzer`) and dependency graphs. 25 | * Performing external web searches via APIs (e.g., `PerplexityAPI`). 26 | * Returns the gathered context in a structured format. 27 | 28 | 4. **`Writer` (`writer.py`)** 29 | * **Role:** Docstring Generation. 30 | * **Functionality:** Takes the original code component and the accumulated context (provided by the `Orchestrator` after `Reader` and `Searcher` steps) as input. Uses its configured LLM and detailed prompts (tailored for classes vs. functions/methods, adhering to Google style guide) to generate the docstring. Outputs the generated docstring within specific XML tags (``). 31 | 32 | 5. **`Verifier` (`verifier.py`)** 33 | * **Role:** Quality Assurance. 34 | * **Functionality:** Evaluates the docstring produced by the `Writer` against the original code and the context used. Checks for clarity, accuracy, completeness, information value (avoiding redundancy), and appropriate level of detail. Determines if the docstring meets quality standards or requires revision. If revision is needed, it specifies whether more context is required or provides direct suggestions for improvement. 35 | 36 | 6. **`Orchestrator` (`orchestrator.py`)** 37 | * **Role:** Workflow Management. 38 | * **Functionality:** Coordinates the entire process. It manages the sequence of agent interactions: 39 | * Calls `Reader` to assess context needs. 40 | * Calls `Searcher` iteratively if more context is requested (up to a limit). 41 | * Calls `Writer` to generate the docstring. 42 | * Calls `Verifier` to evaluate the docstring. 43 | * Manages revision loops based on `Verifier` feedback, potentially involving further searches or refinement by the `Writer` (up to a limit). 44 | * Handles context accumulation, token limit constraints, and status visualization. 45 | 46 | ## Supporting Files 47 | 48 | * **`workflow.py`:** Provides the primary entry point function `generate_docstring` to initiate the docstring generation process for a given code component. 49 | * **`__init__.py`:** Makes the `agent` directory a Python package. 50 | * **`llm/`:** Contains LLM-related code, including the `LLMFactory` and base LLM classes. 51 | * **`tool/`:** Contains tools used by agents, such as the `ASTNodeAnalyzer` for internal code traversal and the `PerplexityAPI` wrapper for external search. -------------------------------------------------------------------------------- /src/agent/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | # Import only essential components to avoid circular imports 3 | from .reader import CodeComponentType 4 | 5 | # Explicitly list what should be accessible, but don't import until needed 6 | # to prevent circular imports 7 | __all__ = ['generate_docstring', 'CodeComponentType'] 8 | 9 | # Lazy load generate_docstring when it's actually needed 10 | def __getattr__(name): 11 | if name == 'generate_docstring': 12 | from .workflow import generate_docstring 13 | return generate_docstring 14 | raise AttributeError(f"module '{__name__}' has no attribute '{name}'") -------------------------------------------------------------------------------- /src/agent/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/base.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/base.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/base.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/base.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/docstring_workflow.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/docstring_workflow.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/orchestrator.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/orchestrator.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/orchestrator.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/orchestrator.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/reader.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/reader.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/reader.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/reader.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/searcher.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/searcher.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/searcher.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/searcher.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/verifier.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/verifier.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/verifier.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/verifier.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/workflow.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/workflow.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/writer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/writer.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/__pycache__/writer.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/__pycache__/writer.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/base.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from abc import ABC, abstractmethod 3 | from typing import Any, Dict, Optional, List 4 | import os 5 | from pathlib import Path 6 | 7 | from .llm.factory import LLMFactory 8 | from .llm.base import BaseLLM 9 | 10 | class BaseAgent(ABC): 11 | """Base class for all agents in the docstring generation system.""" 12 | 13 | def __init__(self, name: str, config_path: Optional[str] = None): 14 | """Initialize the base agent. 15 | 16 | Args: 17 | name: The name of the agent 18 | config_path: Optional path to the configuration file 19 | """ 20 | self.name = name 21 | self._memory: list[Dict[str, Any]] = [] 22 | 23 | # Initialize LLM and parameters from config 24 | self.llm, self.llm_params = self._initialize_llm(name, config_path) 25 | 26 | 27 | def _initialize_llm(self, agent_name: str, config_path: Optional[str] = None) -> tuple[BaseLLM, Dict[str, Any]]: 28 | """Initialize the LLM for this agent. 29 | 30 | Args: 31 | agent_name: Name of the agent 32 | config_path: Optional path to the configuration file 33 | 34 | Returns: 35 | Tuple of (Initialized LLM instance, LLM parameters dictionary) 36 | """ 37 | # Load configuration 38 | if config_path is None: 39 | config_path = "config/agent_config.yaml" 40 | print(f"Using default config from {config_path}") 41 | 42 | config = LLMFactory.load_config(config_path) 43 | 44 | # Check for agent-specific configuration 45 | agent_config = config.get("agent_llms", {}).get(agent_name.lower()) 46 | 47 | # Use agent-specific config if available, otherwise use default 48 | llm_config = agent_config if agent_config else config.get("llm", {}) 49 | 50 | # Verify api_key is provided in config 51 | if ("api_key" not in llm_config or not llm_config["api_key"]) and (llm_config["type"] not in ["huggingface", "local"]): 52 | raise ValueError("API key must be specified directly in the config file") 53 | 54 | # Extract LLM parameters 55 | llm_params = { 56 | "max_output_tokens": llm_config.get("max_output_tokens", 4096), 57 | "temperature": llm_config.get("temperature", 0.1), 58 | "model": llm_config.get("model") 59 | } 60 | 61 | return LLMFactory.create_llm(llm_config), llm_params 62 | 63 | def add_to_memory(self, role: str, content: str) -> None: 64 | """Add a message to the agent's memory. 65 | 66 | Args: 67 | role: The role of the message sender (e.g., 'system', 'user', 'assistant') 68 | content: The content of the message 69 | """ 70 | assert content is not None and content != "", "Content cannot be empty" 71 | self._memory.append(self.llm.format_message(role, content)) 72 | 73 | def refresh_memory(self, new_memory: list[Dict[str, Any]]) -> None: 74 | """Replace the current memory with new memory. 75 | 76 | Args: 77 | new_memory: The new memory to replace the current memory 78 | """ 79 | self._memory = [ 80 | self.llm.format_message(msg["role"], msg["content"]) 81 | for msg in new_memory 82 | ] 83 | 84 | def clear_memory(self) -> None: 85 | """Clear the agent's memory.""" 86 | self._memory = [] 87 | 88 | @property 89 | def memory(self) -> list[Dict[str, Any]]: 90 | """Get the agent's memory. 91 | 92 | Returns: 93 | The agent's memory as a list of message dictionaries 94 | """ 95 | return self._memory.copy() 96 | 97 | def generate_response(self, messages: Optional[List[Dict[str, Any]]] = None) -> str: 98 | """Generate a response using the agent's LLM and memory. 99 | 100 | Args: 101 | messages: Optional list of messages to use instead of memory 102 | 103 | Returns: 104 | Generated response text 105 | """ 106 | return self.llm.generate( 107 | messages=messages if messages is not None else self._memory, 108 | temperature=self.llm_params["temperature"], 109 | max_tokens=self.llm_params["max_output_tokens"] 110 | ) 111 | 112 | @abstractmethod 113 | def process(self, *args, **kwargs) -> Any: 114 | """Process the input and generate output. 115 | 116 | This method should be implemented by each specific agent. 117 | """ 118 | pass -------------------------------------------------------------------------------- /src/agent/llm/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from .base import BaseLLM 3 | from .openai_llm import OpenAILLM 4 | from .claude_llm import ClaudeLLM 5 | from .huggingface_llm import HuggingFaceLLM 6 | from .gemini_llm import GeminiLLM 7 | from .factory import LLMFactory 8 | 9 | __all__ = [ 10 | 'BaseLLM', 11 | 'OpenAILLM', 12 | 'ClaudeLLM', 13 | 'HuggingFaceLLM', 14 | 'GeminiLLM', 15 | 'LLMFactory' 16 | ] -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/base.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/base.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/base.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/base.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/claude_llm.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/claude_llm.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/claude_llm.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/claude_llm.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/factory.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/factory.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/factory.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/factory.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/gemini_llm.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/gemini_llm.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/gemini_llm.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/gemini_llm.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/huggingface_llm.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/huggingface_llm.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/huggingface_llm.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/huggingface_llm.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/openai_llm.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/openai_llm.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/openai_llm.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/openai_llm.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/rate_limiter.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/rate_limiter.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/llm/__pycache__/rate_limiter.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/llm/__pycache__/rate_limiter.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/llm/base.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from abc import ABC, abstractmethod 3 | from typing import List, Dict, Any, Optional 4 | 5 | class BaseLLM(ABC): 6 | """Base class for LLM wrappers.""" 7 | 8 | @abstractmethod 9 | def generate( 10 | self, 11 | messages: List[Dict[str, str]], 12 | temperature: float = 0.7, 13 | max_output_tokens: Optional[int] = None 14 | ) -> str: 15 | """Generate a response from the LLM. 16 | 17 | Args: 18 | messages: List of message dictionaries with 'role' and 'content' keys 19 | temperature: Sampling temperature (0.0 to 1.0) 20 | max_output_tokens: Maximum number of tokens to generate 21 | 22 | Returns: 23 | The generated response text 24 | """ 25 | pass 26 | 27 | @abstractmethod 28 | def format_message(self, role: str, content: str) -> Dict[str, str]: 29 | """Format a message for the specific LLM API. 30 | 31 | Args: 32 | role: The role of the message sender 33 | content: The content of the message 34 | 35 | Returns: 36 | Formatted message dictionary 37 | """ 38 | pass -------------------------------------------------------------------------------- /src/agent/llm/claude_llm.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import List, Dict, Any, Optional 3 | import anthropic 4 | from .base import BaseLLM 5 | from .rate_limiter import RateLimiter 6 | import logging 7 | 8 | class ClaudeLLM(BaseLLM): 9 | """Anthropic Claude API wrapper.""" 10 | 11 | def __init__( 12 | self, 13 | api_key: str, 14 | model: str, 15 | rate_limits: Optional[Dict[str, Any]] = None 16 | ): 17 | """Initialize Claude LLM. 18 | 19 | Args: 20 | api_key: Anthropic API key 21 | model: Model identifier (e.g., "claude-3-sonnet-20240229") 22 | rate_limits: Optional dictionary with rate limit settings 23 | """ 24 | self.client = anthropic.Anthropic(api_key=api_key) 25 | self.model = model 26 | 27 | # Default rate limits for Claude 3.7 Sonnet 28 | default_limits = { 29 | "requests_per_minute": 50, 30 | "input_tokens_per_minute": 20000, 31 | "output_tokens_per_minute": 8000, 32 | "input_token_price_per_million": 3.0, 33 | "output_token_price_per_million": 15.0 34 | } 35 | 36 | # Use provided rate limits or defaults 37 | limits = rate_limits or default_limits 38 | 39 | # Initialize rate limiter 40 | self.rate_limiter = RateLimiter( 41 | provider="Claude", 42 | requests_per_minute=limits.get("requests_per_minute", default_limits["requests_per_minute"]), 43 | input_tokens_per_minute=limits.get("input_tokens_per_minute", default_limits["input_tokens_per_minute"]), 44 | output_tokens_per_minute=limits.get("output_tokens_per_minute", default_limits["output_tokens_per_minute"]), 45 | input_token_price_per_million=limits.get("input_token_price_per_million", default_limits["input_token_price_per_million"]), 46 | output_token_price_per_million=limits.get("output_token_price_per_million", default_limits["output_token_price_per_million"]) 47 | ) 48 | 49 | def _count_tokens(self, text: str) -> int: 50 | """Count tokens in a string using Claude's tokenizer. 51 | 52 | Args: 53 | text: Text to count tokens for 54 | 55 | Returns: 56 | Token count 57 | """ 58 | if not text: 59 | return 0 60 | 61 | try: 62 | # Format text as a message for token counting 63 | count = self.client.beta.messages.count_tokens( 64 | model=self.model, 65 | messages=[ 66 | {"role": "user", "content": text} 67 | ] 68 | ) 69 | return count.input_tokens 70 | except Exception as e: 71 | # Log the error but don't fail 72 | logging.warning(f"Failed to count tokens with Claude tokenizer: {e}") 73 | # Fallback: rough estimate if tokenizer fails 74 | return len(text.split()) * 1.3 75 | 76 | def _count_messages_tokens(self, messages: List[Dict[str, str]], system_message: Optional[str] = None) -> int: 77 | """Count tokens in message list with optional system message. 78 | 79 | Args: 80 | messages: List of message dictionaries 81 | system_message: Optional system message 82 | 83 | Returns: 84 | Total token count 85 | """ 86 | if not messages: 87 | return 0 88 | 89 | # Convert messages to Claude format 90 | claude_messages = [self._convert_to_claude_message(msg) for msg in messages 91 | if msg["role"] != "system"] 92 | 93 | # Format system message if provided 94 | system_content = None 95 | if system_message: 96 | system_content = system_message 97 | 98 | try: 99 | # Use the API to count tokens for all messages at once 100 | count = self.client.beta.messages.count_tokens( 101 | model=self.model, 102 | messages=claude_messages, 103 | system=system_content 104 | ) 105 | return count.input_tokens 106 | except Exception as e: 107 | # Log the error but don't fail 108 | logging.warning(f"Failed to count tokens with Claude tokenizer: {e}") 109 | 110 | # Fallback: count tokens individually 111 | total_tokens = 0 112 | for msg in claude_messages: 113 | if "content" in msg and msg["content"]: 114 | total_tokens += self._count_tokens(msg["content"]) 115 | 116 | # Add system message tokens if provided 117 | if system_message: 118 | total_tokens += self._count_tokens(system_message) 119 | 120 | # Add overhead for message formatting 121 | total_tokens += 10 * len(claude_messages) # Add ~10 tokens per message for formatting 122 | 123 | return total_tokens 124 | 125 | def generate( 126 | self, 127 | messages: List[Dict[str, str]], 128 | temperature: float, 129 | max_tokens: Optional[int] 130 | ) -> str: 131 | """Generate a response using Claude API with rate limiting. 132 | 133 | Args: 134 | messages: List of message dictionaries 135 | temperature: Sampling temperature 136 | max_output_tokens: Maximum tokens to generate 137 | 138 | Returns: 139 | Generated response text 140 | """ 141 | # Extract system message if present 142 | system_message = None 143 | chat_messages = [] 144 | 145 | for msg in messages: 146 | if msg["role"] == "system": 147 | system_message = msg["content"] 148 | else: 149 | chat_messages.append(self._convert_to_claude_message(msg)) 150 | 151 | # Count input tokens 152 | input_tokens = self._count_messages_tokens(messages, system_message) 153 | 154 | # Wait if we're approaching rate limits (estimate output tokens as max_output_tokens) 155 | self.rate_limiter.wait_if_needed(input_tokens, max_tokens) 156 | 157 | # Make the API call 158 | response = self.client.messages.create( 159 | model=self.model, 160 | messages=chat_messages, 161 | system=system_message, 162 | temperature=temperature, 163 | max_tokens=max_tokens 164 | ) 165 | 166 | result_text = response.content[0].text 167 | 168 | # Count output tokens and record request 169 | output_tokens = self._count_tokens(result_text) 170 | self.rate_limiter.record_request(input_tokens, output_tokens) 171 | 172 | return result_text 173 | 174 | def format_message(self, role: str, content: str) -> Dict[str, str]: 175 | """Format message for Claude API. 176 | 177 | Args: 178 | role: Message role (system, user, assistant) 179 | content: Message content 180 | 181 | Returns: 182 | Formatted message dictionary 183 | """ 184 | # Store in standard format, conversion happens in generate() 185 | return {"role": role, "content": content} 186 | 187 | def _convert_to_claude_message(self, message: Dict[str, str]) -> Dict[str, str]: 188 | """Convert standard message format to Claude's format. 189 | 190 | Args: 191 | message: Standard format message 192 | 193 | Returns: 194 | Claude format message 195 | """ 196 | role_mapping = { 197 | "user": "user", 198 | "assistant": "assistant" 199 | } 200 | 201 | role = role_mapping[message["role"]] 202 | content = message["content"] 203 | 204 | return {"role": role, "content": content} -------------------------------------------------------------------------------- /src/agent/llm/factory.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import Dict, Any, Optional 3 | from pathlib import Path 4 | import yaml 5 | 6 | from .base import BaseLLM 7 | from .openai_llm import OpenAILLM 8 | from .claude_llm import ClaudeLLM 9 | from .huggingface_llm import HuggingFaceLLM 10 | from .gemini_llm import GeminiLLM 11 | 12 | class LLMFactory: 13 | """Factory class for creating LLM instances.""" 14 | 15 | @staticmethod 16 | def create_llm(config: Dict[str, Any]) -> BaseLLM: 17 | """Create an LLM instance based on configuration. 18 | 19 | Args: 20 | config: Configuration dictionary containing LLM settings 21 | 22 | Returns: 23 | An instance of BaseLLM 24 | 25 | Raises: 26 | ValueError: If the LLM type is not supported 27 | """ 28 | llm_type = config["type"].lower() 29 | model = config.get("model") 30 | 31 | if not model: 32 | raise ValueError("Model must be specified in the config file") 33 | 34 | # Extract rate limit settings from config 35 | # First check if there are specific rate limits in the LLM config 36 | rate_limits = config.get("rate_limits", {}) 37 | 38 | # If not, check if there are global rate limits for this provider type 39 | global_config = LLMFactory.load_config() 40 | if not rate_limits and "rate_limits" in global_config: 41 | # Map LLM types to provider names in rate_limits section 42 | provider_map = { 43 | "openai": "openai", 44 | "claude": "claude", 45 | "gemini": "gemini" 46 | } 47 | provider_key = provider_map.get(llm_type, llm_type) 48 | provider_limits = global_config.get("rate_limits", {}).get(provider_key, {}) 49 | if provider_limits: 50 | rate_limits = provider_limits 51 | 52 | if llm_type == "openai": 53 | return OpenAILLM( 54 | api_key=config["api_key"], 55 | model=model, 56 | rate_limits=rate_limits 57 | ) 58 | elif llm_type == "claude": 59 | return ClaudeLLM( 60 | api_key=config["api_key"], 61 | model=model, 62 | rate_limits=rate_limits 63 | ) 64 | elif llm_type == "gemini": 65 | return GeminiLLM( 66 | api_key=config["api_key"], 67 | model=model, 68 | rate_limits=rate_limits 69 | ) 70 | elif llm_type == "huggingface": 71 | return HuggingFaceLLM( 72 | model_name=model, 73 | device=config.get("device", "cuda"), 74 | torch_dtype=config.get("torch_dtype", "float16") 75 | ) 76 | else: 77 | raise ValueError(f"Unsupported LLM type: {llm_type}") 78 | 79 | @staticmethod 80 | def load_config(config_path: Optional[str] = None) -> Dict[str, Any]: 81 | """Load LLM configuration from file. 82 | 83 | Args: 84 | config_path: Path to the configuration file. If None, uses default path. 85 | 86 | Returns: 87 | Configuration dictionary 88 | 89 | Raises: 90 | FileNotFoundError: If the configuration file doesn't exist 91 | """ 92 | if config_path is None: 93 | config_path = str(Path(__file__).parent.parent.parent.parent / "config" / "agent_config.yaml") 94 | 95 | if not Path(config_path).exists(): 96 | raise FileNotFoundError(f"Configuration file not found: {config_path}") 97 | 98 | with open(config_path, 'r') as f: 99 | config = yaml.safe_load(f) 100 | 101 | return config -------------------------------------------------------------------------------- /src/agent/llm/openai_llm.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import List, Dict, Any, Optional 3 | import openai 4 | import tiktoken 5 | from .base import BaseLLM 6 | from .rate_limiter import RateLimiter 7 | 8 | class OpenAILLM(BaseLLM): 9 | """OpenAI API wrapper.""" 10 | 11 | def __init__( 12 | self, 13 | api_key: str, 14 | model: str, 15 | rate_limits: Optional[Dict[str, Any]] = None 16 | ): 17 | """Initialize OpenAI LLM. 18 | 19 | Args: 20 | api_key: OpenAI API key 21 | model: Model identifier (e.g., "gpt-4", "gpt-3.5-turbo") 22 | rate_limits: Optional dictionary with rate limit settings 23 | """ 24 | self.client = openai.OpenAI(api_key=api_key) 25 | self.model = model 26 | 27 | try: 28 | # Initialize tokenizer for the model 29 | self.tokenizer = tiktoken.encoding_for_model(model) 30 | except: 31 | # Fallback to cl100k_base for new models 32 | self.tokenizer = tiktoken.get_encoding("cl100k_base") 33 | 34 | # Default rate limits for GPT-4o-mini 35 | default_limits = { 36 | "requests_per_minute": 500, 37 | "input_tokens_per_minute": 200000, 38 | "output_tokens_per_minute": 100000, 39 | "input_token_price_per_million": 0.15, 40 | "output_token_price_per_million": 0.60 41 | } 42 | 43 | # Use provided rate limits or defaults 44 | limits = rate_limits or default_limits 45 | 46 | # Initialize rate limiter 47 | self.rate_limiter = RateLimiter( 48 | provider="OpenAI", 49 | requests_per_minute=limits.get("requests_per_minute", default_limits["requests_per_minute"]), 50 | input_tokens_per_minute=limits.get("input_tokens_per_minute", default_limits["input_tokens_per_minute"]), 51 | output_tokens_per_minute=limits.get("output_tokens_per_minute", default_limits["output_tokens_per_minute"]), 52 | input_token_price_per_million=limits.get("input_token_price_per_million", default_limits["input_token_price_per_million"]), 53 | output_token_price_per_million=limits.get("output_token_price_per_million", default_limits["output_token_price_per_million"]) 54 | ) 55 | 56 | def _count_tokens(self, text: str) -> int: 57 | """Count tokens in a string using the model's tokenizer. 58 | 59 | Args: 60 | text: Text to count tokens for 61 | 62 | Returns: 63 | Token count 64 | """ 65 | if not text: 66 | return 0 67 | 68 | try: 69 | return len(self.tokenizer.encode(text)) 70 | except Exception as e: 71 | # Log the error but don't fail 72 | import logging 73 | logging.warning(f"Failed to count tokens with OpenAI tokenizer: {e}") 74 | # Fallback: rough estimate if tokenizer fails 75 | return len(text.split()) * 1.3 76 | 77 | def _count_messages_tokens(self, messages: List[Dict[str, str]]) -> int: 78 | """Count tokens in all messages. 79 | 80 | Args: 81 | messages: List of message dictionaries 82 | 83 | Returns: 84 | Total token count 85 | """ 86 | if not messages: 87 | return 0 88 | 89 | total_tokens = 0 90 | 91 | # Count tokens in each message 92 | for message in messages: 93 | if "content" in message and message["content"]: 94 | total_tokens += self._count_tokens(message["content"]) 95 | 96 | # Add overhead for message formatting (varies by model, but ~4 tokens per message) 97 | total_tokens += 4 * len(messages) 98 | 99 | # Add tokens for model overhead (varies by model) 100 | total_tokens += 3 # Every reply is primed with <|start|>assistant<|message|> 101 | 102 | return total_tokens 103 | 104 | def generate( 105 | self, 106 | messages: List[Dict[str, str]], 107 | temperature: float, 108 | max_tokens: Optional[int] 109 | ) -> str: 110 | """Generate a response using OpenAI API with rate limiting. 111 | 112 | Args: 113 | messages: List of message dictionaries 114 | temperature: Sampling temperature 115 | max_output_tokens: Maximum tokens to generate 116 | 117 | Returns: 118 | Generated response text 119 | """ 120 | # Count input tokens 121 | input_tokens = self._count_messages_tokens(messages) 122 | 123 | # Wait if we're approaching rate limits (estimate output tokens as max_output_tokens) 124 | self.rate_limiter.wait_if_needed(input_tokens, max_tokens) 125 | 126 | # Make the API call 127 | response = self.client.chat.completions.create( 128 | model=self.model, 129 | messages=messages, 130 | temperature=temperature, 131 | max_tokens=max_tokens if max_tokens else None 132 | ) 133 | 134 | result_text = response.choices[0].message.content 135 | 136 | # Count output tokens and record request 137 | output_tokens = response.usage.completion_tokens if hasattr(response, 'usage') else self._count_tokens(result_text) 138 | input_tokens = response.usage.prompt_tokens if hasattr(response, 'usage') else input_tokens 139 | 140 | self.rate_limiter.record_request(input_tokens, output_tokens) 141 | 142 | return result_text 143 | 144 | def format_message(self, role: str, content: str) -> Dict[str, str]: 145 | """Format message for OpenAI API. 146 | 147 | Args: 148 | role: Message role (system, user, assistant) 149 | content: Message content 150 | 151 | Returns: 152 | Formatted message dictionary 153 | """ 154 | # OpenAI uses standard role names 155 | return {"role": role, "content": content} -------------------------------------------------------------------------------- /src/agent/reader.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from dataclasses import dataclass 3 | from enum import Enum 4 | from typing import Any, Dict, List, Optional, Tuple 5 | 6 | from .base import BaseAgent 7 | 8 | 9 | class CodeComponentType(Enum): 10 | """Enum for different types of code components.""" 11 | 12 | FUNCTION = "function" 13 | METHOD = "method" 14 | CLASS = "class" 15 | 16 | 17 | @dataclass 18 | class InformationRequest: 19 | """Data class for structured information requests.""" 20 | 21 | internal_requests: List[str] 22 | external_requests: List[str] 23 | 24 | 25 | class Reader(BaseAgent): 26 | """Agent responsible for determining if more context is needed for docstring generation.""" 27 | 28 | def __init__(self, config_path: Optional[str] = None): 29 | """Initialize the Reader agent. 30 | 31 | Args: 32 | config_path: Optional path to the configuration file 33 | """ 34 | super().__init__("Reader", config_path) 35 | self.system_prompt = """You are a Reader agent responsible for determining if more context 36 | is needed to generate a high-quality docstring. You should analyze the code component and 37 | current context to make this determination. 38 | 39 | You have access to two types of information sources: 40 | 41 | 1. Internal Codebase Information (from local code repository): 42 | For Functions: 43 | - Code components called within the function body 44 | - Places where this function is called 45 | 46 | For Methods: 47 | - Code components called within the method body 48 | - Places where this method is called 49 | - The class this method belongs to 50 | 51 | For Classes: 52 | - Code components called in the __init__ method 53 | - Places where this class is instantiated 54 | - Complete class implementation beyond __init__ 55 | 56 | 2. External Open Internet retrieval Information: 57 | - External Retrieval is extremely expensive. Only request external open internet retrieval information if the component involves a novel, state of the art, recently-proposed algorithms or techniques. 58 | (e.g. computing a novel loss function (NDCG Loss, Alignment and Uniformity Loss, etc), certain novel metrics (Cohen's Kappa, etc), specialized novel ideas) 59 | - Each query should be a clear, natural language question 60 | 61 | Your response should: 62 | 1. First provide a free text analysis of the current code and context 63 | 2. Explain what additional information might be needed (if any) 64 | 3. Include an true tag if more information is needed, 65 | or false if current context is sufficient 66 | 4. If more information is needed, end your response with a structured request in XML format: 67 | 68 | 69 | 70 | 71 | class1,class2 72 | func1,func2 73 | self.method1,instance.method2,class.method3 74 | 75 | true/false 76 | 77 | 78 | query1,query2 79 | 80 | 81 | 82 | Important rules for structured request: 83 | 1. For CALLS sections, only include names that are explicitly needed 84 | 2. If no items exist for a category, use empty tags (e.g., ) 85 | 3. CALL_BY should be "true" only if you need to know what calls/uses a component 86 | 4. Each external QUERY should be a concise, clear, natural language search query 87 | 5. Use comma-separated values without spaces for multiple items 88 | 6. For METHODS, keep dot notation in the same format as the input. 89 | 7. Only first-level calls of the focal code component are accessible. Do not request information on code components that are not directly called by the focal component. 90 | 8. External Open-Internet Retrieval is extremely expensive. Only request external open internet retrieval information if the component involves a novel, state of the art, recently-proposed algorithms or techniques. 91 | (e.g. computing a novel loss function (NDCG Loss, Alignment and Uniformity Loss, etc), certain novel metrics (Cohen's Kappa, etc), specialized novel ideas) 92 | 93 | 94 | Important rules: 95 | 1. Only request internal codebase information that you think is necessary for docstring generation task. For some components that is simple and obvious, you do not need any other information for docstring generation. 96 | 2. External Open-Internet retrieval request is extremely expensive. Only request information that you think is absolutely necessary for docstring generation task. 97 | 98 | 99 | The current code shows a database connection function. To write a comprehensive docstring, we need to understand: 100 | 1. Where this function is called - this will reveal the expected input patterns and common use cases 101 | 2. What internal database functions it relies on - this will help document any dependencies or prerequisites 102 | 103 | This additional context is necessary because database connections often have specific setup requirements and usage patterns that should be documented for proper implementation. 104 | 105 | true 106 | 107 | 108 | 109 | 110 | 111 | execute_query,connect_db 112 | self.process_data,data_processor._internal_process 113 | 114 | true 115 | 116 | 117 | 118 | 119 | 120 | 121 | 122 | 123 | Keep in mind that: 124 | 125 | 3. You do not need to generate docstring for the component. Just determine if more information is needed. 126 | """ 127 | self.add_to_memory("system", self.system_prompt) 128 | 129 | def process(self, focal_component: str, context: str = "") -> str: 130 | """Process the input and determine if more context is needed. 131 | 132 | Args: 133 | instruction: The instruction for docstring generation 134 | focal_component: The code component needing a docstring (full code snippet) 135 | component_type: The type of the code component (function, method, or class) 136 | context: Current context information (if any) 137 | 138 | Returns: 139 | A string containing the analysis and tag indicating if more information is needed 140 | """ 141 | # Add the current task to memory 142 | task_description = f""" 143 | 144 | Current context: 145 | {context if context else 'No context provided yet.'} 146 | 147 | 148 | 149 | Analyze the following code component: 150 | 151 | {focal_component} 152 | 153 | """ 154 | self.add_to_memory("user", task_description) 155 | 156 | # Generate response using LLM 157 | response = self.generate_response() 158 | return response 159 | -------------------------------------------------------------------------------- /src/agent/tool/README.md: -------------------------------------------------------------------------------- 1 | # AST Call Graph Analysis Tool 2 | 3 | This tool provides functionality to analyze Python codebases by building and querying call graphs using Abstract Syntax Tree (AST) parsing. It helps in understanding code relationships and dependencies between functions, methods, and classes. 4 | 5 | ## Features 6 | 7 | ### Call Graph Building 8 | - Automatically builds a complete call graph for a Python repository 9 | - Tracks relationships between functions, methods, and classes 10 | - Handles cross-file dependencies 11 | - Caches AST parsing results for better performance 12 | 13 | ### Code Component Analysis 14 | 15 | The tool provides six main functionalities for analyzing code relationships: 16 | 17 | 1. **Child Function Analysis** (`get_child_function`) 18 | - Input: Component signature, file path, and child function name 19 | - Output: Full code of the function being called 20 | - Use case: Finding implementation of functions called within your code 21 | 22 | 2. **Child Method Analysis** (`get_child_method`) 23 | - Input: Component signature, file path, and child method name 24 | - Output: Full code of the method being called 25 | - Use case: Finding implementation of methods called on objects 26 | 27 | 3. **Child Class Analysis** (`get_child_class`) 28 | - Input: Component signature, file path, and child class name 29 | - Output: Class signature and initialization code 30 | - Use case: Finding class definitions for instantiated objects 31 | 32 | 4. **Parent Function Analysis** (`get_parent_function`) 33 | - Input: Component signature, file path, and parent function name 34 | - Output: Full code of the function that calls the component 35 | - Use case: Finding where a function is being used 36 | 37 | 5. **Parent Method Analysis** (`get_parent_method`) 38 | - Input: Component signature, file path, and parent method name 39 | - Output: Full code of the method that calls the component 40 | - Use case: Finding where a method is being called 41 | 42 | 6. **Parent Class Analysis** (`get_parent_class`) 43 | - Input: Component signature, file path, and parent class name 44 | - Output: Full code of the class that uses the component 45 | - Use case: Finding classes that depend on other classes 46 | 47 | ## Usage Example 48 | 49 | ```python 50 | from agent.tool.ast import CallGraphBuilder 51 | 52 | # Initialize the builder with repository path 53 | builder = CallGraphBuilder("/path/to/repo") 54 | 55 | # Find where a function is called 56 | parent_code = builder.get_parent_function( 57 | "def process_data(self):", 58 | "src/data/processor.py", 59 | "main_function" 60 | ) 61 | 62 | # Find what methods a class uses 63 | child_code = builder.get_child_method( 64 | "class DataProcessor:", 65 | "src/data/processor.py", 66 | "transform_data" 67 | ) 68 | ``` 69 | 70 | ## Implementation Details 71 | 72 | - Uses Python's built-in `ast` module for code parsing 73 | - Maintains parent-child relationships in AST nodes 74 | - Handles various Python constructs: 75 | - Function definitions and calls 76 | - Class definitions and instantiations 77 | - Method calls (both direct and through objects) 78 | - Static methods 79 | - Internal methods 80 | - Cross-file dependencies 81 | 82 | ## Limitations 83 | 84 | - Currently only supports Python files 85 | - Requires valid Python syntax in source files 86 | - Does not handle dynamic code execution (eval, exec) 87 | - Method resolution is name-based (doesn't handle complex inheritance) 88 | - Doesn't track calls through variables or complex expressions -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/ast.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/ast.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/ast.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/ast.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/ast_analyzer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/ast_analyzer.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/internal_traverse.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/internal_traverse.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/internal_traverse.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/internal_traverse.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/perplexity_api.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/perplexity_api.cpython-310.pyc -------------------------------------------------------------------------------- /src/agent/tool/__pycache__/perplexity_api.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/agent/tool/__pycache__/perplexity_api.cpython-311.pyc -------------------------------------------------------------------------------- /src/agent/tool/perplexity_api.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | import os 3 | import requests 4 | from typing import List, Dict, Any 5 | from dataclasses import dataclass 6 | import yaml 7 | 8 | @dataclass 9 | class PerplexityResponse: 10 | """Structured response from Perplexity API""" 11 | content: str 12 | raw_response: Dict[str, Any] 13 | 14 | class PerplexityAPI: 15 | """Wrapper for Perplexity API interactions""" 16 | 17 | def __init__(self, api_key: str | None = None, config_path: str = "config/agent_config.yaml"): 18 | """Initialize the API wrapper. 19 | 20 | Args: 21 | api_key: Perplexity API key. If None, will try to get from config. 22 | config_path: Path to the configuration file 23 | """ 24 | self.config = self._load_config(config_path) 25 | self.api_key = api_key or self.config.get('api_key') 26 | if not self.api_key: 27 | raise ValueError("Perplexity API key not provided and not found in config") 28 | 29 | self.base_url = "https://api.perplexity.ai/chat/completions" 30 | self.headers = { 31 | "Authorization": f"Bearer {self.api_key}", 32 | "Content-Type": "application/json" 33 | } 34 | 35 | def _load_config(self, config_path: str) -> Dict[str, Any]: 36 | """Load configuration from yaml file.""" 37 | try: 38 | with open(config_path, 'r') as f: 39 | config = yaml.safe_load(f) 40 | return config.get('perplexity', {}) 41 | except Exception as e: 42 | print(f"Warning: Could not load config file: {e}") 43 | return {} 44 | 45 | def query(self, 46 | question: str, 47 | system_prompt: str = "Be precise and concise.", 48 | temperature: float | None = None, 49 | model: str | None = None, 50 | max_output_tokens: int | None = 4096) -> PerplexityResponse: 51 | """Send a single query to Perplexity API. 52 | 53 | Args: 54 | question: The question to ask 55 | system_prompt: System prompt to guide the response 56 | temperature: Temperature for response generation (0.0-1.0) 57 | model: Model to use for generation 58 | max_output_tokens: Maximum tokens in response 59 | 60 | Returns: 61 | PerplexityResponse containing the response content and raw API response 62 | 63 | Raises: 64 | requests.exceptions.RequestException: If API request fails 65 | ValueError: If API response is invalid 66 | """ 67 | payload = { 68 | "model": model or self.config.get('model', 'sonar'), 69 | "messages": [ 70 | { 71 | "role": "system", 72 | "content": system_prompt 73 | }, 74 | { 75 | "role": "user", 76 | "content": question 77 | } 78 | ], 79 | "temperature": temperature or self.config.get('temperature', 0.1), 80 | "max_tokens": max_output_tokens or self.config.get('max_output_tokens', 200), 81 | "top_p": 0.9, 82 | "return_images": False, 83 | "return_related_questions": False 84 | } 85 | 86 | response = requests.post(self.base_url, json=payload, headers=self.headers) 87 | response.raise_for_status() 88 | 89 | response_data = response.json() 90 | if "choices" not in response_data or not response_data["choices"]: 91 | raise ValueError("Invalid API response: missing choices") 92 | 93 | content = response_data["choices"][0].get("message", {}).get("content", "") 94 | if not content: 95 | raise ValueError("Invalid API response: missing content") 96 | 97 | return PerplexityResponse(content=content, raw_response=response_data) 98 | 99 | def batch_query(self, 100 | questions: List[str], 101 | system_prompt: str = "Be precise and concise.", 102 | temperature: float | None = None, 103 | model: str | None = None, 104 | max_output_tokens: int | None = None) -> List[PerplexityResponse]: 105 | """Send multiple queries to Perplexity API. 106 | 107 | Args: 108 | questions: List of questions to ask 109 | system_prompt: System prompt to guide the responses 110 | temperature: Temperature for response generation (0.0-1.0) 111 | model: Model to use for generation 112 | max_output_tokens: Maximum tokens in response 113 | 114 | Returns: 115 | List of PerplexityResponse objects 116 | """ 117 | responses = [] 118 | for question in questions: 119 | try: 120 | response = self.query( 121 | question=question, 122 | system_prompt=system_prompt, 123 | temperature=temperature, 124 | model=model, 125 | max_output_tokens=max_output_tokens 126 | ) 127 | responses.append(response) 128 | except Exception as e: 129 | # If a query fails, add None to maintain order with input questions 130 | print(f"Error querying Perplexity API: {str(e)}") 131 | responses.append(None) 132 | 133 | return responses -------------------------------------------------------------------------------- /src/agent/verifier.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | 3 | from typing import Optional, List 4 | from .base import BaseAgent 5 | 6 | 7 | class Verifier(BaseAgent): 8 | """Agent responsible for verifying the quality of generated docstrings.""" 9 | 10 | def __init__(self, config_path: Optional[str] = None): 11 | """Initialize the Verifier agent. 12 | 13 | Args: 14 | config_path: Optional path to the configuration file 15 | """ 16 | super().__init__("Verifier", config_path=config_path) 17 | self.system_prompt = """You are a Verifier agent responsible for ensuring the quality of generated docstrings. 18 | Your role is to evaluate docstrings from the perspective of a first-time user encountering the code component. 19 | 20 | Analysis Process: 21 | 1. First read the code component as if you're seeing it for the first time 22 | 2. Read the docstring and analyze how well it helps you understand the code 23 | 3. Evaluate if the docstring provides the right level of abstraction and information 24 | 25 | Verification Criteria: 26 | 1. Information Value: 27 | - Identify parts that merely repeat the code without adding value 28 | - Flag docstrings that state the obvious without providing insights 29 | - Check if explanations actually help understand the purpose and usage 30 | 31 | 2. Appropriate Detail Level: 32 | - Flag overly detailed technical explanations of implementation 33 | - Ensure focus is on usage and purpose, not line-by-line explanation 34 | - Check if internal implementation details are unnecessarily exposed 35 | 36 | 3. Completeness Check: 37 | - Verify all required sections are present (summary, args, returns, etc.) 38 | - Check if each section provides meaningful information 39 | - Ensure critical usage information is not missing 40 | 41 | Output Format: 42 | Your analysis must include: 43 | 1. true/false 44 | - Indicates if docstring needs improvement 45 | 46 | 2. If revision needed: 47 | true/false 48 | - Indicates if additional context is required for improvement 49 | - Keep in mind that collecting context is very expensive and may fail, so only use it when absolutely necessary 50 | 51 | 3. Based on MORE_CONTEXT, provide suggestions at the end of your response: 52 | If true: 53 | explain why and what specific context is needed 54 | 55 | If false: 56 | specific improvement suggestions 57 | 58 | Do not generate other things after or . 59 | """ 60 | self.add_to_memory("system", self.system_prompt) 61 | 62 | def process( 63 | self, 64 | focal_component: str, 65 | docstring: str, 66 | context: str = "" 67 | ) -> str: 68 | """Verify the quality of a generated docstring. 69 | 70 | Args: 71 | instruction: The original instruction for docstring generation 72 | focal_component: The code component with the docstring 73 | component_type: The type of the code component 74 | docstring: The generated docstring to verify 75 | context: The context used to generate the docstring 76 | 77 | Returns: 78 | List of VerificationFeedback objects for each aspect that needs improvement 79 | """ 80 | task_description = f""" 81 | Context Used: 82 | {context if context else 'No context was used.'} 83 | 84 | Verify the quality of the following docstring for the following Code Component: 85 | 86 | Code Component: 87 | {focal_component} 88 | 89 | Generated Docstring: 90 | {docstring} 91 | 92 | """ 93 | self.add_to_memory("user", task_description) 94 | 95 | full_response = self.generate_response() 96 | return full_response 97 | -------------------------------------------------------------------------------- /src/agent/workflow.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import Optional 3 | from pathlib import Path 4 | from .orchestrator import Orchestrator 5 | from .reader import CodeComponentType 6 | 7 | def generate_docstring( 8 | repo_path: str, 9 | file_path: str, 10 | focal_component: str, 11 | component_type: CodeComponentType, 12 | instruction: Optional[str] = None 13 | ) -> str: 14 | """Generate a high-quality docstring for a code component using the multi-agent system. 15 | 16 | Args: 17 | repo_path: Path to the repository containing the code 18 | file_path: Path to the file containing the component 19 | focal_component: The code component needing a docstring 20 | component_type: The type of the code component (function, method, or class) 21 | instruction: Optional specific instructions for docstring generation 22 | 23 | Returns: 24 | The generated and verified docstring 25 | 26 | Raises: 27 | FileNotFoundError: If the repository or file path doesn't exist 28 | ValueError: If the component type is invalid 29 | """ 30 | # Validate inputs 31 | repo_path = str(Path(repo_path).resolve()) 32 | file_path = str(Path(file_path).resolve()) 33 | 34 | if not Path(repo_path).exists(): 35 | raise FileNotFoundError(f"Repository path does not exist: {repo_path}") 36 | if not Path(file_path).exists(): 37 | raise FileNotFoundError(f"File path does not exist: {file_path}") 38 | 39 | # Use default instruction if none provided 40 | if instruction is None: 41 | instruction = """Generate a comprehensive and helpful docstring that includes: 42 | 1. A clear description of what the component does 43 | 2. All parameters and their types 44 | 3. Return value and type 45 | 4. Any exceptions that may be raised 46 | 5. Usage examples where appropriate 47 | The docstring should follow PEP 257 style guidelines.""" 48 | 49 | # Create orchestrator and generate docstring 50 | orchestrator = Orchestrator(repo_path) 51 | return orchestrator.process( 52 | instruction=instruction, 53 | focal_component=focal_component, 54 | component_type=component_type, 55 | file_path=file_path 56 | ) -------------------------------------------------------------------------------- /src/data/parse/repo_tree.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # Copyright (c) Meta Platforms, Inc. and affiliates 3 | import os 4 | import argparse 5 | from pathlib import Path 6 | import json 7 | from typing import Dict, List, Optional 8 | 9 | class ProjectStructureGenerator: 10 | def __init__(self, ignore_patterns: List[str] = None): 11 | self.ignore_patterns = ignore_patterns or [ 12 | '.git', '__pycache__', '.pytest_cache', 13 | '.env', 'venv', 'node_modules', '.DS_Store', 14 | '*.pyc', '*.pyo', '*.pyd', '.Python', '*.so' 15 | ] 16 | 17 | def should_ignore(self, path: str) -> bool: 18 | """Check if the path should be ignored based on patterns.""" 19 | path_obj = Path(path) 20 | return any( 21 | path_obj.match(pattern) or 22 | any(parent.match(pattern) for parent in path_obj.parents) 23 | for pattern in self.ignore_patterns 24 | ) 25 | 26 | def generate_structure(self, root_path: str, max_depth: Optional[int] = None) -> Dict: 27 | """Generate a hierarchical structure of the project.""" 28 | root_path = os.path.abspath(root_path) 29 | root_name = os.path.basename(root_path) 30 | 31 | def explore_directory(current_path: str, current_depth: int = 0) -> Dict: 32 | if max_depth is not None and current_depth > max_depth: 33 | return {"type": "directory", "name": os.path.basename(current_path), "truncated": True} 34 | 35 | structure = { 36 | "type": "directory", 37 | "name": os.path.basename(current_path), 38 | "contents": [] 39 | } 40 | 41 | try: 42 | for item in sorted(os.listdir(current_path)): 43 | item_path = os.path.join(current_path, item) 44 | 45 | if self.should_ignore(item_path): 46 | continue 47 | 48 | if os.path.isfile(item_path): 49 | file_info = { 50 | "type": "file", 51 | "name": item, 52 | "extension": os.path.splitext(item)[1][1:] or "none" 53 | } 54 | structure["contents"].append(file_info) 55 | elif os.path.isdir(item_path): 56 | subdir = explore_directory(item_path, current_depth + 1) 57 | if subdir.get("contents") or not subdir.get("truncated"): 58 | structure["contents"].append(subdir) 59 | 60 | except PermissionError: 61 | structure["error"] = "Permission denied" 62 | 63 | return structure 64 | 65 | return explore_directory(root_path) 66 | 67 | def format_structure(self, structure: Dict, indent: int = 0) -> str: 68 | """Format the structure in a hierarchical text format.""" 69 | output = [] 70 | prefix = "│ " * (indent - 1) + "├── " if indent > 0 else "" 71 | 72 | if structure.get("truncated"): 73 | output.append(f"{prefix}{structure['name']} [...]") 74 | return "\n".join(output) 75 | 76 | output.append(f"{prefix}{structure['name']}/") 77 | 78 | if "contents" in structure: 79 | for i, item in enumerate(structure["contents"]): 80 | is_last = i == len(structure["contents"]) - 1 81 | if item["type"] == "file": 82 | item_prefix = "│ " * indent + ("└── " if is_last else "├── ") 83 | output.append(f"{item_prefix}{item['name']}") 84 | else: 85 | output.append(self.format_structure(item, indent + 1)) 86 | 87 | return "\n".join(output) 88 | 89 | def main(): 90 | parser = argparse.ArgumentParser( 91 | description="Generate a project structure in LLM-friendly format" 92 | ) 93 | parser.add_argument( 94 | "path", 95 | nargs="?", 96 | default=".", 97 | help="Path to the project directory (default: current directory)" 98 | ) 99 | parser.add_argument( 100 | "--max-depth", 101 | type=int, 102 | help="Maximum depth to traverse (default: no limit)" 103 | ) 104 | parser.add_argument( 105 | "--output", 106 | choices=["text", "json"], 107 | default="text", 108 | help="Output format (default: text)" 109 | ) 110 | parser.add_argument( 111 | "--ignore", 112 | nargs="+", 113 | help="Additional patterns to ignore" 114 | ) 115 | 116 | args = parser.parse_args() 117 | 118 | generator = ProjectStructureGenerator() 119 | if args.ignore: 120 | generator.ignore_patterns.extend(args.ignore) 121 | 122 | structure = generator.generate_structure(args.path, args.max_depth) 123 | 124 | if args.output == "json": 125 | print(json.dumps(structure, indent=2)) 126 | else: 127 | print(generator.format_structure(structure)) 128 | 129 | if __name__ == "__main__": 130 | main() -------------------------------------------------------------------------------- /src/dependency_analyzer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Dependency analyzer module for building and processing import dependency graphs 4 | between Python code components. 5 | """ 6 | 7 | from .ast_parser import CodeComponent, DependencyParser 8 | from .topo_sort import topological_sort, resolve_cycles, build_graph_from_components, dependency_first_dfs 9 | 10 | __all__ = [ 11 | 'CodeComponent', 12 | 'DependencyParser', 13 | 'topological_sort', 14 | 'resolve_cycles', 15 | 'build_graph_from_components', 16 | 'dependency_first_dfs' 17 | ] -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/__init__.cpython-312.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/__init__.cpython-312.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/analyzer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/analyzer.cpython-310.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/ast_parser.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/ast_parser.cpython-310.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/ast_parser.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/ast_parser.cpython-311.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/ast_parser.cpython-312.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/ast_parser.cpython-312.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/graph.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/graph.cpython-310.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/topo_sort.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/topo_sort.cpython-310.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/topo_sort.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/topo_sort.cpython-311.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/topo_sort.cpython-312.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/topo_sort.cpython-312.pyc -------------------------------------------------------------------------------- /src/dependency_analyzer/__pycache__/visualizer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/dependency_analyzer/__pycache__/visualizer.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluate_helpfulness.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Copyright (c) Meta Platforms, Inc. and affiliates 3 | """ 4 | Script to evaluate the helpfulness of docstrings generated by different systems. 5 | 6 | Usage: 7 | conda activate docstringgen 8 | python src/evaluate_helpfulness.py 9 | """ 10 | 11 | import os 12 | import yaml 13 | import argparse 14 | import sys 15 | from pathlib import Path 16 | 17 | # Add the src directory to the path so we can import modules 18 | src_dir = Path(__file__).parent.parent 19 | sys.path.insert(0, str(src_dir)) 20 | 21 | from src.evaluator.helpfulness_evaluator import DocstringHelpfulnessEvaluator 22 | 23 | def main(): 24 | parser = argparse.ArgumentParser(description="Evaluate docstring helpfulness") 25 | parser.add_argument("--data-path", type=str, 26 | default="experiments/eval/results/completeness_evaluation_cleaned.json", 27 | help="Path to the completeness evaluation data") 28 | parser.add_argument("--output-dir", type=str, 29 | default="experiments/eval/results/helpfulness", 30 | help="Directory to store evaluation results") 31 | parser.add_argument("--n-samples", type=int, default=50, 32 | help="Number of components to sample") 33 | parser.add_argument("--seed", type=int, default=42, 34 | help="Random seed for reproducibility") 35 | parser.add_argument("--model", type=str, default=None, 36 | help="LLM model to use (defaults to model in config)") 37 | args = parser.parse_args() 38 | 39 | # Create output directory if it doesn't exist 40 | os.makedirs(args.output_dir, exist_ok=True) 41 | 42 | # Get configuration 43 | config_path = "config/agent_config.yaml" 44 | with open(config_path, 'r') as f: 45 | config = yaml.safe_load(f) 46 | 47 | # Get API key and model from config 48 | api_key = config["llm"]["api_key"] 49 | model = args.model or config["llm"]["model"] 50 | 51 | print(f"Using model: {model}") 52 | print(f"Sampling {args.n_samples} components with seed {args.seed}") 53 | 54 | # Initialize evaluator 55 | evaluator = DocstringHelpfulnessEvaluator( 56 | data_path=args.data_path, 57 | output_dir=args.output_dir, 58 | api_key=api_key, 59 | model=model 60 | ) 61 | 62 | # Run evaluation 63 | results = evaluator.run_evaluation( 64 | n_samples=args.n_samples, 65 | seed=args.seed 66 | ) 67 | 68 | # Print summary 69 | print("\n=== Evaluation Complete ===") 70 | print(f"Results saved to {args.output_dir}") 71 | print(f"Total evaluations: {len(results['results'])}") 72 | 73 | # Calculate average score 74 | scores = [r["score"] for r in results["results"]] 75 | avg_score = sum(scores) / len(scores) if scores else 0 76 | print(f"Overall average score: {avg_score:.2f}") 77 | 78 | # Calculate average by system 79 | systems = evaluator.SYSTEMS 80 | for system in systems: 81 | system_scores = [r["score"] for r in results["results"] if r["system"] == system] 82 | if system_scores: 83 | avg = sum(system_scores) / len(system_scores) 84 | print(f"{system}: {avg:.2f} (n={len(system_scores)})") 85 | 86 | if __name__ == "__main__": 87 | main() -------------------------------------------------------------------------------- /src/evaluator/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from .base import BaseEvaluator 3 | from .completeness import ( # Remove 'evaluators.' from the path 4 | CompletenessEvaluator, 5 | ClassCompletenessEvaluator, 6 | FunctionCompletenessEvaluator 7 | ) 8 | 9 | __all__ = [ 10 | 'BaseEvaluator', 11 | 'CompletenessEvaluator', 12 | 'ClassCompletenessEvaluator', 13 | 'FunctionCompletenessEvaluator' 14 | ] -------------------------------------------------------------------------------- /src/evaluator/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/__init__.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/__init__.cpython-39.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/base.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/base.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/base.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/base.cpython-38.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/base.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/base.cpython-39.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/completeness.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/completeness.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/completeness.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/completeness.cpython-38.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/completeness.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/completeness.cpython-39.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/evaluation_common.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/evaluation_common.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_arguments.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_arguments.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_attributes.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_attributes.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_description.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_description.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_evaluator.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_evaluator.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_evaluator_ablation.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_evaluator_ablation.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_examples.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_examples.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_parameters.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_parameters.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/helpfulness_summary.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/helpfulness_summary.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/__pycache__/segment.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/__pycache__/segment.cpython-39.pyc -------------------------------------------------------------------------------- /src/evaluator/base.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from abc import ABC, abstractmethod 3 | import ast 4 | from typing import Optional, Dict, Any 5 | 6 | class BaseEvaluator(ABC): 7 | """ 8 | Base class for all docstring evaluators. 9 | 10 | This class provides the foundation for implementing various docstring quality 11 | evaluators. Each evaluator should focus on a specific aspect of docstring 12 | quality such as completeness, helpfulness, or redundancy. 13 | 14 | Attributes: 15 | score (float): The evaluation score, ranging from 0 to 1. 16 | name (str): The name of the evaluator. 17 | description (str): A description of what this evaluator checks. 18 | """ 19 | 20 | def __init__(self, name: str, description: str): 21 | self._score: float = 0.0 22 | self._name = name 23 | self._description = description 24 | 25 | @property 26 | def score(self) -> float: 27 | """ 28 | Returns the current evaluation score. 29 | 30 | Returns: 31 | float: A score between 0 and 1 indicating the quality measure. 32 | """ 33 | return self._score 34 | 35 | @score.setter 36 | def score(self, value: float) -> None: 37 | """ 38 | Sets the evaluation score. 39 | 40 | Args: 41 | value (float): The score to set, must be between 0 and 1. 42 | 43 | Raises: 44 | ValueError: If the score is not between 0 and 1. 45 | """ 46 | if not 0 <= value <= 1: 47 | raise ValueError("Score must be between 0 and 1") 48 | self._score = value 49 | 50 | @abstractmethod 51 | def evaluate(self, node: ast.AST) -> float: 52 | """ 53 | Evaluates the quality of a docstring based on specific criteria. 54 | 55 | Args: 56 | node (ast.AST): The AST node containing the docstring to evaluate. 57 | 58 | Returns: 59 | float: The evaluation score between 0 and 1. 60 | """ 61 | pass -------------------------------------------------------------------------------- /src/evaluator/evaluation_common.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """Common utilities and classes for docstring evaluation.""" 3 | 4 | from typing import Dict, Any, List, Optional, Tuple 5 | from dataclasses import dataclass 6 | from enum import Enum 7 | 8 | class ScoreLevel(Enum): 9 | """Defines the possible score levels for docstring evaluation.""" 10 | POOR = 1 11 | FAIR = 2 12 | GOOD = 3 13 | VERY_GOOD = 4 14 | EXCELLENT = 5 15 | 16 | @dataclass 17 | class SummaryEvaluationExample: 18 | """Stores an example of docstring summary evaluation with different quality levels.""" 19 | function_signature: str 20 | summaries: Dict[ScoreLevel, str] 21 | explanations: Dict[ScoreLevel, str] 22 | 23 | @dataclass 24 | class DescriptionEvaluationExample: 25 | """Stores an example of docstring description evaluation with different quality levels.""" 26 | function_signature: str 27 | descriptions: Dict[ScoreLevel, str] 28 | explanations: Dict[ScoreLevel, str] 29 | 30 | @dataclass 31 | class ParameterEvaluationExample: 32 | """Stores an example of docstring parameter evaluation with different quality levels.""" 33 | parameters: Dict[str, str] 34 | quality_examples: Dict[ScoreLevel, Dict[str, str]] 35 | explanations: Dict[ScoreLevel, str] -------------------------------------------------------------------------------- /src/evaluator/helper/__pycache__/context_finder.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/evaluator/helper/__pycache__/context_finder.cpython-310.pyc -------------------------------------------------------------------------------- /src/evaluator/segment.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | import re 3 | 4 | def parse_google_style_docstring(docstring): 5 | """ 6 | A robust parser for Google-style docstrings that have multiple possible 7 | labels for each section. 8 | 9 | For example, any of the lines in EXAMPLE_LABELS indicates the start of the "examples" section. 10 | """ 11 | 12 | # Define all recognized sections. The key is the canonical name (lowercase). 13 | # The value is a set of synonyms (also lowercase). 14 | SECTION_LABELS = { 15 | "summary": {"summary:", "short description:", "brief:", "overview:"}, 16 | "description": {"description:", "desc:", "details:", "detailed description:", "long description:"}, 17 | "parameters": {"parameters:", "params:", "args:", "arguments:", "keyword args:", "keyword arguments:", "**kwargs:"}, 18 | "attributes": {"attributes:", "members:", "member variables:", "instance variables:", "properties:", "vars:", "variables:"}, 19 | "returns": {"returns:", "return:", "return value:", "return values:"}, 20 | "raises": {"raises:", "exceptions:", "throws:", "raise:", "exception:", "throw:"}, 21 | "examples": {"example:", "examples:", "usage:", "usage example:", "usage examples:", "example usage:"}, 22 | } 23 | 24 | # Prepare a dictionary to hold the parsed content for each canonical key 25 | parsed_content = {key: [] for key in SECTION_LABELS.keys()} 26 | 27 | # Split by lines; if docstring uses Windows line endings, .splitlines() handles that gracefully 28 | lines = docstring.strip().splitlines() 29 | 30 | # -- 1) Fallback: no explicit sections at all in the entire docstring -- 31 | # If no recognized label appears anywhere, treat the first line as summary, rest as description. 32 | has_section_labels = False 33 | for line in lines: 34 | line_lower = line.strip().lower() 35 | for labels in SECTION_LABELS.values(): 36 | for label in labels: 37 | if line_lower.startswith(label): 38 | has_section_labels = True 39 | break 40 | if has_section_labels: 41 | break 42 | if has_section_labels: 43 | break 44 | 45 | if len(lines) > 0 and not has_section_labels: 46 | parsed_content["summary"] = [lines[0]] 47 | if len(lines) > 1: 48 | parsed_content["description"] = lines[1:] 49 | # Convert lists to single strings 50 | return {key: "\n".join(value).strip() for key, value in parsed_content.items()} 51 | 52 | # -- 2) Partial Fallback for the first line only -- 53 | # If the first line doesn't match any known label, treat it as summary and then 54 | # switch to "description" until an explicit label is found. 55 | current_section = None # keep track of which section we're in 56 | 57 | first_line = lines[0].strip().lower() if lines else "" 58 | if not any(first_line.startswith(label) for labels in SECTION_LABELS.values() for label in labels): 59 | if lines: 60 | # Save first line as summary 61 | parsed_content["summary"] = [lines[0]] 62 | # Make the current section "description" 63 | current_section = "description" 64 | lines = lines[1:] # We'll handle the rest below 65 | 66 | for line in lines: 67 | # We'll do a trimmed, lowercase version of the line to check for a header 68 | # but keep original_line if you want to preserve original indentation or case. 69 | trimmed_line = line.strip().lower() 70 | 71 | # Check if the trimmed line (minus trailing colon, if present) matches a known section 72 | # We'll also handle any trailing colon, extra spaces, etc. 73 | # e.g. " Parameters: " -> "parameters:" 74 | # We only match a line if it starts exactly with that label. 75 | # If you want more flexible matching (like partial lines), you can adapt this. 76 | matched_section = None 77 | for canonical_name, synonyms in SECTION_LABELS.items(): 78 | # Each synonym might be "parameters:", "args:", etc. 79 | # We'll see if the trimmed_line starts exactly with one of them. 80 | for synonym in synonyms: 81 | # If line starts with the synonym, we treat it as a new section. 82 | # Example: "PARAMETERS:" -> synonyms might contain "parameters:" in lowercase 83 | if trimmed_line.startswith(synonym): 84 | matched_section = canonical_name 85 | # Extract leftover text on the same line, after the label 86 | leftover = line.strip()[len(synonym):].strip() 87 | if leftover: 88 | parsed_content[matched_section].append(leftover) 89 | break 90 | 91 | if matched_section: 92 | break 93 | 94 | # If matched_section is not None, we found a new section header 95 | if matched_section is not None: 96 | # Switch to that section 97 | current_section = matched_section 98 | # No need to append the header line to content - we've already handled any content after the label 99 | else: 100 | # Otherwise, accumulate this line under the current section if we have one 101 | if current_section is not None: 102 | parsed_content[current_section].append(line) 103 | 104 | # Convert list of lines to a single string for each section, 105 | # with consistent line breaks, and strip extra whitespace 106 | for section in parsed_content: 107 | parsed_content[section] = "\n".join(parsed_content[section]).strip() 108 | 109 | return parsed_content 110 | 111 | 112 | # ------------------------------ Example Usage ------------------------------ 113 | if __name__ == "__main__": 114 | sample_docstring = """ 115 | Summary: 116 | Provides a utility for processing and managing data through a structured workflow. 117 | 118 | Description: 119 | This class is designed to facilitate data processing tasks by integrating with the `DataProcessor` class. 120 | It retrieves and manipulates data. 121 | 122 | Parameters: 123 | param1: This is the first parameter. 124 | param2: This is the second parameter. 125 | 126 | Attributes: 127 | data: Stores the current data. 128 | 129 | Example: 130 | ```python 131 | helper = HelperClass() 132 | helper.process_data() 133 | print(helper.data) 134 | ``` 135 | """ 136 | 137 | result = parse_google_style_docstring(sample_docstring) 138 | 139 | # Print out each section 140 | for section_name, content in result.items(): 141 | print("SECTION:", section_name.upper()) 142 | print("CONTENT:\n", content) 143 | print("-" * 40) 144 | -------------------------------------------------------------------------------- /src/visualizer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from .status import StatusVisualizer 3 | from .progress import ProgressVisualizer 4 | 5 | __all__ = ['StatusVisualizer', 'ProgressVisualizer'] -------------------------------------------------------------------------------- /src/visualizer/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/__init__.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/__init__.cpython-311.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/graph_visualizer.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/graph_visualizer.cpython-310.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/progress.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/progress.cpython-310.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/progress.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/progress.cpython-311.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/status.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/status.cpython-310.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/status.cpython-311.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/status.cpython-311.pyc -------------------------------------------------------------------------------- /src/visualizer/__pycache__/web_bridge.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/visualizer/__pycache__/web_bridge.cpython-310.pyc -------------------------------------------------------------------------------- /src/visualizer/status.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | from typing import Dict, Set 3 | from colorama import Fore, Back, Style, init 4 | import sys 5 | import time 6 | import ast 7 | from agent.tool.ast import _get_component_name_from_code 8 | class StatusVisualizer: 9 | """Visualizes the workflow status of DocAssist agents in the terminal.""" 10 | 11 | def __init__(self): 12 | """Initialize the status visualizer.""" 13 | init() # Initialize colorama 14 | self.active_agent = None # Track only the currently active agent 15 | self._agent_art = { 16 | 'reader': [ 17 | "┌─────────┐", 18 | "│ READER │", 19 | "└─────────┘" 20 | ], 21 | 'searcher': [ 22 | "┌─────────┐", 23 | "│SEARCHER │", 24 | "└─────────┘" 25 | ], 26 | 'writer': [ 27 | "┌─────────┐", 28 | "│ WRITER │", 29 | "└─────────┘" 30 | ], 31 | 'verifier': [ 32 | "┌─────────┐", 33 | "│VERIFIER │", 34 | "└─────────┘" 35 | ] 36 | } 37 | self._status_message = "" 38 | self._current_component = "" 39 | self._current_file = "" 40 | 41 | def _clear_screen(self): 42 | """Clear the terminal screen.""" 43 | sys.stdout.write("\033[2J\033[H") 44 | sys.stdout.flush() 45 | 46 | def _get_agent_color(self, agent: str) -> str: 47 | """Get the color for an agent based on its state.""" 48 | return Fore.GREEN if agent == self.active_agent else Fore.WHITE 49 | 50 | def set_current_component(self, focal_component: str, file_path: str): 51 | """Set the current component being processed and display its information. 52 | 53 | Args: 54 | focal_component: The code component being processed 55 | file_path: Relative path to the file containing the component 56 | """ 57 | # Try to extract the component name from the code 58 | try: 59 | self._current_component = _get_component_name_from_code(focal_component) 60 | except: 61 | # If parsing fails, just use a generic name 62 | self._current_component = "unknown component" 63 | 64 | self._current_file = file_path 65 | self._display_component_info() 66 | 67 | def _display_component_info(self): 68 | """Display information about the current component being processed.""" 69 | # print(f"\n{Fore.CYAN}Currently Processing:{Style.RESET_ALL}") 70 | print(f"Component: {self._current_component}") 71 | print(f"File: {self._current_file}\n") 72 | 73 | def update(self, active_agent: str, status_message: str = ""): 74 | """Update the visualization with the current active agent and status. 75 | 76 | Args: 77 | active_agent: Name of the currently active agent 78 | status_message: Current status message to display 79 | """ 80 | self.active_agent = active_agent # Update the single active agent 81 | self._status_message = status_message 82 | self._clear_screen() 83 | 84 | # Build the visualization 85 | lines = [] 86 | 87 | # Add header 88 | # lines.append(f"{Fore.CYAN}DocAssist Workflow Status{Style.RESET_ALL}") 89 | # lines.append("") 90 | 91 | # Display current component info if available 92 | if self._current_component and self._current_file: 93 | lines.append(f"Processing: {self._current_component}") 94 | lines.append(f"File: {self._current_file}") 95 | lines.append("") 96 | 97 | # Input arrow to Reader 98 | # lines.append(" Input") 99 | # lines.append(" ↓") 100 | 101 | # First row: Reader and Searcher with loop 102 | for i in range(3): 103 | line = (f"{self._get_agent_color('reader')}{self._agent_art['reader'][i]}" 104 | f" ←→ " 105 | f"{self._get_agent_color('searcher')}{self._agent_art['searcher'][i]}" 106 | f"{Style.RESET_ALL}") 107 | lines.append(line) 108 | 109 | # Arrow from Reader to Writer 110 | # lines.append(" ↓") 111 | 112 | # Second row: Writer 113 | for i in range(3): 114 | line = (f" {self._get_agent_color('writer')}{self._agent_art['writer'][i]}{Style.RESET_ALL}") 115 | lines.append(line) 116 | 117 | # Arrow from Writer to Verifier 118 | # lines.append(" ↓") 119 | 120 | # Third row: Verifier with output 121 | for i in range(3): 122 | if i == 1: 123 | line = (f" {self._get_agent_color('verifier')}{self._agent_art['verifier'][i]}{Style.RESET_ALL} → Output") 124 | else: 125 | line = (f" {self._get_agent_color('verifier')}{self._agent_art['verifier'][i]}{Style.RESET_ALL}") 126 | lines.append(line) 127 | 128 | # # Feedback arrows from Verifier 129 | # lines.append(" ↑") 130 | # lines.append(" ↗ ↑") 131 | 132 | # Add status message 133 | if self._status_message: 134 | lines.append("") 135 | lines.append(f"{Fore.YELLOW}Status: {self._status_message}{Style.RESET_ALL}") 136 | 137 | # Print the visualization 138 | print("\n".join(lines)) 139 | sys.stdout.flush() 140 | 141 | def reset(self): 142 | """Reset the visualization state.""" 143 | self.active_agent = None 144 | self._status_message = "" 145 | self._current_component = "" 146 | self._current_file = "" 147 | self._clear_screen() -------------------------------------------------------------------------------- /src/web/README.md: -------------------------------------------------------------------------------- 1 | # DocAgent Web Interface 2 | 3 | A real-time web visualization system for the DocAgent docstring generation tool. 4 | 5 | ## Overview 6 | 7 | The DocAgent Web Interface provides a modern, interactive web UI for generating and tracking Python docstring generation. The application visualizes the agent-based docstring generation process in real-time, allowing users to monitor progress, view code structure, track completeness metrics, and manage the configuration. 8 | 9 | ## Features 10 | 11 | - **Configuration Management**: Easily configure all aspects of the docstring generation process (Repository Path, LLM settings, Flow Control, Docstring Options) through a user-friendly web form. Test LLM API connectivity before starting. 12 | - **Real-time Visualization**: Observe the docstring generation process as it happens. 13 | - **Agent Status Tracking**: View which agent (Reader, Searcher, Writer, Verifier) is currently active in the generation workflow via a visual graph. 14 | - **Repository Structure Visualization**: Interactive tree visualization of your Python codebase, highlighting files as they are processed (White: unprocessed, Yellow: processing, Green: completed). 15 | - **Dynamic Progress Tracking**: Real-time progress bars and component completion tracking. 16 | - **Completeness Metrics Visualization**: Visual representation of docstring completeness across your codebase, updated as the generation progresses (visible in the left sidebar). 17 | - **Log Viewer**: Consolidated view of the generation process logs. 18 | - **Process Control**: Start and stop the generation process via UI buttons. 19 | 20 | ## Architecture 21 | 22 | ### Backend 23 | 24 | The web application is built using: 25 | 26 | - **Flask**: Web framework for the backend server 27 | - **Socket.IO**: Real-time bidirectional communication between client and server 28 | - **Eventlet**: Asynchronous networking library for handling concurrent connections 29 | 30 | ### Frontend 31 | 32 | The frontend uses: 33 | 34 | - **Bootstrap 5**: CSS framework for responsive design 35 | - **D3.js**: Data visualization library for interactive repository and agent visualizations 36 | - **Socket.IO Client**: Real-time communication with the backend 37 | - **jQuery**: DOM manipulation and event handling 38 | 39 | ### Directory Structure 40 | 41 | ``` 42 | src/web/ 43 | ├── app.py - Main Flask application 44 | ├── config_handler.py - Handles configuration loading/saving 45 | ├── process_handler.py - Manages the docstring generation process 46 | ├── visualization_handler.py - Handles visualization state management 47 | ├── static/ - Static assets 48 | │ ├── css/ - CSS stylesheets 49 | │ │ └── style.css - Custom styling 50 | │ └── js/ - JavaScript files 51 | │ ├── completeness.js - Completeness visualization 52 | │ ├── config.js - Configuration handling 53 | │ ├── log-handler.js - Log display handling 54 | │ ├── main.js - Main application logic 55 | │ ├── repo-structure.js - Repository structure visualization 56 | │ └── status-visualizer.js - Agent status visualization 57 | └── templates/ - HTML templates 58 | └── index.html - Main application page 59 | ``` 60 | 61 | ## Data Flow 62 | 63 | 1. User configures settings via the web form. 64 | 2. User clicks "Start Generation". 65 | 3. Flask backend spawns a subprocess running the `generate_docstrings.py` script (expected in the project root). 66 | 4. Process output (status updates, logs, metrics) is captured and parsed in real-time by the backend. 67 | 5. Parsed events are emitted via Socket.IO to the frontend. 68 | 6. Frontend components (Agent Status, Repo Structure, Logs, Progress, Completeness) update dynamically based on the received events. 69 | 7. User receives real-time feedback on the generation process. 70 | 8. User can stop the process using the "Stop Generation" button. 71 | 72 | 73 | 74 | ## Usage Guide 75 | 76 | ### 1. Starting the Web Interface 77 | 78 | Run the web application from the project root directory: 79 | 80 | ```bash 81 | python run_web_ui.py 82 | ``` 83 | 84 | By default, the web interface will be available at `http://127.0.0.1:5000`. 85 | 86 | You can customize the host and port: 87 | 88 | ```bash 89 | # Example: Run on port 8080, accessible externally 90 | python run_web_ui.py --host 0.0.0.0 --port 8080 91 | ``` 92 | 93 | ### 2. Configuration 94 | 95 | The initial screen presents configuration options: 96 | 97 | - **Repository Path**: Path to the Python codebase for docstring generation. 98 | - **LLM Configuration**: Settings for the language model (Type, API Key, Model, Temperature, Max Tokens). Use the "Test API" button to verify credentials. 99 | - **Flow Control**: Advanced settings for the generation process. 100 | - **Docstring Options**: Control options like overwriting existing docstrings. 101 | 102 | ### 3. Starting the Generation Process 103 | 104 | 1. Fill in the configuration form accurately. 105 | 2. Click "Start Generation". 106 | 3. The interface will switch to the monitoring/visualization view. 107 | 108 | ### 4. Monitoring the Generation Process 109 | 110 | The visualization interface consists of several panels: 111 | 112 | - **Agent Status Panel**: Shows the current active agent in the workflow graph. 113 | - **Repository Structure Panel**: Displays the interactive codebase tree, highlighting the currently processed file. 114 | - **Logs and Progress Panel**: Shows real-time logs and overall progress. 115 | - **Completeness Panel (Sidebar)**: Shows statistics about docstring completeness. 116 | 117 | ### 5. Stopping the Process 118 | 119 | Click the "Stop Generation" button in the header to terminate the process early. 120 | -------------------------------------------------------------------------------- /src/web/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Web application for docstring generation visualization. 4 | 5 | This module provides a web-based interface for configuring and visualizing 6 | the progress of docstring generation in a Python codebase. 7 | """ 8 | 9 | from .app import create_app 10 | 11 | __all__ = ['create_app'] -------------------------------------------------------------------------------- /src/web/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/__pycache__/app.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/app.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/__pycache__/config_handler.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/config_handler.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/__pycache__/process_handler.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/process_handler.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/__pycache__/run.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/run.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/__pycache__/visualization_handler.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/__pycache__/visualization_handler.cpython-310.pyc -------------------------------------------------------------------------------- /src/web/config_handler.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Configuration handler for the docstring generation web interface. 4 | 5 | This module handles reading, writing, and validating the configuration for 6 | the docstring generation process. 7 | """ 8 | 9 | import os 10 | import yaml 11 | import json 12 | import tempfile 13 | from pathlib import Path 14 | 15 | def get_default_config(): 16 | """ 17 | Get the default configuration from agent_config.yaml. 18 | 19 | Returns: 20 | Dictionary containing the default configuration 21 | """ 22 | default_config_path = Path('config/agent_config.yaml') 23 | 24 | if not default_config_path.exists(): 25 | return { 26 | 'llm': { 27 | 'type': 'claude', 28 | 'api_key': '', 29 | 'model': 'claude-3-5-haiku-latest', 30 | 'temperature': 0.1, 31 | 'max_tokens': 4096 32 | }, 33 | 'flow_control': { 34 | 'max_reader_search_attempts': 2, 35 | 'max_verifier_rejections': 1, 36 | 'status_sleep_time': 1 37 | }, 38 | 'docstring_options': { 39 | 'overwrite_docstrings': False 40 | } 41 | } 42 | 43 | with open(default_config_path, 'r') as f: 44 | config = yaml.safe_load(f) 45 | 46 | return config 47 | 48 | def validate_config(config): 49 | """ 50 | Validate that the configuration has the required fields. 51 | 52 | Args: 53 | config: Dictionary containing the configuration to validate 54 | 55 | Returns: 56 | Tuple of (is_valid, error_message) 57 | """ 58 | required_keys = ['llm', 'flow_control', 'docstring_options'] 59 | 60 | for key in required_keys: 61 | if key not in config: 62 | return False, f"Missing required configuration section: {key}" 63 | 64 | # Check specific required fields in llm section 65 | llm_required = ['type', 'api_key', 'model'] 66 | for key in llm_required: 67 | if key not in config['llm']: 68 | return False, f"Missing required field in llm section: {key}" 69 | 70 | return True, "" 71 | 72 | def save_config(config): 73 | """ 74 | Save the configuration to a temporary file for use by the generation process. 75 | 76 | Args: 77 | config: Dictionary containing the configuration to save 78 | 79 | Returns: 80 | Path to the saved configuration file 81 | """ 82 | # Validate configuration 83 | is_valid, error_message = validate_config(config) 84 | if not is_valid: 85 | raise ValueError(f"Invalid configuration: {error_message}") 86 | 87 | # Create a temporary file 88 | temp_dir = tempfile.gettempdir() 89 | config_file = os.path.join(temp_dir, 'docstring_generator_config.yaml') 90 | 91 | with open(config_file, 'w') as f: 92 | yaml.dump(config, f, default_flow_style=False) 93 | 94 | return config_file -------------------------------------------------------------------------------- /src/web/run.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Entry point for running the docstring generation visualization web application. 4 | 5 | This script creates and starts the Flask application for visualizing the 6 | docstring generation process. 7 | """ 8 | 9 | import os 10 | import sys 11 | import argparse 12 | from pathlib import Path 13 | 14 | from .app import create_app 15 | 16 | def main(): 17 | """ 18 | Parse command line arguments and start the web application. 19 | """ 20 | parser = argparse.ArgumentParser(description='Start the docstring generation visualization web application') 21 | parser.add_argument('--host', default='127.0.0.1', help='Host to bind the server to') 22 | parser.add_argument('--port', type=int, default=5000, help='Port to bind the server to') 23 | parser.add_argument('--debug', action='store_true', help='Run the application in debug mode') 24 | 25 | args = parser.parse_args() 26 | 27 | # Create the Flask application 28 | app, socketio = create_app(debug=args.debug) 29 | 30 | print(f"Starting docstring generation visualization web application on http://{args.host}:{args.port}") 31 | print("Press Ctrl+C to stop the server") 32 | 33 | # Start the server 34 | socketio.run(app, host=args.host, port=args.port, debug=args.debug, allow_unsafe_werkzeug=True) 35 | 36 | if __name__ == '__main__': 37 | # Add the parent directory to the path so we can import the module 38 | sys.path.insert(0, str(Path(__file__).parent.parent.parent)) 39 | main() -------------------------------------------------------------------------------- /src/web/static/assets/meta_logo_white.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web/static/assets/meta_logo_white.png -------------------------------------------------------------------------------- /src/web/static/css/style.css: -------------------------------------------------------------------------------- 1 | /* Copyright (c) Meta Platforms, Inc. and affiliates */ 2 | /* Main layout styles */ 3 | body { 4 | overflow-x: hidden; 5 | } 6 | 7 | .sidebar { 8 | transition: width 0.3s ease; 9 | box-shadow: 2px 0 5px rgba(0, 0, 0, 0.1); 10 | } 11 | 12 | /* Header logo styles */ 13 | .header-logo { 14 | max-height: 30px; 15 | margin-right: 10px; 16 | } 17 | 18 | /* Transition for main content when sidebar changes */ 19 | .main-content-transition { 20 | transition: all 0.3s ease; 21 | } 22 | 23 | /* Status visualizer styles */ 24 | .agent-box { 25 | border: 1px solid #ccc; 26 | border-radius: 5px; 27 | padding: 10px; 28 | margin-bottom: 10px; 29 | text-align: center; 30 | transition: all 0.3s ease; 31 | } 32 | 33 | .agent-box.active { 34 | border-color: #198754; 35 | box-shadow: 0 0 5px rgba(25, 135, 84, 0.5); 36 | background-color: rgba(25, 135, 84, 0.1); 37 | } 38 | 39 | .agent-box h3 { 40 | margin-top: 5px; 41 | font-size: 1.2rem; 42 | } 43 | 44 | .component-info { 45 | margin-top: 20px; 46 | padding: 10px; 47 | background-color: #f8f9fa; 48 | border-radius: 5px; 49 | border-left: 3px solid #007bff; 50 | } 51 | 52 | /* Agent workflow visualization styles */ 53 | #agent-workflow { 54 | min-height: 200px; 55 | } 56 | 57 | .workflow-node circle { 58 | fill: #ffffff; /* White background by default */ 59 | stroke: #6c757d; 60 | stroke-width: 1.5px; 61 | transition: all 0.3s ease; 62 | } 63 | 64 | .workflow-node.active circle { 65 | fill: #198754; /* Green background when active */ 66 | stroke: #0d6efd; 67 | stroke-width: 2px; 68 | } 69 | 70 | .workflow-link { 71 | stroke: #adb5bd; 72 | stroke-width: 2px; 73 | fill: none; 74 | marker-end: url(#arrowhead); 75 | } 76 | 77 | .workflow-label { 78 | font-size: 12px; 79 | text-anchor: middle; 80 | dominant-baseline: middle; 81 | fill: #212529; 82 | pointer-events: none; 83 | transition: all 0.3s ease; 84 | } 85 | 86 | .workflow-node.active .workflow-label { 87 | fill: #fff; 88 | font-weight: bold; 89 | } 90 | 91 | .workflow-text-label { 92 | font-size: 14px; 93 | text-anchor: middle; 94 | dominant-baseline: middle; 95 | fill: #666; 96 | font-weight: bold; 97 | } 98 | 99 | /* Repository structure styles */ 100 | .repo-node { 101 | cursor: pointer; 102 | transition: all 0.2s ease; 103 | } 104 | 105 | .repo-node:hover { 106 | filter: brightness(0.9); 107 | } 108 | 109 | .repo-node-label { 110 | font-size: 0.9rem; 111 | overflow: hidden; 112 | text-overflow: ellipsis; 113 | white-space: nowrap; 114 | } 115 | 116 | .repo-node-complete { 117 | fill: #198754; /* Green */ 118 | } 119 | 120 | .repo-node-in-progress { 121 | fill: #ffc107; /* Yellow */ 122 | } 123 | 124 | .repo-node-not-started { 125 | fill: #f8f9fa; /* Light grey */ 126 | } 127 | 128 | .repo-node-focus { 129 | stroke: #dc3545; /* Red */ 130 | stroke-width: 2; 131 | } 132 | 133 | /* Log container styles */ 134 | #log-container { 135 | font-family: monospace; 136 | font-size: 0.85rem; 137 | line-height: 1.5; 138 | background-color: #212529; 139 | color: #f8f9fa; 140 | border-radius: 5px; 141 | height: 250px; 142 | max-height: 250px; 143 | } 144 | 145 | .log-line { 146 | margin-bottom: 2px; 147 | white-space: pre-wrap; 148 | word-break: break-word; 149 | } 150 | 151 | .log-info { 152 | color: #f8f9fa; 153 | } 154 | 155 | .log-warning { 156 | color: #ffc107; 157 | } 158 | 159 | .log-error { 160 | color: #dc3545; 161 | } 162 | 163 | .log-debug { 164 | color: #6c757d; 165 | } 166 | 167 | /* Completeness table styles */ 168 | .completeness-table { 169 | font-size: 0.9rem; 170 | } 171 | 172 | .progress-cell { 173 | width: 100px; 174 | } 175 | 176 | .progress-bar-mini { 177 | height: 10px; 178 | margin-top: 5px; 179 | border-radius: 5px; 180 | } 181 | 182 | /* Animation for focus transitions */ 183 | @keyframes pulse { 184 | 0% { 185 | transform: scale(1); 186 | opacity: 1; 187 | } 188 | 50% { 189 | transform: scale(1.05); 190 | opacity: 0.8; 191 | } 192 | 100% { 193 | transform: scale(1); 194 | opacity: 1; 195 | } 196 | } 197 | 198 | .highlight-focus { 199 | animation: pulse 1s; 200 | } -------------------------------------------------------------------------------- /src/web/static/js/completeness.js: -------------------------------------------------------------------------------- 1 | // Copyright (c) Meta Platforms, Inc. and affiliates 2 | /** 3 | * Completeness visualization for the docstring generation web application. 4 | * 5 | * This file provides functions for rendering and updating the completeness 6 | * visualization in the web interface. 7 | */ 8 | 9 | /** 10 | * Update the completeness view with the evaluation results. 11 | * 12 | * @param {Object} completenessData - The completeness evaluation data from the server 13 | */ 14 | function updateCompletenessView(completenessData) { 15 | if (!completenessData || !completenessData.files) { 16 | $('#completeness-data').html(` 17 |
18 | No completeness data available 19 |
20 | `); 21 | return; 22 | } 23 | 24 | // Calculate overall statistics 25 | const totalFiles = completenessData.files.length; 26 | let totalClasses = 0; 27 | let totalClassesWithDocs = 0; 28 | let totalFunctions = 0; 29 | let totalFunctionsWithDocs = 0; 30 | 31 | completenessData.files.forEach(file => { 32 | if (file.classes) { 33 | totalClasses += file.classes.length; 34 | totalClassesWithDocs += file.classes.filter(c => c.has_docstring).length; 35 | } 36 | if (file.functions) { 37 | totalFunctions += file.functions.length; 38 | totalFunctionsWithDocs += file.functions.filter(f => f.has_docstring).length; 39 | } 40 | }); 41 | 42 | const classCompleteness = totalClasses > 0 ? Math.round((totalClassesWithDocs / totalClasses) * 100) : 0; 43 | const functionCompleteness = totalFunctions > 0 ? Math.round((totalFunctionsWithDocs / totalFunctions) * 100) : 0; 44 | const totalComponents = totalClasses + totalFunctions; 45 | const totalComponentsWithDocs = totalClassesWithDocs + totalFunctionsWithDocs; 46 | const overallCompleteness = totalComponents > 0 ? Math.round((totalComponentsWithDocs / totalComponents) * 100) : 0; 47 | 48 | // Create the HTML for the completeness view 49 | let html = ` 50 |
51 |
Overall Completeness: ${overallCompleteness}%
52 |
53 |
${overallCompleteness}%
54 |
55 |
56 |
57 | Classes: ${totalClassesWithDocs}/${totalClasses} (${classCompleteness}%) 58 |
59 |
60 | Functions: ${totalFunctionsWithDocs}/${totalFunctions} (${functionCompleteness}%) 61 |
62 |
63 |
64 | 65 |
Files (${totalFiles})
66 |
67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | `; 78 | 79 | // Sort files by completeness (lowest first) 80 | const sortedFiles = [...completenessData.files].sort((a, b) => { 81 | const aTotal = (a.classes?.length || 0) + (a.functions?.length || 0); 82 | const aWithDocs = (a.classes?.filter(c => c.has_docstring).length || 0) + 83 | (a.functions?.filter(f => f.has_docstring).length || 0); 84 | const aPercentage = aTotal > 0 ? (aWithDocs / aTotal) : 1; 85 | 86 | const bTotal = (b.classes?.length || 0) + (b.functions?.length || 0); 87 | const bWithDocs = (b.classes?.filter(c => c.has_docstring).length || 0) + 88 | (b.functions?.filter(f => f.has_docstring).length || 0); 89 | const bPercentage = bTotal > 0 ? (bWithDocs / bTotal) : 1; 90 | 91 | return aPercentage - bPercentage; 92 | }); 93 | 94 | // Add rows for each file 95 | sortedFiles.forEach(file => { 96 | const classes = file.classes || []; 97 | const functions = file.functions || []; 98 | const classesWithDocs = classes.filter(c => c.has_docstring).length; 99 | const functionsWithDocs = functions.filter(f => f.has_docstring).length; 100 | const totalInFile = classes.length + functions.length; 101 | const totalWithDocsInFile = classesWithDocs + functionsWithDocs; 102 | const fileCompleteness = totalInFile > 0 ? Math.round((totalWithDocsInFile / totalInFile) * 100) : 100; 103 | 104 | // Determine the row color based on completeness 105 | let rowClass = ''; 106 | if (fileCompleteness === 100) { 107 | rowClass = 'table-success'; 108 | } else if (fileCompleteness >= 50) { 109 | rowClass = 'table-warning'; 110 | } else { 111 | rowClass = 'table-danger'; 112 | } 113 | 114 | html += ` 115 | 116 | 117 | 118 | 119 | 125 | 126 | `; 127 | }); 128 | 129 | html += ` 130 | 131 |
FileClassesFunctionsCompleteness
${file.file.split('/').pop()}${classesWithDocs}/${classes.length}${functionsWithDocs}/${functions.length} 120 |
121 |
122 |
123 | ${fileCompleteness}% 124 |
132 |
133 | `; 134 | 135 | // Update the completeness data container 136 | $('#completeness-data').html(html); 137 | } -------------------------------------------------------------------------------- /src/web/static/js/config.js: -------------------------------------------------------------------------------- 1 | // Copyright (c) Meta Platforms, Inc. and affiliates 2 | /** 3 | * Configuration handling for the docstring generation web application. 4 | * 5 | * This file provides functions for loading and saving configuration for the 6 | * docstring generation process. 7 | */ 8 | 9 | /** 10 | * Load the default configuration from the server. 11 | */ 12 | function loadDefaultConfig() { 13 | $.ajax({ 14 | url: '/api/default_config', 15 | type: 'GET', 16 | success: function(config) { 17 | applyConfigToForm(config); 18 | }, 19 | error: function(xhr, status, error) { 20 | console.error('Error loading default configuration:', error); 21 | showMessage('warning', 'Failed to load default configuration. Using fallback values.'); 22 | } 23 | }); 24 | } 25 | 26 | /** 27 | * Apply a configuration object to the form inputs. 28 | * 29 | * @param {Object} config - The configuration object to apply 30 | */ 31 | function applyConfigToForm(config) { 32 | // Set LLM configuration 33 | if (config.llm) { 34 | $('#llm-type').val(config.llm.type || 'claude'); 35 | $('#llm-api-key').val(config.llm.api_key || ''); 36 | $('#llm-model').val(config.llm.model || 'claude-3-5-haiku-latest'); 37 | $('#llm-temperature').val(config.llm.temperature || 0.1); 38 | $('#llm-max-tokens').val(config.llm.max_tokens || 4096); 39 | } 40 | 41 | // Set flow control configuration 42 | if (config.flow_control) { 43 | $('#max-reader-search-attempts').val(config.flow_control.max_reader_search_attempts || 2); 44 | $('#max-verifier-rejections').val(config.flow_control.max_verifier_rejections || 1); 45 | $('#status-sleep-time').val(config.flow_control.status_sleep_time || 1); 46 | } 47 | 48 | // Set docstring options 49 | if (config.docstring_options) { 50 | $('#overwrite-docstrings').prop('checked', config.docstring_options.overwrite_docstrings || false); 51 | } 52 | } 53 | 54 | /** 55 | * Build a configuration object from the form inputs. 56 | * 57 | * @returns {Object} The configuration object 58 | */ 59 | function buildConfigFromForm() { 60 | return { 61 | llm: { 62 | type: $('#llm-type').val(), 63 | api_key: $('#llm-api-key').val(), 64 | model: $('#llm-model').val(), 65 | temperature: parseFloat($('#llm-temperature').val()), 66 | max_tokens: parseInt($('#llm-max-tokens').val()) 67 | }, 68 | flow_control: { 69 | max_reader_search_attempts: parseInt($('#max-reader-search-attempts').val()), 70 | max_verifier_rejections: parseInt($('#max-verifier-rejections').val()), 71 | status_sleep_time: parseFloat($('#status-sleep-time').val()) 72 | }, 73 | docstring_options: { 74 | overwrite_docstrings: $('#overwrite-docstrings').is(':checked') 75 | } 76 | }; 77 | } -------------------------------------------------------------------------------- /src/web/static/js/log-handler.js: -------------------------------------------------------------------------------- 1 | // Copyright (c) Meta Platforms, Inc. and affiliates 2 | /** 3 | * Log message handler for the docstring generation web application. 4 | * 5 | * This file provides functions for displaying and managing log messages 6 | * in the web interface. 7 | */ 8 | 9 | // Maximum number of log lines to keep in the UI 10 | const MAX_LOG_LINES = 5000; 11 | 12 | /** 13 | * Add a log message to the log container. 14 | * 15 | * @param {string} level - The log level (info, warning, error, debug) 16 | * @param {string} message - The log message to display 17 | */ 18 | function addLogMessage(level, message) { 19 | // Create a CSS class based on the log level 20 | let logClass = 'log-info'; 21 | switch (level.toLowerCase()) { 22 | case 'warning': 23 | case 'warn': 24 | logClass = 'log-warning'; 25 | break; 26 | case 'error': 27 | case 'critical': 28 | logClass = 'log-error'; 29 | break; 30 | case 'debug': 31 | logClass = 'log-debug'; 32 | break; 33 | } 34 | 35 | // Create the log line element 36 | const logLine = $(`
`); 37 | logLine.text(message); 38 | 39 | // Add the log line to the log content 40 | $('#log-content').append(logLine); 41 | 42 | // Trim log lines if necessary 43 | const logLines = $('#log-content .log-line'); 44 | if (logLines.length > MAX_LOG_LINES) { 45 | // Remove the oldest lines 46 | logLines.slice(0, logLines.length - MAX_LOG_LINES).remove(); 47 | } 48 | 49 | // Scroll to the bottom of the log container 50 | const logContainer = $('#log-container'); 51 | logContainer.scrollTop(logContainer[0].scrollHeight); 52 | } -------------------------------------------------------------------------------- /src/web_eval/README.md: -------------------------------------------------------------------------------- 1 | # DocAgent - Docstring Evaluation System 2 | 3 | A web application for evaluating the quality of Python docstrings in your codebase, providing objective metrics and actionable feedback. 4 | 5 | 6 | ## Overview 7 | 8 | DocAgentis a powerful tool that analyzes Python docstrings in a repository and evaluates them based on two key metrics: 9 | 10 | 1. **Completeness**: Automatically checks if docstrings contain all required components (summary, description, arguments, returns, etc.) 11 | 2. **Helpfulness**: Uses LLM-based evaluation to assess how helpful and informative each docstring component is on a scale of 1-5 12 | 13 | The system provides an intuitive web interface for configuring evaluation settings, viewing results, and getting actionable feedback to improve your codebase documentation. 14 | 15 | ## Features 16 | 17 | - **Configuration Interface**: User-friendly setup for LLM API (OpenAI or Claude) and repository path 18 | - **API Connection Testing**: Verify API credentials before running evaluations 19 | - **Automated Completeness Evaluation**: Scan all Python files in a repository to check for required docstring components 20 | - **Interactive Results Dashboard**: View completeness scores for all classes and functions with detailed breakdowns 21 | - **On-demand Helpfulness Assessment**: Use LLM-powered evaluation for specific docstring components 22 | - **Visual Status Indicators**: Clear visual feedback for required vs. optional components and their quality 23 | - **Component-specific Evaluations**: Different criteria for evaluating summaries, descriptions, parameters, etc. 24 | - **Refresh Functionality**: Re-run evaluation after making code changes 25 | - **Detailed Explanations**: Get specific feedback on why a component received its score and how to improve it 26 | 27 | ## System Architecture 28 | 29 | DocAgent's web evaluation system consists of several key components: 30 | 31 | ``` 32 | src/web_eval/ 33 | │ 34 | ├── app.py # Main Flask application 35 | ├── helpers.py # Utility functions (parsing, extraction, etc.) 36 | ├── requirements.txt # Python dependencies 37 | ├── start_server.sh # Convenience script for starting the server 38 | ├── test_docstring_parser.py # Tests for the docstring parser 39 | │ 40 | ├── templates/ # HTML templates 41 | │ ├── index.html # Configuration page 42 | │ └── results.html # Results display page 43 | │ 44 | └── static/ # Static assets 45 | ├── css/ # CSS stylesheets 46 | ├── js/ # JavaScript files 47 | └── assets/ # Images and other assets 48 | ``` 49 | 50 | The system follows a Model-View-Controller architecture: 51 | 52 | - **Model**: Evaluation logic in the imported evaluator modules and parsing functions in helpers.py 53 | - **View**: HTML templates with Jinja2 for rendering the UI 54 | - **Controller**: Flask routes in app.py that handle requests and connect the model with views 55 | 56 | The application integrates with two key external components: 57 | 58 | 1. **DocAgent Evaluator Modules**: Core evaluation logic for assessing docstring quality 59 | 2. **LLM APIs**: OpenAI or Anthropic Claude for helpfulness evaluation 60 | 61 | -------------------------------------------------------------------------------- /src/web_eval/__pycache__/app.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/__pycache__/app.cpython-310.pyc -------------------------------------------------------------------------------- /src/web_eval/__pycache__/app.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/__pycache__/app.cpython-38.pyc -------------------------------------------------------------------------------- /src/web_eval/__pycache__/helpers.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/__pycache__/helpers.cpython-310.pyc -------------------------------------------------------------------------------- /src/web_eval/__pycache__/helpers.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/__pycache__/helpers.cpython-38.pyc -------------------------------------------------------------------------------- /src/web_eval/__pycache__/helpers.cpython-39.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/__pycache__/helpers.cpython-39.pyc -------------------------------------------------------------------------------- /src/web_eval/helpers.py: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | """ 3 | Helper functions for the DocAgent web application 4 | """ 5 | 6 | import re 7 | from typing import Tuple, Optional, Dict, List 8 | 9 | def parse_llm_score_from_text(text: str) -> Tuple[int, str]: 10 | """ 11 | Parse score and explanation from LLM response text. 12 | 13 | Args: 14 | text: The raw LLM response text 15 | 16 | Returns: 17 | Tuple containing (score, explanation) 18 | """ 19 | # Try to extract score from tags 20 | score_match = re.search(r'(\d+)', text) 21 | if score_match: 22 | score = int(score_match.group(1)) 23 | else: 24 | # Try looking for the score in various formats 25 | score_patterns = [ 26 | r'score:?\s*(\d+)/5', 27 | r'score:?\s*(\d+)', 28 | r'rating:?\s*(\d+)/5', 29 | r'rating:?\s*(\d+)', 30 | r'(\d+)/5', 31 | r'I would rate this as a (\d+)', 32 | r'I would give this a (\d+)' 33 | ] 34 | 35 | for pattern in score_patterns: 36 | match = re.search(pattern, text, re.IGNORECASE) 37 | if match: 38 | score = int(match.group(1)) 39 | break 40 | else: 41 | # Default score if we can't find one 42 | score = 3 43 | 44 | # Limit score to 1-5 range 45 | score = max(1, min(5, score)) 46 | 47 | # Extract explanation (everything except the score tags) 48 | explanation = re.sub(r'\d+', '', text).strip() 49 | 50 | # If explanation is very long, truncate it 51 | if len(explanation) > 500: 52 | explanation = explanation[:497] + "..." 53 | 54 | return score, explanation 55 | 56 | from typing import Dict 57 | 58 | def parse_google_style_docstring(docstring: str) -> Dict[str, str]: 59 | """ 60 | A robust parser for Google-style docstrings that handles multiple possible 61 | labels for each section. 62 | 63 | Args: 64 | docstring: The docstring to parse 65 | 66 | Returns: 67 | Dictionary with canonical section names as keys and their content as values 68 | """ 69 | # If docstring is empty or None, return empty sections 70 | if not docstring: 71 | return {key: "" for key in ['summary', 'description', 'parameters', 'attributes', 'returns', 'raises', 'examples']} 72 | 73 | # Define all recognized sections. The key is the canonical name (lowercase). 74 | # The value is a set of synonyms (also lowercase). 75 | SECTION_LABELS = { 76 | "summary": {"summary:", "brief:", "overview:"}, 77 | "description": {"description:", "desc:", "details:", "long description:"}, 78 | "parameters": {"parameters:", "params:", "args:", "arguments:", "keyword args:", "keyword arguments:", "**kwargs:"}, 79 | "attributes": {"attributes:", "members:", "member variables:", "instance variables:", "properties:", "vars:", "variables:"}, 80 | "returns": {"returns:", "return:", "return value:", "return values:"}, 81 | "raises": {"raises:", "exceptions:", "throws:", "raise:", "exception:", "throw:"}, 82 | "examples": {"example:", "examples:", "usage:", "usage example:", "usage examples:", "example usage:"}, 83 | } 84 | 85 | # Prepare a dictionary to hold the parsed content for each canonical key 86 | parsed_content = {key: [] for key in SECTION_LABELS.keys()} 87 | 88 | # Split by lines; if docstring uses Windows line endings, .splitlines() handles that gracefully 89 | lines = docstring.strip().splitlines() 90 | 91 | # -- 1) Fallback: no explicit sections at all in the entire docstring -- 92 | # If no recognized label appears anywhere, treat the first line as summary, rest as description. 93 | has_section_labels = False 94 | for line in lines: 95 | line_lower = line.strip().lower() 96 | for labels in SECTION_LABELS.values(): 97 | for label in labels: 98 | if line_lower.startswith(label): 99 | has_section_labels = True 100 | break 101 | if has_section_labels: 102 | break 103 | if has_section_labels: 104 | break 105 | 106 | if len(lines) > 0 and not has_section_labels: 107 | parsed_content["summary"] = [lines[0]] 108 | if len(lines) > 1: 109 | parsed_content["description"] = lines[1:] 110 | # Convert lists to single strings 111 | return {key: "\n".join(value).strip() for key, value in parsed_content.items()} 112 | 113 | # We'll track the current section as we parse line by line 114 | current_section = None 115 | 116 | # -- 2) Partial Fallback for the first line only -- 117 | # If the first line doesn't match any known label, treat it as summary and then 118 | # switch to "description" until an explicit label is found. 119 | first_line = lines[0].strip().lower() if lines else "" 120 | if not any(first_line.startswith(label) for labels in SECTION_LABELS.values() for label in labels): 121 | if lines: 122 | # Save first line as summary 123 | parsed_content["summary"] = [lines[0]] 124 | # Make the current section "description" 125 | current_section = "description" 126 | lines = lines[1:] # We'll handle the rest below 127 | 128 | # -- 3) Main Parsing Loop -- 129 | for line in lines: 130 | trimmed_line = line.strip().lower() 131 | matched_section = None 132 | 133 | # Check if this line begins with a known label (case-insensitive) 134 | # If so, we identify that as a new section. 135 | for canonical_name, synonyms in SECTION_LABELS.items(): 136 | for synonym in synonyms: 137 | if trimmed_line.startswith(synonym): 138 | matched_section = canonical_name 139 | # Extract leftover text on the same line, after the label 140 | leftover = line.strip()[len(synonym):].strip() 141 | if leftover: 142 | parsed_content[matched_section].append(leftover) 143 | break 144 | if matched_section: 145 | break 146 | 147 | if matched_section is not None: 148 | # We found a new section header on this line 149 | current_section = matched_section 150 | # No need to append the header line to content - we've already handled any content after the label 151 | else: 152 | # Otherwise, continue appending lines to the current section 153 | if current_section is not None: 154 | parsed_content[current_section].append(line) 155 | 156 | # -- 4) Convert list of lines to single string, preserving line breaks -- 157 | for section in parsed_content: 158 | parsed_content[section] = "\n".join(parsed_content[section]).strip() 159 | 160 | return parsed_content 161 | 162 | 163 | def extract_docstring_component(docstring: str, component: str) -> Optional[str]: 164 | """ 165 | Extract a specific component from a docstring using the robust parser. 166 | 167 | Args: 168 | docstring: The full docstring text 169 | component: The component to extract (summary, description, etc.) 170 | 171 | Returns: 172 | The extracted component text, or None if not found 173 | """ 174 | if not docstring: 175 | return None 176 | 177 | # Map component name to canonical name used in the parser 178 | component_map = { 179 | 'summary': 'summary', 180 | 'description': 'description', 181 | # 'arguments': 'parameters', 182 | 'params': 'parameters', 183 | 'parameters': 'parameters', 184 | 'attributes': 'attributes', 185 | 'returns': 'returns', 186 | 'raises': 'raises', 187 | 'examples': 'examples' 188 | } 189 | 190 | canonical_component = component_map.get(component.lower(), component.lower()) 191 | 192 | # Parse the docstring 193 | parsed = parse_google_style_docstring(docstring) 194 | 195 | # Return the requested component 196 | if canonical_component in parsed: 197 | return parsed[canonical_component] or None 198 | 199 | return None -------------------------------------------------------------------------------- /src/web_eval/requirements.txt: -------------------------------------------------------------------------------- 1 | flask>=2.0.0 2 | openai>=1.0.0 3 | anthropic>=0.5.0 4 | tabulate>=0.8.0 -------------------------------------------------------------------------------- /src/web_eval/start_server.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) Meta Platforms, Inc. and affiliates 3 | 4 | # Default values 5 | HOST="0.0.0.0" 6 | PORT="8080" 7 | DEBUG="" 8 | 9 | # Show help function 10 | show_help() { 11 | echo "Usage: ./start_server.sh [options]" 12 | echo "" 13 | echo "Options:" 14 | echo " -h, --host HOST Host address to bind to (default: 0.0.0.0)" 15 | echo " -p, --port PORT Port to run the server on (default: 8080)" 16 | echo " -d, --debug Run in debug mode" 17 | echo " --help Show this help message" 18 | echo "" 19 | echo "Examples:" 20 | echo " ./start_server.sh # Run on default host:port (0.0.0.0:8080)" 21 | echo " ./start_server.sh -p 9090 # Run on port 9090" 22 | echo " ./start_server.sh -h 127.0.0.1 # Run on localhost only" 23 | echo " ./start_server.sh -d # Run in debug mode" 24 | echo "" 25 | } 26 | 27 | # Parse command line arguments 28 | while [[ $# -gt 0 ]]; do 29 | case "$1" in 30 | -h|--host) 31 | HOST="$2" 32 | shift 2 33 | ;; 34 | -p|--port) 35 | PORT="$2" 36 | shift 2 37 | ;; 38 | -d|--debug) 39 | DEBUG="--debug" 40 | shift 41 | ;; 42 | --help) 43 | show_help 44 | exit 0 45 | ;; 46 | *) 47 | echo "Unknown option: $1" 48 | show_help 49 | exit 1 50 | ;; 51 | esac 52 | done 53 | 54 | # Display startup message 55 | echo "Starting DocAgent Web Server..." 56 | echo "Host: $HOST" 57 | echo "Port: $PORT" 58 | if [ -n "$DEBUG" ]; then 59 | echo "Mode: DEBUG (not recommended for production)" 60 | else 61 | echo "Mode: Production" 62 | fi 63 | echo "" 64 | 65 | # Run the Flask app with the specified options 66 | python app.py --host "$HOST" --port "$PORT" $DEBUG -------------------------------------------------------------------------------- /src/web_eval/static/assets/meta_logo_white.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/facebookresearch/DocAgent/f27f68574ee18a9a64460229e5abcb4ce77b8c24/src/web_eval/static/assets/meta_logo_white.png -------------------------------------------------------------------------------- /src/web_eval/static/css/style.css: -------------------------------------------------------------------------------- 1 | /* Copyright (c) Meta Platforms, Inc. and affiliates */ 2 | /* DocAgent - Docstring Evaluation System Styles */ 3 | 4 | /* General Styles */ 5 | body { 6 | background-color: #f8f9fa; 7 | } 8 | 9 | .card { 10 | border-radius: 0.5rem; 11 | overflow: hidden; 12 | } 13 | 14 | .card-header { 15 | border-bottom: none; 16 | } 17 | 18 | /* Table Styles */ 19 | .table { 20 | font-size: 0.9rem; 21 | } 22 | 23 | .table th { 24 | font-weight: 600; 25 | } 26 | 27 | .table-responsive { 28 | max-height: 70vh; 29 | overflow-y: auto; 30 | } 31 | 32 | /* Button Styles */ 33 | .evaluate-btn { 34 | font-size: 0.75rem; 35 | padding: 0.2rem 0.5rem; 36 | } 37 | 38 | /* Modal Styles */ 39 | .modal-content { 40 | border-radius: 0.5rem; 41 | overflow: hidden; 42 | } 43 | 44 | .modal-header { 45 | border-bottom: none; 46 | } 47 | 48 | .modal-footer { 49 | border-top: none; 50 | } 51 | 52 | /* Docstring content display */ 53 | pre#docstringContent { 54 | max-height: 300px; 55 | overflow-y: auto; 56 | font-size: 0.9rem; 57 | white-space: pre-wrap; 58 | } 59 | 60 | /* Badges */ 61 | .badge { 62 | font-weight: 500; 63 | padding: 0.35rem 0.65rem; 64 | } 65 | 66 | /* Alert Styles */ 67 | .alert { 68 | border-radius: 0.5rem; 69 | } 70 | 71 | /* Responsive Adjustments */ 72 | @media (max-width: 992px) { 73 | .table { 74 | font-size: 0.8rem; 75 | } 76 | 77 | .evaluate-btn { 78 | font-size: 0.7rem; 79 | padding: 0.15rem 0.4rem; 80 | } 81 | 82 | .badge { 83 | font-size: 0.7rem; 84 | padding: 0.25rem 0.5rem; 85 | } 86 | } 87 | 88 | /* Custom scrollbar */ 89 | ::-webkit-scrollbar { 90 | width: 8px; 91 | height: 8px; 92 | } 93 | 94 | ::-webkit-scrollbar-track { 95 | background: #f1f1f1; 96 | border-radius: 4px; 97 | } 98 | 99 | ::-webkit-scrollbar-thumb { 100 | background: #888; 101 | border-radius: 4px; 102 | } 103 | 104 | ::-webkit-scrollbar-thumb:hover { 105 | background: #555; 106 | } -------------------------------------------------------------------------------- /tool/remove_docstrings.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | # Copyright (c) Meta Platforms, Inc. and affiliates 3 | """ 4 | Tool to remove docstrings from Python files in a repository. 5 | """ 6 | 7 | import os 8 | import ast 9 | import astor 10 | import argparse 11 | from typing import List, Tuple 12 | 13 | 14 | class DocstringRemover(ast.NodeTransformer): 15 | """ 16 | AST NodeTransformer that removes docstrings from classes, methods, and functions. 17 | """ 18 | 19 | def visit_ClassDef(self, node): 20 | """Remove docstrings from class definitions.""" 21 | # Process class body first (recursive) 22 | node = self.generic_visit(node) 23 | 24 | # Remove docstring if present 25 | if (node.body and isinstance(node.body[0], ast.Expr) and 26 | isinstance(node.body[0].value, ast.Str)): 27 | node.body = node.body[1:] 28 | 29 | return node 30 | 31 | def visit_FunctionDef(self, node): 32 | """Remove docstrings from function/method definitions.""" 33 | # Process function body first (recursive) 34 | node = self.generic_visit(node) 35 | 36 | # Remove docstring if present 37 | if (node.body and isinstance(node.body[0], ast.Expr) and 38 | isinstance(node.body[0].value, ast.Str)): 39 | node.body = node.body[1:] 40 | 41 | return node 42 | 43 | def visit_AsyncFunctionDef(self, node): 44 | """Remove docstrings from async function/method definitions.""" 45 | # Process function body first (recursive) 46 | node = self.generic_visit(node) 47 | 48 | # Remove docstring if present 49 | if (node.body and isinstance(node.body[0], ast.Expr) and 50 | isinstance(node.body[0].value, ast.Str)): 51 | node.body = node.body[1:] 52 | 53 | return node 54 | 55 | 56 | def find_python_files(directory: str) -> List[str]: 57 | """Find all Python files in the given directory and its subdirectories.""" 58 | python_files = [] 59 | 60 | for root, _, files in os.walk(directory): 61 | for file in files: 62 | if file.endswith('.py'): 63 | python_files.append(os.path.join(root, file)) 64 | 65 | return python_files 66 | 67 | 68 | def remove_docstrings_from_file(file_path: str, dry_run: bool = False) -> Tuple[bool, str]: 69 | """ 70 | Remove docstrings from a Python file. 71 | 72 | Args: 73 | file_path: Path to the Python file 74 | dry_run: If True, don't actually write changes to file 75 | 76 | Returns: 77 | Tuple of (success, message) 78 | """ 79 | try: 80 | with open(file_path, 'r', encoding='utf-8') as f: 81 | source = f.read() 82 | 83 | # Parse the source code into an AST 84 | tree = ast.parse(source) 85 | 86 | # Remove docstrings 87 | transformer = DocstringRemover() 88 | new_tree = transformer.visit(tree) 89 | 90 | # Generate the modified source code 91 | new_source = astor.to_source(new_tree) 92 | 93 | if not dry_run: 94 | with open(file_path, 'w', encoding='utf-8') as f: 95 | f.write(new_source) 96 | 97 | return True, f"Successfully removed docstrings from {file_path}" 98 | else: 99 | return True, f"Would remove docstrings from {file_path} (dry run)" 100 | 101 | except Exception as e: 102 | return False, f"Error processing {file_path}: {str(e)}" 103 | 104 | 105 | def main(): 106 | parser = argparse.ArgumentParser(description="Remove docstrings from Python files in a repository") 107 | parser.add_argument("directory", help="Directory containing Python files to process") 108 | parser.add_argument("--dry-run", action="store_true", help="Don't actually modify files, just show what would be done") 109 | args = parser.parse_args() 110 | 111 | # Find all Python files 112 | python_files = find_python_files(args.directory) 113 | print(f"Found {len(python_files)} Python files to process") 114 | 115 | # Process each file 116 | success_count = 0 117 | for file_path in python_files: 118 | success, message = remove_docstrings_from_file(file_path, args.dry_run) 119 | print(message) 120 | if success: 121 | success_count += 1 122 | 123 | # Summary 124 | print(f"\nProcessed {len(python_files)} files, {success_count} successful") 125 | 126 | 127 | if __name__ == "__main__": 128 | main() -------------------------------------------------------------------------------- /tool/remove_docstrings.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | # Copyright (c) Meta Platforms, Inc. and affiliates 3 | 4 | # Shell script wrapper for the remove_docstrings.py tool 5 | 6 | set -e 7 | 8 | # Script directory 9 | SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" 10 | 11 | # Show usage 12 | function show_usage { 13 | echo "Usage: $(basename $0) [options] DIRECTORY" 14 | echo "" 15 | echo "Options:" 16 | echo " -h, --help Show this help message" 17 | echo " -d, --dry-run Perform a dry run (no changes are made)" 18 | echo " -b, --backup Create backup files before making changes" 19 | echo "" 20 | echo "Example:" 21 | echo " $(basename $0) ~/my-python-project" 22 | echo " $(basename $0) --dry-run ~/my-python-project" 23 | exit 1 24 | } 25 | 26 | # Parse arguments 27 | DRY_RUN="" 28 | BACKUP=false 29 | DIRECTORY="" 30 | 31 | while [[ $# -gt 0 ]]; do 32 | case $1 in 33 | -h|--help) 34 | show_usage 35 | ;; 36 | -d|--dry-run) 37 | DRY_RUN="--dry-run" 38 | shift 39 | ;; 40 | -b|--backup) 41 | BACKUP=true 42 | shift 43 | ;; 44 | *) 45 | if [[ -z "$DIRECTORY" ]]; then 46 | DIRECTORY="$1" 47 | else 48 | echo "Error: Too many arguments" 49 | show_usage 50 | fi 51 | shift 52 | ;; 53 | esac 54 | done 55 | 56 | # Check if directory is provided 57 | if [[ -z "$DIRECTORY" ]]; then 58 | echo "Error: No directory specified" 59 | show_usage 60 | fi 61 | 62 | # Check if directory exists 63 | if [[ ! -d "$DIRECTORY" ]]; then 64 | echo "Error: Directory does not exist: $DIRECTORY" 65 | exit 1 66 | fi 67 | 68 | # Create backups if requested 69 | if [[ "$BACKUP" = true ]]; then 70 | echo "Creating backups of Python files..." 71 | find "$DIRECTORY" -name "*.py" -type f -exec cp {} {}.bak \; 72 | echo "Backups created with .bak extension" 73 | fi 74 | 75 | # Run the Python script 76 | python3 "$SCRIPT_DIR/remove_docstrings.py" $DRY_RUN "$DIRECTORY" 77 | 78 | echo "Done!" -------------------------------------------------------------------------------- /tool/serve_local_llm.sh: -------------------------------------------------------------------------------- 1 | # Copyright (c) Meta Platforms, Inc. and affiliates 2 | CUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server \ 3 | --model Your-Model-Name \ 4 | --tensor-parallel-size 8 \ 5 | --quantization fp8 \ 6 | --gpu-memory-utilization 0.9 \ 7 | --dtype bfloat16 \ 8 | --host 0.0.0.0 \ 9 | --port 8000 --------------------------------------------------------------------------------