Showing preview only (821K chars total). Download the full file or copy to clipboard to get everything.
Repository: facebookresearch/DocAgent
Branch: main
Commit: f27f68574ee1
Files: 109
Total size: 781.2 KB
Directory structure:
gitextract_774lw2f5/
├── .gitignore
├── CHANGELOG.md
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── INSTALL.md
├── LICENSE
├── README.md
├── 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
│ ├── agent/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── llm/
│ │ │ ├── __init__.py
│ │ │ ├── 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
│ │ │ ├── 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
│ │ ├── ast_parser.py
│ │ └── topo_sort.py
│ ├── evaluate_helpfulness.py
│ ├── evaluator/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── completeness.py
│ │ ├── evaluation_common.py
│ │ ├── helper/
│ │ │ └── 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
│ │ ├── progress.py
│ │ ├── status.py
│ │ └── web_bridge.py
│ ├── web/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── app.py
│ │ ├── config_handler.py
│ │ ├── process_handler.py
│ │ ├── run.py
│ │ ├── static/
│ │ │ ├── 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
│ ├── app.py
│ ├── helpers.py
│ ├── requirements.txt
│ ├── start_server.sh
│ ├── static/
│ │ └── css/
│ │ └── style.css
│ ├── templates/
│ │ ├── index.html
│ │ └── results.html
│ └── test_docstring_parser.py
└── tool/
├── remove_docstrings.py
├── remove_docstrings.sh
└── serve_local_llm.sh
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
config/agent_config.yaml
tool/add_header.sh
================================================
FILE: CHANGELOG.md
================================================
0.0.1 (April 17, 2025)
### First Version
Include web UI, CLI for DocAgent.
================================================
FILE: CODE_OF_CONDUCT.md
================================================
# Code of Conduct
## Our Pledge
In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to make participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, sex characteristics, gender identity and expression,
level of experience, education, socio-economic status, nationality, personal
appearance, race, religion, or sexual identity and orientation.
## Our Standards
Examples of behavior that contributes to creating a positive environment
include:
* Using welcoming and inclusive language
* Being respectful of differing viewpoints and experiences
* Gracefully accepting constructive criticism
* Focusing on what is best for the community
* Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
* The use of sexualized language or imagery and unwelcome sexual attention or
advances
* Trolling, insulting/derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or electronic
address, without explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or
reject comments, commits, code, wiki edits, issues, and other contributions
that are not aligned to this Code of Conduct, or to ban temporarily or
permanently any contributor for other behaviors that they deem inappropriate,
threatening, offensive, or harmful.
## Scope
This Code of Conduct applies within all project spaces, and it also applies when
an individual is representing the project or its community in public spaces.
Examples of representing a project or community include using an official
project e-mail address, posting via an official social media account, or acting
as an appointed representative at an online or offline event. Representation of
a project may be further defined and clarified by project maintainers.
This Code of Conduct also applies outside the project spaces when there is a
reasonable belief that an individual's behavior may have a negative impact on
the project or its community.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported by contacting the project team at <opensource-conduct@meta.com>. All
complaints will be reviewed and investigated and will result in a response that
is deemed necessary and appropriate to the circumstances. The project team is
obligated to maintain confidentiality with regard to the reporter of an incident.
Further details of specific enforcement policies may be posted separately.
Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage], version 1.4,
available at https://www.contributor-covenant.org/version/1/4/code-of-conduct.html
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see
https://www.contributor-covenant.org/faq
================================================
FILE: CONTRIBUTING.md
================================================
# Contributing to DocAgent
We want to make contributing to this project as easy and transparent as
possible.
## Pull Requests
We actively welcome your pull requests.
1. Fork the repo and create your branch from `main`.
2. If you've added code that should be tested, add tests.
3. If you've changed APIs, update the documentation.
4. Ensure the test suite passes.
5. Make sure your code lints.
6. If you haven't already, complete the Contributor License Agreement ("CLA").
## Contributor License Agreement ("CLA")
In order to accept your pull request, we need you to submit a CLA. You only need
to do this once to work on any of Meta's open source projects.
Complete your CLA here: <https://code.facebook.com/cla>
## Issues
We use GitHub issues to track public bugs. Please ensure your description is
clear and has sufficient instructions to be able to reproduce the issue.
Meta has a [bounty program](https://bugbounty.meta.com/) for the safe
disclosure of security bugs. In those cases, please go through the process
outlined on that page and do not file a public issue.
## Coding Style
* 2 spaces for indentation rather than tabs
* 80 character line length
* Use [Black](https://github.com/psf/black) for code formatting.
* Use [Flake8](https://flake8.pycqa.org/en/latest/) for linting.
* Follow [PEP 8](https://www.python.org/dev/peps/pep-0008/) style guidelines.
* Use snake_case for variable and function names.
* Use PascalCase for class names.
* Write docstrings for all public modules, classes, functions, and methods using Google style.
* Use type hints for function signatures.
* 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.
## License
By contributing to DocAgent, you agree that your contributions will be licensed
under the LICENSE file in the root directory of this source tree.
================================================
FILE: INSTALL.md
================================================
# Installation Guide
This guide details how to set up the environment for DocAgent.
## Option 1: Installation with pip (Recommended)
### Basic Installation
To install the basic package with core dependencies:
```bash
# For all dependencies
pip install -e ".[all]"
```
## Development Setup
For development, we recommend installing in editable mode with dev dependencies:
```bash
# Install the package in editable mode with dev dependencies
pip install -e ".[dev]"
# Run tests
pytest
```
## Troubleshooting
### GraphViz Dependencies
For visualization components, you may need to install system-level dependencies for GraphViz:
```bash
# Ubuntu/Debian
sudo apt-get install graphviz graphviz-dev
# CentOS/RHEL
sudo yum install graphviz graphviz-devel
# macOS
brew install graphviz
```
### CUDA Support
If you're using CUDA for accelerated processing, ensure you have the correct CUDA toolkit installed that matches your PyTorch version.
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) Meta Platforms, Inc. and affiliates.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: README.md
================================================
# DocAgent: Agentic Hierarchical Docstring Generation System
<p align="center">
<img src="assets/meta_logo_white.png" width="20%" alt="Meta Logo">
</p>
DocAgent is a system designed to generate high-quality, context-aware docstrings for Python codebases using a multi-agent approach and hierarchical processing.
## Citation
If you use DocAgent in your research, please cite our paper:
```bibtex
@misc{yang2025docagent,
title={DocAgent: A Multi-Agent System for Automated Code Documentation Generation},
author={Dayu Yang and Antoine Simoulin and Xin Qian and Xiaoyi Liu and Yuwei Cao and Zhaopu Teng and Grey Yang},
year={2025},
eprint={2504.08725},
archivePrefix={arXiv},
primaryClass={cs.SE}
}
```
You can find the paper on arXiv: [https://arxiv.org/abs/2504.08725](https://arxiv.org/abs/2504.08725)
## Table of Contents
- [Motivation](#motivation)
- [Methodology](#methodology)
- [Installation](#installation)
- [Components](#components)
- [Configuration](#configuration)
- [Usage](#usage)
- [Running the Evaluation System](#running-the-evaluation-system)
- [Optional: Using a Local LLM](#optional-using-a-local-llm)
## Motivation
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.
## Methodology
DocAgent employs two key strategies:
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.
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.
<img src="assets/system.png" width="100%" alt="System Overview">
For more details on the agentic framework, see the [Agent Component README](./src/agent/README.md).
## Installation
1. Clone the repository:
```bash
git clone <repository_url>
cd DocAgent
```
2. Install the necessary dependencies. It's recommended to use a virtual environment:
```bash
python -m venv venv
source venv/bin/activate # if you use venv, you can also use conda
pip install -e .
```
*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]"`).*
## Components
DocAgent is composed of several key parts:
- **[Core Agent Framework](./src/agent/README.md)**: Implements the multi-agent system (Reader, Searcher, Writer, Verifier, Orchestrator) responsible for the generation logic.
- **[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.*
- **[Generation Web UI](./src/web/README.md)**: A web interface for configuring, running, and monitoring the docstring *generation* process in real-time.
## Configuration
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.
1. **Copy the Example**: An example configuration file is provided at `config/example_config.yaml`. Copy this file to `config/agent_config.yaml`:
```bash
cp config/example_config.yaml config/agent_config.yaml
```
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.
## Usage
You can run the docstring generation process using either the command line or the web UI.
**1. Command Line Interface (CLI)**
This is the primary method for running the generation process directly.
```bash
# Example: Run on a test repo (remove existing docstrings first if desired)
./test/tool/remove_docstrings.sh data/raw_test_repo
python generate_docstrings.py --repo-path data/raw_test_repo
```
Use `python generate_docstrings.py --help` to see available options, such as specifying different configurations or test modes.
**2. Generation Web UI**
The web UI provides a graphical interface to configure, run, and monitor the process.
- Note that when input repo path, always put complete absolute path.
```bash
# Launch the web UI server
python run_web_ui.py --host 0.0.0.0 --port 5000
```
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)).
*Basic SSH Tunneling (if running server remotely):*
```bash
# In your local terminal
ssh -L 5000:localhost:5000 <your_remote_username>@<your_remote_host>
# Then access http://localhost:5000 in your local browser
```
## Running the Evaluation System
DocAgent includes a separate web-based interface for evaluating the quality of generated docstrings.
**1. Running Locally**
To run the evaluation system on your local machine:
```bash
python src/web_eval/app.py
```
Then, access the evaluation UI in your web browser at `http://localhost:5001`.
**2. Running on a Remote Server**
To run the evaluation system on a remote server:
```bash
python src/web_eval/app.py --host 0.0.0.0 --port 5001
```
Then, set up SSH tunneling to access the remote server from your local machine:
```bash
ssh -L 5001:localhost:5001 <your_remote_username>@<your_remote_host>
```
Once the tunnel is established, access the evaluation UI in your local web browser at `http://localhost:5001`.
## Optional: Using a Local LLM
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.
1. **Serve the Local LLM**: Use a tool like `vllm` to serve your model. A convenience script is provided:
```bash
# Ensure vllm is installed: pip install vllm
bash tool/serve_local_llm.sh
```
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`).
2. **Configure DocAgent**: Update your `config/agent_config.yaml` to point to the local LLM API endpoint. You'll typically need to set:
- The `provider` to `openai` (if using an OpenAI-compatible server like vllm's default).
- The `api_base` or equivalent URL parameter to your local server address (e.g., `http://localhost:8000/v1`).
- The `model_name` to the appropriate identifier for your local model.
- Set the `api_key` to `None` or an empty string if no key is required by your local server.
3. **Run DocAgent**: Run the generation process as usual (CLI or Web UI). DocAgent will now send requests to your local LLM.
## License
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
================================================
FILE: config/example_config.yaml
================================================
# Example configuration file for DocAgent
# Copy this file to agent_config.yaml and add your own API keys
# LLM configuration for all agents
llm:
# Choose ONE of the following LLM provider configurations by uncommenting
# Option 1: Claude (Anthropic)
type: "claude"
api_key: "your-anthropic-api-key-here"
model: "claude-3-5-haiku-latest" # Options: claude-3-5-sonnet, claude-3-opus, etc.
temperature: 0.1
max_output_tokens: 4096
max_input_tokens: 100000 # Maximum number of tokens for input context
# Option 2: OpenAI
# type: "openai"
# api_key: "your-openai-api-key-here"
# model: "gpt-4o" # Options: gpt-4o, gpt-4-turbo, gpt-3.5-turbo, etc.
# temperature: 0.1
# max_output_tokens: 4096
# max_input_tokens: 100000
# Option 3: Gemini
# type: "gemini"
# api_key: "your-gemini-api-key-here"
# model: "gemini-1.5-pro"
# temperature: 0.1
# max_output_tokens: 4096
# max_input_tokens: 100000
# Option 4: HuggingFace (for local models)
# type: "huggingface"
# model: "codellama/CodeLlama-34b-Instruct-hf"
# api_base: "http://localhost:8000/v1" # Local API endpoint
# api_key: "EMPTY" # Can be empty for local models
# device: "cuda" # Options: cuda, cpu
# torch_dtype: "float16"
# temperature: 0.1
# max_output_tokens: 4096
# max_input_tokens: 32000
# Rate limit settings for different LLM providers
# These are default values - adjust based on your specific API tier
rate_limits:
# Claude rate limits
claude:
requests_per_minute: 50
input_tokens_per_minute: 20000
output_tokens_per_minute: 8000
input_token_price_per_million: 3.0
output_token_price_per_million: 15.0
# OpenAI rate limits
openai:
requests_per_minute: 500
input_tokens_per_minute: 200000
output_tokens_per_minute: 100000
input_token_price_per_million: 0.15
output_token_price_per_million: 0.60
# Gemini rate limits
gemini:
requests_per_minute: 60
input_tokens_per_minute: 30000
output_tokens_per_minute: 10000
input_token_price_per_million: 0.125
output_token_price_per_million: 0.375
# Flow control parameters
flow_control:
max_reader_search_attempts: 2 # Maximum times reader can call searcher
max_verifier_rejections: 1 # Maximum times verifier can reject a docstring
status_sleep_time: 1 # Time to sleep between status updates (seconds)
# Docstring generation options
docstring_options:
overwrite_docstrings: false # Whether to overwrite existing docstrings (default: false)
# Perplexity API configuration (for web search capability)
perplexity:
api_key: "your-perplexity-api-key-here" # Replace with your actual Perplexity API key
model: "sonar" # Default model
temperature: 0.1
max_output_tokens: 250
================================================
FILE: data/raw_test_repo/README.md
================================================
# Vending Machine Test Repository
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.
## Project Structure
```
test_repo_vm/
├── __init__.py # Main package initialization
├── example.py # Example usage demonstration
├── vending_machine.py # Main vending machine implementation
├── models/ # Data models
│ ├── __init__.py
│ └── product.py # Product class definition
├── payment/ # Payment processing
│ ├── __init__.py
│ └── payment_processor.py # Payment-related classes
└── inventory/ # Inventory management
├── __init__.py
└── inventory_manager.py # Inventory tracking system
```
## Components
### 1. Product Management (`models/product.py`)
- `Product` class with attributes like ID, name, price, quantity, and expiry date
- Methods for checking availability and managing stock
### 2. Payment Processing (`payment/payment_processor.py`)
- Abstract `PaymentMethod` base class for different payment types
- `CashPayment` implementation for handling cash transactions
- `PaymentTransaction` class for tracking payment status
- `PaymentStatus` enum for transaction states
### 3. Inventory Management (`inventory/inventory_manager.py`)
- `InventoryManager` class for product storage and retrieval
- Slot-based product organization
- Stock level tracking
- Product availability checking
### 4. Main Vending Machine (`vending_machine.py`)
- `VendingMachine` class that coordinates all components
- Product selection and purchase workflow
- Payment processing and change calculation
- Exception handling for error cases
## Code Features
This repository demonstrates various Python programming features:
1. **Object-Oriented Design**
- Abstract base classes
- Inheritance
- Encapsulation
- Interface definitions
2. **Modern Python Features**
- Type hints
- Dataclasses
- Enums
- Optional types
- Package organization
3. **Documentation**
- Comprehensive docstrings
- Type annotations
- Code organization
- Exception documentation
4. **Best Practices**
- SOLID principles
- Clean code architecture
- Error handling
- Modular design
## Usage Example
```python
from decimal import Decimal
from vending_machine import VendingMachine
from models.product import Product
# Create a vending machine
vm = VendingMachine()
# Add products
product = Product(
id="COLA001",
name="Cola Classic",
price=1.50,
quantity=10,
category="drinks"
)
vm.inventory.add_product(product, slot=0)
# Insert money
vm.insert_money(Decimal('2.00'))
# Purchase product
product, change = vm.purchase_product(slot=0)
print(f"Purchased: {product.name}")
print(f"Change: ${change:.2f}")
```
## Running the Example
To run the example implementation:
```bash
python example.py
```
This will demonstrate:
1. Creating a vending machine
2. Adding products to inventory
3. Displaying available products
4. Making a purchase
5. Handling change
6. Updating inventory
## Testing Documentation Generation
This repository is structured to test various aspects of documentation generation:
1. **Complex Imports**
- Cross-module dependencies
- Package-level imports
- Relative imports
2. **Documentation Styles**
- Function documentation
- Class documentation
- Module documentation
- Package documentation
3. **Code Complexity**
- Multiple inheritance
- Abstract classes
- Type annotations
- Exception hierarchies
## Requirements
- Python 3.7+
- No external dependencies required
## License
This project is open source and available under the MIT License.
================================================
FILE: data/raw_test_repo/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
"""
Vending Machine Package
A comprehensive vending machine implementation with:
- Product management
- Inventory tracking
- Payment processing
- Transaction handling
"""
================================================
FILE: data/raw_test_repo/example.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from decimal import Decimal
from datetime import datetime, timedelta
from models.product import Item
from vending_machine import Sys, SysErr
def main():
s = Sys()
items = [Item(code='D1', label='Drink1', val=1.5, count=10, grp='d',
exp=datetime.now() + timedelta(days=90)), Item(code='S1', label=
'Snack1', val=1.0, count=15, grp='s', exp=datetime.now() +
timedelta(days=30)), Item(code='S2', label='Snack2', val=2.0, count
=8, grp='s', exp=datetime.now() + timedelta(days=60))]
for i, item in enumerate(items):
s.store.put(item, i)
try:
print('Items:')
for pos, item in s.ls():
print(f'Pos {pos}: {item.label} - ${item.val:.2f}')
pos = 0
print('\nAdding $2.00...')
s.add_money(Decimal('2.00'))
item, ret = s.buy(pos)
print(f'\nBought: {item.label}')
if ret:
print(f'Return: ${ret:.2f}')
print('\nUpdated Items:')
for pos, item in s.ls():
print(
f'Pos {pos}: {item.label} - ${item.val:.2f} (Count: {item.count})'
)
except SysErr as e:
print(f'Err: {str(e)}')
if __name__ == '__main__':
main()
================================================
FILE: data/raw_test_repo/inventory/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
"""Inventory management package for product stock tracking."""
================================================
FILE: data/raw_test_repo/inventory/inventory_manager.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import Dict, List, Optional
from ..models.product import Item
class Store:
def __init__(self, cap: int=20):
self.cap = cap
self._data: Dict[str, Item] = {}
self._map: Dict[int, str] = {}
def put(self, obj: Item, pos: Optional[int]=None) ->bool:
if obj.code in self._data:
curr = self._data[obj.code]
curr.count += obj.count
return True
if pos is not None:
if pos < 0 or pos >= self.cap:
return False
if pos in self._map:
return False
self._map[pos] = obj.code
else:
for i in range(self.cap):
if i not in self._map:
self._map[i] = obj.code
break
else:
return False
self._data[obj.code] = obj
return True
def rm(self, code: str) ->bool:
if code not in self._data:
return False
for k, v in list(self._map.items()):
if v == code:
del self._map[k]
del self._data[code]
return True
def get(self, code: str) ->Optional[Item]:
return self._data.get(code)
def get_at(self, pos: int) ->Optional[Item]:
if pos not in self._map:
return None
code = self._map[pos]
return self._data.get(code)
def ls(self) ->List[Item]:
return [obj for obj in self._data.values() if obj.check()]
def find(self, code: str) ->Optional[int]:
for k, v in self._map.items():
if v == code:
return k
return None
================================================
FILE: data/raw_test_repo/models/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
"""Models package for data structures used in the vending machine."""
================================================
FILE: data/raw_test_repo/models/product.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from dataclasses import dataclass
from typing import Optional
from datetime import datetime
@dataclass
class Item:
"""
Summary:
Represents an item with associated attributes for tracking and management in various contexts.
Description:
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.
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.
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.
Example:
```python
from datetime import datetime, timedelta
# Create an item with a specific expiration date
item = Item(code='A123', label='Sample Item', val=10.0, count=5, exp=datetime.now() + timedelta(days=1))
# Check if the item is valid
is_valid = item.check() # Returns True if count > 0 and not expired
# Modify the count of the item
item.mod(2) # Decreases count by 2, returns True
```
Parameters:
- code (str): A unique identifier for the item.
- label (str): A descriptive name for the item.
- val (float): The value associated with the item, representing its worth.
- count (int): The quantity of the item available. Must be a non-negative integer.
- exp (Optional[datetime]): An optional expiration date for the item. If set, the item will be considered invalid after this date.
- grp (str): A classification group for the item, defaulting to 'misc'.
Attributes:
- code (str): The unique identifier for the item.
- label (str): The name or description of the item.
- val (float): The monetary or functional value of the item.
- count (int): The current quantity of the item available, must be non-negative.
- exp (Optional[datetime]): The expiration date of the item, if applicable.
- grp (str): The group classification of the item, useful for categorization.
"""
code: str
label: str
val: float
count: int
exp: Optional[datetime] = None
grp: str = 'misc'
def check(self) -> bool:
"""
Validates the current object's state based on count and expiration.
Checks whether the object is still valid by verifying two key conditions:
1. The object's count is greater than zero
2. The object has not exceeded its expiration timestamp
This method is typically used to determine if an object is still usable
or has become stale/invalid. It provides a quick state validation check
that can be used in resource management, token validation, or lifecycle
tracking scenarios.
Returns:
bool: True if the object is valid (count > 0 and not expired),
False otherwise.
"""
if self.count <= 0:
return False
if self.exp and datetime.now() > self.exp:
return False
return True
def mod(self, n: int=1) -> bool:
"""
Summary:
Determines if the current count can be decremented by a specified value.
Description:
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.
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.
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.
Args:
n (int, optional): The value to decrement from `count`. Must be a positive integer that does not exceed the current `count`. Default is 1.
Returns:
bool: Returns `True` if the decrement was successful (i.e., `count` was greater than or equal to `n`), otherwise returns `False`.
Raises:
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.
Examples:
```python
obj = YourClass()
obj.count = 5
result = obj.mod(2) # result will be True, obj.count will be 3
result = obj.mod(4) # result will be False, obj.count remains 3
result = obj.mod(0) # result will be False, as n should be greater than 0
result = obj.mod(-1) # result will be False, as n should be a positive integer
```
"""
if self.count >= n:
self.count -= n
return True
return False
================================================
FILE: data/raw_test_repo/payment/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
"""Payment processing package for handling different payment methods."""
================================================
FILE: data/raw_test_repo/payment/payment_processor.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from abc import ABC, abstractmethod
from dataclasses import dataclass
from enum import Enum
from typing import Optional
from decimal import Decimal
class TxStatus(Enum):
WAIT = 'pending'
DONE = 'completed'
ERR = 'failed'
RET = 'refunded'
@dataclass
class Tx:
id: str
amt: Decimal
st: TxStatus
mth: str
msg: Optional[str] = None
class Handler(ABC):
@abstractmethod
def proc(self, amt: Decimal) ->Tx:
pass
@abstractmethod
def rev(self, tx: Tx) ->bool:
pass
class Cash(Handler):
def __init__(self):
self.bal: Decimal = Decimal('0.00')
def add(self, amt: Decimal) ->None:
self.bal += amt
def proc(self, amt: Decimal) ->Tx:
if self.bal >= amt:
self.bal -= amt
return Tx(id=f'C_{id(self)}', amt=amt, st=TxStatus.DONE, mth='cash'
)
return Tx(id=f'C_{id(self)}', amt=amt, st=TxStatus.ERR, mth='cash',
msg='insufficient')
def rev(self, tx: Tx) ->bool:
if tx.st == TxStatus.DONE:
self.bal += tx.amt
tx.st = TxStatus.RET
return True
return False
def ret(self) ->Decimal:
tmp = self.bal
self.bal = Decimal('0.00')
return tmp
================================================
FILE: data/raw_test_repo/vending_machine.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from decimal import Decimal
from typing import Optional, List, Tuple
from .models.product import Item
from .payment.payment_processor import Handler, Tx, TxStatus, Cash
from .inventory.inventory_manager import Store
class SysErr(Exception):
pass
class Sys:
def __init__(self, h: Optional[Handler]=None):
self.store = Store()
self.h = h or Cash()
self._tx: Optional[Tx] = None
def ls(self) ->List[Tuple[int, Item]]:
items = []
for item in self.store.ls():
pos = self.store.find(item.code)
if pos is not None:
items.append((pos, item))
return sorted(items, key=lambda x: x[0])
def pick(self, pos: int) ->Optional[Item]:
item = self.store.get_at(pos)
if not item:
raise SysErr('invalid pos')
if not item.check():
raise SysErr('unavailable')
return item
def add_money(self, amt: Decimal) ->None:
if not isinstance(self.h, Cash):
raise SysErr('cash not supported')
self.h.add(amt)
def buy(self, pos: int) ->Tuple[Item, Optional[Decimal]]:
item = self.pick(pos)
tx = self.h.proc(Decimal(str(item.val)))
self._tx = tx
if tx.st != TxStatus.DONE:
raise SysErr(tx.msg or 'tx failed')
if not item.mod():
self.h.rev(tx)
raise SysErr('dispense failed')
ret = None
if isinstance(self.h, Cash):
ret = self.h.ret()
return item, ret
def cancel(self) ->Optional[Decimal]:
if not self._tx:
raise SysErr('no tx')
ok = self.h.rev(self._tx)
if not ok:
raise SysErr('rev failed')
ret = None
if isinstance(self.h, Cash):
ret = self.h.ret()
self._tx = None
return ret
================================================
FILE: data/raw_test_repo_simple/helper.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
class HelperClass:
"""
Represents a utility for managing and processing data.
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.
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.
Example:
# Initialize the HelperClass
helper = HelperClass()
# Process data using the helper
helper.process_data()
# Retrieve the processed data result
result = helper.get_result()
print(result) # Output: '[1, 2, 3]'
Attributes:
data (list): Stores the processed data, initially an empty list.
"""
def __init__(self):
self.data = []
def process_data(self):
"""
Processes and updates the internal data.
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.
Returns:
None: This method updates the internal state and does not return a value.
"""
self.data = DataProcessor.process()
self._internal_process()
def _internal_process(self):
"""
No docstring provided.
"""
return self.data
def get_result(self):
"""
No docstring provided.
"""
return str(self.data)
class DataProcessor:
'''
"""Handles basic data processing tasks within a system.
This class is designed to perform simple data processing operations, providing
utility methods that can be used in various scenarios where basic data manipulation
is required. It is particularly useful in contexts where a straightforward list of
integers is needed for further processing or testing.
The `DataProcessor` class fits into the larger system architecture as a utility
component, offering static and internal methods to handle specific processing tasks.
It achieves its purpose by providing a static method for general use and an internal
method for class-specific operations.
Example:
# Initialize the DataProcessor class
processor = DataProcessor()
# Use the static method to process data
result = DataProcessor.process()
print(result) # Output: [1, 2, 3]
# Use the internal method for internal processing
internal_result = processor._internal_process()
print(internal_result) # Output: 'processed'
"""
'''
@staticmethod
def process():
'''
"""Processes data and returns a list of integers.
This static method is designed to perform a basic data processing task
and return a predefined list of integers. It can be used whenever a simple
list of integers is required for further operations or testing purposes.
Returns:
list of int: A list containing the integers [1, 2, 3].
"""
'''
return [1, 2, 3]
def _internal_process(self):
'''
"""Processes internal data and returns a status message.
This method is used internally within the `DataProcessor` class to perform
specific data processing tasks that are not exposed publicly. It is typically
called by other methods within the class to handle intermediate processing
steps.
Returns:
str: A string indicating the processing status, specifically 'processed'.
"""
'''
return 'processed'
================================================
FILE: data/raw_test_repo_simple/inner/inner_functions.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
def inner_function():
"""
Returns a greeting message from an inner function.
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.
Returns:
str: A greeting message stating 'Hello from inner function!'
Example:
>>> message = inner_function()
>>> print(message)
'Hello from inner function!'
"""
return 'Hello from inner function!'
def get_random_quote():
"""
Fetches a predefined inspirational quote.
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.
Returns:
str: A quote string stating 'The best way to predict the future is to create it.'
Example:
>>> quote = get_random_quote()
>>> print(quote)
'The best way to predict the future is to create it.'
"""
return 'The best way to predict the future is to create it.'
def generate_timestamp():
"""
Generates and returns a static timestamp.
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.
Returns:
str: A string representing the static timestamp '2023-05-15 14:30:22'.
"""
return '2023-05-15 14:30:22'
def get_system_status():
"""
Provides a static message indicating the operational status of systems.
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.
Returns:
str: A status message stating 'All systems operational.'
Example:
>>> status = get_system_status()
>>> print(status)
'All systems operational'
"""
return 'All systems operational'
def fetch_user_message():
'''
"""Fetches a predefined user message indicating notifications.
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.
Returns:
str: A message string stating 'Welcome back! You have 3 notifications.'
Example:
>>> message = fetch_user_message()
>>> print(message)
'Welcome back! You have 3 notifications.'
"""
'''
return 'Welcome back! You have 3 notifications.'
================================================
FILE: data/raw_test_repo_simple/main.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from helper import HelperClass
from inner.inner_functions import inner_function, get_random_quote, generate_timestamp, get_system_status, fetch_user_message
def main_function():
"""
Executes data processing and utility operations, returning the processed data as a string.
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.
Returns:
str: The processed data result as a string, derived from the `HelperClass` instance after executing the data processing and utility functions.
Example:
# Execute the main function to process data and retrieve the result
result = main_function()
print(result) # Output: '[1, 2, 3]'
"""
helper = HelperClass()
helper.process_data()
utility_function()
generate_timestamp()
return helper.get_result()
def utility_function():
"""
Returns a utility string.
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.
Returns:
str: The string 'utility', serving as a generic utility value.
"""
return 'utility'
================================================
FILE: data/raw_test_repo_simple/processor.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from helper import HelperClass
from processor import DataProcessor
from main import utility_function
class AdvancedProcessor:
"""
Facilitates advanced data processing by coordinating multiple processing components.
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.
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.
Example:
# Initialize the AdvancedProcessor
processor = AdvancedProcessor()
# Execute the processing workflow
result = processor.run()
print(result) # Output: 'utility'
Attributes:
helper (HelperClass): An instance of `HelperClass` used to manage data processing tasks.
data_processor (DataProcessor): An instance of `DataProcessor` used to perform specific data processing operations.
"""
def __init__(self):
self.helper = HelperClass()
self.data_processor = DataProcessor()
def run(self):
"""
Executes the complete data processing workflow and returns the result.
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.
Returns:
str: The result of the processing workflow, typically a utility string indicating successful completion.
Example:
# Create an instance of AdvancedProcessor
processor = AdvancedProcessor()
# Run the processing workflow
result = processor.run()
print(result) # Output: 'utility'
"""
self.helper.process_data()
self.data_processor._internal_process()
return self.process_result()
def process_result(self):
"""
Returns a utility string as the result of processing.
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.
Returns:
str: The string 'utility', serving as a generic utility value to indicate the completion of processing tasks.
"""
return utility_function()
================================================
FILE: data/raw_test_repo_simple/test_file.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
def test_function():
"""
Returns a boolean value indicating a successful test condition.
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.
Returns:
bool: The boolean value `True`, indicating a successful or positive condition.
"""
return True
================================================
FILE: eval_completeness.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
import ast
import os
from typing import Dict, Any, List, Union
from pathlib import Path
from evaluator.completeness import ClassCompletenessEvaluator, FunctionCompletenessEvaluator
from tabulate import tabulate
def run_docstring_tests(source_file: str) -> Dict[str, Any]:
"""
Run comprehensive docstring evaluation tests on a Python source file.
This function reads a Python file and evaluates docstrings for all classes,
functions, and methods found within. It provides detailed evaluation results
using different evaluators.
Args:
source_file: Path to the Python file to analyze
Returns:
Dictionary containing evaluation results for each found element
Example:
>>> results = run_docstring_tests('my_module.py')
>>> print(results['functions'][0])
1.0
"""
with open(source_file, 'r', encoding='utf-8') as f:
source = f.read()
try:
tree = ast.parse(source)
except SyntaxError as e:
return {
'status': 'error',
'message': f'Failed to parse {source_file}: {str(e)}'
}
results = {
'status': 'success',
'file': source_file,
'classes': [],
'functions': [],
'debug_info': {}
}
# Instantiate evaluators
class_evaluator = ClassCompletenessEvaluator()
func_evaluator = FunctionCompletenessEvaluator()
# Process all nodes in the AST
for node in ast.iter_child_nodes(tree):
if isinstance(node, ast.ClassDef):
class_result = {
'name': node.name,
'type': 'class',
'completeness_score': class_evaluator.evaluate(node),
'completeness_elements': class_evaluator.element_scores,
'element_required': class_evaluator.element_required
}
results['classes'].append(class_result)
# Evaluate methods within the class
for method in [n for n in ast.iter_child_nodes(node) if isinstance(n, ast.FunctionDef)]:
# Skip __init__ methods
if method.name == '__init__':
continue
method_result = {
'name': f"{node.name}.{method.name}",
'type': 'method',
'completeness_score': func_evaluator.evaluate(method),
'completeness_elements': func_evaluator.element_scores,
'element_required': func_evaluator.element_required
}
results['functions'].append(method_result)
elif isinstance(node, ast.FunctionDef):
# Only process top-level functions
func_result = {
'name': node.name,
'type': 'function',
'completeness_score': func_evaluator.evaluate(node),
'completeness_elements': func_evaluator.element_scores,
'element_required': func_evaluator.element_required
}
results['functions'].append(func_result)
# Add overall statistics
results['statistics'] = {
'total_classes': len(results['classes']),
'total_functions': len(results['functions']),
'average_class_score': sum(r['completeness_score'] for r in results['classes']) /
max(1, len(results['classes'])),
'average_function_score': sum(r['completeness_score'] for r in results['functions']) /
max(1, len(results['functions']))
}
return results
def process_directory(directory_path: str) -> Dict[str, Any]:
"""
Process all Python files in a directory and its subdirectories.
Args:
directory_path: Path to the directory to analyze
Returns:
Dictionary containing aggregated evaluation results for all files
"""
directory = Path(directory_path)
# Initialize aggregate results
aggregate_results = {
'status': 'success',
'directory': str(directory),
'files': [],
'file_results': [],
'classes': [],
'functions': [],
'statistics': {
'total_files': 0,
'successful_files': 0,
'failed_files': 0,
'total_classes': 0,
'total_functions': 0,
'average_class_score': 0.0,
'average_function_score': 0.0,
'overall_average_score': 0.0
}
}
# Find all Python files recursively
python_files = []
for root, _, files in os.walk(directory):
for file in files:
if file.endswith('.py'):
python_files.append(os.path.join(root, file))
if not python_files:
aggregate_results['status'] = 'error'
aggregate_results['message'] = f'No Python files found in {directory_path}'
return aggregate_results
aggregate_results['statistics']['total_files'] = len(python_files)
# Process each Python file
all_class_scores = []
all_function_scores = []
for py_file in python_files:
file_result = run_docstring_tests(py_file)
if file_result['status'] == 'success':
aggregate_results['successful_files'] = aggregate_results['statistics']['successful_files'] + 1
aggregate_results['file_results'].append(file_result)
aggregate_results['files'].append(py_file)
# Accumulate classes and functions with file path context
for class_result in file_result['classes']:
class_result['file'] = py_file
aggregate_results['classes'].append(class_result)
all_class_scores.append(class_result['completeness_score'])
for func_result in file_result['functions']:
func_result['file'] = py_file
aggregate_results['functions'].append(func_result)
all_function_scores.append(func_result['completeness_score'])
# Update statistics
aggregate_results['statistics']['total_classes'] += file_result['statistics']['total_classes']
aggregate_results['statistics']['total_functions'] += file_result['statistics']['total_functions']
else:
aggregate_results['statistics']['failed_files'] += 1
# Calculate average scores
if all_class_scores:
aggregate_results['statistics']['average_class_score'] = sum(all_class_scores) / len(all_class_scores)
if all_function_scores:
aggregate_results['statistics']['average_function_score'] = sum(all_function_scores) / len(all_function_scores)
# Calculate overall average score (classes and functions combined)
all_scores = all_class_scores + all_function_scores
if all_scores:
aggregate_results['statistics']['overall_average_score'] = sum(all_scores) / len(all_scores)
return aggregate_results
def print_evaluation_results(results: Dict[str, Any]) -> None:
"""
Pretty print the evaluation results in a readable format with colors.
Args:
results: Dictionary containing evaluation results from run_docstring_tests
"""
# ANSI color codes
GREEN = '\033[92m'
RED = '\033[91m'
BLUE = '\033[94m'
YELLOW = '\033[93m'
BOLD = '\033[1m'
ENDC = '\033[0m'
# Check if this is a directory result or a file result
is_directory = 'directory' in results
if is_directory:
# Print directory path
print(f"\n{BOLD}Evaluating Python files in directory: {results['directory']}{ENDC}")
print("=" * 80)
# Print file summary
print(f"\n{BLUE}{BOLD}FILE SUMMARY:{ENDC}")
stats_data = [
['Total Files', results['statistics']['total_files']],
['Successfully Processed Files', results['statistics']['successful_files']],
['Failed Files', results['statistics']['failed_files']]
]
print(tabulate(stats_data, tablefmt='simple'))
# Print overall statistics
print(f"\n{BLUE}{BOLD}OVERALL STATISTICS:{ENDC}")
# Add colored statistics
class_score = results['statistics']['average_class_score']
if class_score >= 0.8:
class_score_str = f"{GREEN}{class_score:.2f}{ENDC}"
elif class_score >= 0.5:
class_score_str = f"{YELLOW}{class_score:.2f}{ENDC}"
else:
class_score_str = f"{RED}{class_score:.2f}{ENDC}"
func_score = results['statistics']['average_function_score']
if func_score >= 0.8:
func_score_str = f"{GREEN}{func_score:.2f}{ENDC}"
elif func_score >= 0.5:
func_score_str = f"{YELLOW}{func_score:.2f}{ENDC}"
else:
func_score_str = f"{RED}{func_score:.2f}{ENDC}"
overall_score = results['statistics']['overall_average_score']
if overall_score >= 0.8:
overall_score_str = f"{GREEN}{overall_score:.2f}{ENDC}"
elif overall_score >= 0.5:
overall_score_str = f"{YELLOW}{overall_score:.2f}{ENDC}"
else:
overall_score_str = f"{RED}{overall_score:.2f}{ENDC}"
stats_data = [
['Total Classes', results['statistics']['total_classes']],
['Total Functions/Methods', results['statistics']['total_functions']],
['Average Class Score', class_score_str],
['Average Function Score', func_score_str],
['Overall Average Score', overall_score_str]
]
print(tabulate(stats_data, tablefmt='simple'))
# Ask if the user wants to see details for individual files
print(f"\nUse python {os.path.basename(__file__)} <specific_file_path> to see detailed results for a specific file.")
else:
# Original single file display logic
# Print file path
print(f"\n{BOLD}Evaluating Python file: {results['file']}{ENDC}")
print("=" * 80)
# Print class results table
if results['classes']:
print(f"\n{BLUE}{BOLD}CLASSES:{ENDC}")
headers = ['Class Name', 'Score']
elements = list(results['classes'][0]['completeness_elements'].keys())
headers.extend(elements)
table_data = []
for class_result in results['classes']:
row = [class_result['name']]
score = class_result['completeness_score']
# Color the score based on value
if score >= 0.8:
score_str = f"{GREEN}{score:.2f}{ENDC}"
elif score >= 0.5:
score_str = f"{YELLOW}{score:.2f}{ENDC}"
else:
score_str = f"{RED}{score:.2f}{ENDC}"
row.append(score_str)
for element in elements:
required = class_result['element_required'][element]
has_element = class_result['completeness_elements'][element]
if has_element:
check = f"{GREEN}✓{ENDC}"
else:
check = f"{RED}✗{ENDC}"
cell = f"{YELLOW if required else '-'}{'R' if required else ''}{ENDC if required else ''} | {check}"
row.append(cell)
table_data.append(row)
print(tabulate(table_data, headers=headers, tablefmt='grid'))
# Print function/method results table
if results['functions']:
print(f"\n{BLUE}{BOLD}FUNCTIONS/METHODS:{ENDC}")
headers = ['Function Name', 'Type', 'Score']
elements = list(results['functions'][0]['completeness_elements'].keys())
headers.extend(elements)
table_data = []
for func_result in results['functions']:
row = [func_result['name'], func_result['type']]
score = func_result['completeness_score']
# Color the score based on value
if score >= 0.8:
score_str = f"{GREEN}{score:.2f}{ENDC}"
elif score >= 0.5:
score_str = f"{YELLOW}{score:.2f}{ENDC}"
else:
score_str = f"{RED}{score:.2f}{ENDC}"
row.append(score_str)
for element in elements:
required = func_result['element_required'][element]
has_element = func_result['completeness_elements'][element]
if has_element:
check = f"{GREEN}✓{ENDC}"
else:
check = f"{RED}✗{ENDC}"
cell = f"{YELLOW if required else '-'}{'R' if required else ''}{ENDC if required else ''} | {check}"
row.append(cell)
table_data.append(row)
print(tabulate(table_data, headers=headers, tablefmt='grid'))
# Print overall statistics
print(f"\n{BLUE}{BOLD}OVERALL STATISTICS:{ENDC}")
stats_data = []
# Add colored statistics
class_score = results['statistics']['average_class_score']
if class_score >= 0.8:
class_score_str = f"{GREEN}{class_score:.2f}{ENDC}"
elif class_score >= 0.5:
class_score_str = f"{YELLOW}{class_score:.2f}{ENDC}"
else:
class_score_str = f"{RED}{class_score:.2f}{ENDC}"
func_score = results['statistics']['average_function_score']
if func_score >= 0.8:
func_score_str = f"{GREEN}{func_score:.2f}{ENDC}"
elif func_score >= 0.5:
func_score_str = f"{YELLOW}{func_score:.2f}{ENDC}"
else:
func_score_str = f"{RED}{func_score:.2f}{ENDC}"
stats_data = [
['Total Classes', results['statistics']['total_classes']],
['Total Functions/Methods', results['statistics']['total_functions']],
['Average Class Score', class_score_str],
['Average Function Score', func_score_str]
]
print(tabulate(stats_data, tablefmt='simple'))
if __name__ == "__main__":
# Example usage
import sys
if len(sys.argv) < 2:
print("Usage: python eval_completeness.py <path_to_python_file_or_directory>")
sys.exit(1)
path = sys.argv[1]
if not Path(path).exists():
print(f"Error: Path not found: {path}")
sys.exit(1)
if Path(path).is_dir():
# Process directory
results = process_directory(path)
if results['status'] == 'success':
print_evaluation_results(results)
else:
print(f"Error: {results['message']}")
else:
# Process single file
results = run_docstring_tests(path)
if results['status'] == 'success':
print_evaluation_results(results)
else:
print(f"Error: {results['message']}")
================================================
FILE: generate_docstrings.py
================================================
#!/usr/bin/env python3
# Copyright (c) Meta Platforms, Inc. and affiliates
"""
Docstring Generator with Dependency-Based Ordering
This script generates docstrings for Python code components (functions, classes, methods)
using a DFS-based approach that starts from components with no dependencies.
Key features:
1. Parses Python code to identify components and their dependencies
2. Builds a dependency graph where A→B means "A depends on B"
3. Performs DFS traversal starting from components with no dependencies
4. Processes dependencies before the components that depend on them
5. Ensures classes depend on their methods, not vice versa
6. Skips __init__ methods as they typically don't need separate docstrings
7. Provides visual representation of progress in the terminal
Usage:
python generate_docstrings.py --repo-path PATH --config-path PATH [--test-mode]
"""
import os
import sys
import time
import ast
import json
import argparse
import logging
import random
from pathlib import Path
from typing import Dict, List, Set, Optional, Any
from collections import defaultdict
import tiktoken # Add this import for token counting
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.StreamHandler(sys.stdout)
]
)
logger = logging.getLogger("docstring_generator")
# Import dependency analyzer modules
from src.dependency_analyzer import (
CodeComponent,
DependencyParser,
dependency_first_dfs,
build_graph_from_components
)
from src.visualizer import ProgressVisualizer
from src.agent.orchestrator import Orchestrator
def generate_test_docstring(component: CodeComponent) -> str:
"""
Generate a placeholder docstring for test mode.
Args:
component: The code component to generate a placeholder docstring for.
Returns:
A placeholder docstring based on the component type.
"""
comp_type = component.component_type
name = component.id.split(".")[-1]
if comp_type == "function":
return f"""
Test docstring for function '{name}'.
This is a placeholder docstring generated in test mode.
In normal mode, this would be replaced with an AI-generated docstring.
Args:
arg1: Description of first argument
arg2: Description of second argument
Returns:
Description of return value
"""
elif comp_type == "class":
return f"""
Test docstring for class '{name}'.
This is a placeholder docstring generated in test mode.
In normal mode, this would be replaced with an AI-generated docstring.
Attributes:
attr1: Description of first attribute
attr2: Description of second attribute
"""
elif comp_type == "method":
class_name = component.id.split(".")[-2]
return f"""
Test docstring for method '{name}' in class '{class_name}'.
This is a placeholder docstring generated in test mode.
In normal mode, this would be replaced with an AI-generated docstring.
Args:
arg1: Description of first argument
arg2: Description of second argument
Returns:
Description of return value
"""
else:
return f"""
Test docstring for {comp_type} '{name}'.
This is a placeholder docstring generated in test mode.
"""
def generate_docstring_for_component(component: CodeComponent, orchestrator: Optional[Orchestrator], test_mode: str = 'none',
dependency_graph: Optional[Dict[str, List[str]]] = None) -> str:
"""
Generate a docstring for a single component.
Args:
component: The component to generate a docstring for.
orchestrator: The orchestrator instance.
test_mode: The test mode to use.
dependency_graph: Optional dependency graph.
Returns:
The generated docstring.
"""
# do not use try/except here, we want to fail if there is an error
if not orchestrator:
return ""
file_path = component.file_path
# Get the component code
component_code = component.source_code
# Estimate token count of the focal component
encoding = tiktoken.get_encoding("cl100k_base") # Default OpenAI encoding
token_consume_focal = len(encoding.encode(component_code))
# Skip if the component is too large (> 10000 tokens)
if token_consume_focal > 10000:
# truncate the component code to 10000 tokens
component_code = encoding.decode(encoding.encode(component_code)[:10000])
# Parse the file
with open(file_path, "r", encoding="utf-8") as f:
file_content = f.read()
ast_tree = ast.parse(file_content)
ast_node = None
# Locate the AST node for the component
component_parts = component.id.split(".")
component_name = component_parts[-1]
if component.component_type == "function":
# Find top-level function
for node in ast.iter_child_nodes(ast_tree):
if (isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
and node.name == component_name):
ast_node = node
break
elif component.component_type == "class":
# Find class
for node in ast.iter_child_nodes(ast_tree):
if isinstance(node, ast.ClassDef) and node.name == component_name:
ast_node = node
break
elif component.component_type == "method":
# Find method inside class
class_name, method_name = component_parts[-2:]
for node in ast.iter_child_nodes(ast_tree):
if isinstance(node, ast.ClassDef) and node.name == class_name:
for item in node.body:
if (isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
and item.name == method_name):
ast_node = item
break
break
try:
# Pass component.id as the focal_node_dependency_path
docstring = orchestrator.process(
focal_component=component_code,
file_path=file_path,
ast_node=ast_node,
ast_tree=ast_tree,
dependency_graph=dependency_graph,
focal_node_dependency_path=component.id,
token_consume_focal=token_consume_focal # Pass token count to orchestrator
)
return docstring
except Exception as e:
print(f"Error generating docstring for {component.id}: {str(e)}")
return ""
def set_docstring_in_file(file_path: str, component: CodeComponent, docstring: str) -> bool:
"""
Update a Python file with a newly generated docstring for a component.
Args:
file_path: Path to the file to update.
component: The component to update with a docstring.
docstring: The docstring to insert.
Returns:
True if successful, False otherwise.
"""
# Do not use Try/Except here, we want to fail if there is an error
# Read the file
with open(file_path, "r", encoding="utf-8") as f:
source = f.read()
# Parse the file
tree = ast.parse(source)
# Find the component in the parsed AST
component_node = None
component_parts = component.id.split(".")
component_name = component_parts[-1]
if component.component_type == "function":
# Find top-level function
for node in ast.iter_child_nodes(tree):
if (isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
and node.name == component_name):
component_node = node
break
elif component.component_type == "class":
# Find class
for node in ast.iter_child_nodes(tree):
if isinstance(node, ast.ClassDef) and node.name == component_name:
component_node = node
break
elif component.component_type == "method":
# Find method inside class
class_name, method_name = component_parts[-2:]
for node in ast.iter_child_nodes(tree):
if isinstance(node, ast.ClassDef) and node.name == class_name:
for item in node.body:
if (isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef))
and item.name == method_name):
component_node = item
break
break
if not component_node:
logger.error(f"Could not find component {component.id} in {file_path}")
return False
# Set the docstring
set_node_docstring(component_node, docstring)
# Unparse the AST back to source code
if hasattr(ast, "unparse"):
new_source = ast.unparse(tree)
else:
try:
import astor
new_source = astor.to_source(tree)
except ImportError:
logger.error(
"Error: You need to install 'astor' or use Python 3.9+ to unparse the AST. "
f"Skipping file: {file_path}"
)
return False
# Write back to the file
with open(file_path, "w", encoding="utf-8") as f:
f.write(new_source)
return True
def set_node_docstring(node: ast.AST, docstring: str):
"""
Safely set or update the docstring on an AST node (ClassDef, FunctionDef, etc.).
Also adjusts indentation relative to the node's existing indentation level,
ensuring both the opening and closing triple quotes are properly aligned.
Args:
node: The AST node to modify (ClassDef, FunctionDef, etc.).
docstring: The new docstring (as a plain string) to insert.
"""
import textwrap
# 1. Strip leading/trailing empty lines in the provided docstring
# to avoid spurious blank lines.
stripped_docstring = docstring.strip('\n')
if not stripped_docstring:
# If empty or all whitespace, provide a placeholder
stripped_docstring = "No docstring provided."
# 2. Dedent possible indentation in docstring (so a multiline docstring
# doesn't carry undesired left margins).
dedented = textwrap.dedent(stripped_docstring)
# 3. Determine how many spaces to indent for doc lines plus triple quotes.
existing_indent = getattr(node, 'col_offset', 0)
doc_indent_str = ' ' * (existing_indent + 4)
# 4. Build the final string:
# - Start with a newline (so triple quotes appear on a new line).
# - Indent all docstring lines.
# - End with a newline+same indentation (so the closing triple quotes
# line also has the doc_indent_str).
prepared_docstring = (
"\n"
+ textwrap.indent(dedented, doc_indent_str)
+ "\n"
+ doc_indent_str
)
# 5. Create an AST Expr node to store this docstring as a constant.
docstring_node = ast.Expr(value=ast.Constant(value=prepared_docstring, kind=None))
# If there's no body, just make one containing our new docstring.
if not hasattr(node, "body") or not isinstance(node.body, list) or len(node.body) == 0:
node.body = [docstring_node]
else:
# If the first statement is an existing docstring, replace it;
# otherwise, insert the new docstring as the first statement.
first_stmt = node.body[0]
if (
isinstance(first_stmt, ast.Expr)
and isinstance(first_stmt.value, ast.Constant)
and isinstance(first_stmt.value.value, str)
):
node.body[0] = docstring_node
else:
node.body.insert(0, docstring_node)
def main():
"""
Main entry point for the docstring generation script with flexible component ordering.
The script supports different ordering modes through the --order-mode flag:
- 'topo' (default): Dependency-based ordering using a DFS-based approach:
1. If A depends on B, the graph has an edge A→B (meaning "A depends on B")
2. Root nodes (nodes with no dependencies) are processed first
3. Dependencies are always processed before the components that depend on them
4. This ensures proper docstring generation order
- 'random_node': Randomly shuffles all Python components, ignoring dependencies
- 'random_file': Processes files in random order, but preserves component order within files
Class methods are processed before the classes that depend on them (not vice versa) in 'topo' mode,
ensuring proper docstring generation order. Special __init__ methods are skipped as
they typically don't need separate docstrings.
The script provides options to skip or overwrite existing docstrings:
- By default, components with existing docstrings are skipped
- With --overwrite-docstrings flag, existing docstrings will be overwritten
- This behavior can also be configured in the config.yaml file under docstring_options.overwrite_docstrings
Web interface integration:
- With --enable-web flag, the script enables integration with the web UI
- This allows visualization of the docstring generation process in a web browser
- Run the web UI separately using the run_web_ui.py script
"""
# Parse command line arguments
parser = argparse.ArgumentParser(
description='Generate docstrings for Python components in dependency order.'
)
parser.add_argument(
'--repo-path',
type=str,
default='data/raw_test_repo',
help='Path to the repository (default: data/raw_test_repo)'
)
parser.add_argument(
'--config-path',
type=str,
default='config/agent_config.yaml',
help='Path to the configuration file (default: config/agent_config.yaml)'
)
parser.add_argument(
'--test-mode',
type=str,
choices=['placeholder', 'context_print', 'none'],
default='none',
help='Test mode to run: "placeholder" for placeholder docstrings (no LLM calls), "context_print" to print context before writer calls, "none" for normal operation'
)
parser.add_argument(
'--order-mode',
type=str,
choices=['topo', 'random_node', 'random_file'],
default='topo',
help='Order mode for docstring generation: "topo" follows dependency order (default), "random_node" selects random Python nodes, "random_file" processes files in random order'
)
parser.add_argument(
'--enable-web',
action='store_true',
help='Enable integration with the web interface'
)
parser.add_argument(
'--overwrite-docstrings',
action='store_true',
help='Overwrite existing docstrings instead of skipping them (default: False)'
)
args = parser.parse_args()
repo_path = args.repo_path
config_path = args.config_path
test_mode = args.test_mode
order_mode = args.order_mode
overwrite_docstrings = args.overwrite_docstrings
# Create output directory for dependency graph
output_dir = os.path.join("output", "dependency_graphs")
os.makedirs(output_dir, exist_ok=True)
# Extract repository name from path for creating a unique filename
repo_name = os.path.basename(os.path.normpath(repo_path))
# Create a sanitized version of the repo name (remove special characters)
sanitized_repo_name = ''.join(c if c.isalnum() else '_' for c in repo_name)
dependency_graph_path = os.path.join(output_dir, f"{sanitized_repo_name}_dependency_graph.json")
# Initialize the orchestrator for docstring generation
orchestrator = None
# Initialize orchestrator unless we're in placeholder test mode
if test_mode != 'placeholder':
logger.info(f"Initializing orchestrator with config: {config_path}")
# Pass the test_mode to the orchestrator if it's "context_print"
orchestrator_test_mode = test_mode if test_mode != 'none' else None
orchestrator = Orchestrator(repo_path=repo_path, config_path=config_path, test_mode=orchestrator_test_mode)
# Check if the overwrite_docstrings option is in the config file
# If it's there, it overrides the command-line argument
if hasattr(orchestrator, 'config'):
docstring_options = orchestrator.config.get('docstring_options', {})
config_overwrite = docstring_options.get('overwrite_docstrings')
if config_overwrite is not None:
overwrite_docstrings = config_overwrite
logger.info(f"Using config file setting for overwrite_docstrings: {overwrite_docstrings}")
else:
logger.info("Running in PLACEHOLDER TEST MODE with placeholder docstrings (no LLM calls)")
# Parse the repository to build the dependency graph
logger.info(f"Parsing repository: {repo_path}")
parser = DependencyParser(repo_path)
components = parser.parse_repository()
# Save the dependency graph for future reference
parser.save_dependency_graph(dependency_graph_path)
logger.info(f"Dependency graph saved to: {dependency_graph_path}")
# Build the graph for traversal
graph = build_graph_from_components(components)
# Create a dependency graph in the format expected by the orchestrator
# Dictionary mapping component paths to their dependencies
dependency_graph = {}
for component_id, deps in graph.items():
dependency_graph[component_id] = list(deps)
# Perform DFS-based traversal
logger.info("Performing DFS traversal on the dependency graph (starting from nodes with no dependencies)")
sorted_components = dependency_first_dfs(graph)
logger.info(f"Sorted {len(sorted_components)} components for processing")
# Apply the selected ordering mode
if order_mode == 'random_node':
# Randomly shuffle all components
logger.info("Using random node ordering mode - shuffling all components")
random.shuffle(sorted_components)
elif order_mode == 'random_file':
# Group components by file path
logger.info("Using random file ordering mode - processing files in random order")
# Group components by file
file_to_components = defaultdict(list)
for component_id in sorted_components:
component = components.get(component_id)
if component:
file_to_components[component.file_path].append(component_id)
# Randomly shuffle the file order but maintain the order of components within each file
file_paths = list(file_to_components.keys())
random.shuffle(file_paths)
# Create a new ordering based on randomly shuffled files
sorted_components = []
for file_path in file_paths:
sorted_components.extend(file_to_components[file_path])
else:
# Default to topological order (already set in sorted_components)
logger.info("Using topological ordering mode - processing components based on dependencies")
# Check if web interface is enabled
if args.enable_web:
try:
from src.visualizer.web_bridge import patch_visualizers
logger.info("Web interface integration enabled")
patch_visualizers()
except ImportError as e:
logger.warning(f"Failed to enable web interface integration: {e}")
logger.warning("Make sure you have installed the required web dependencies")
# Initialize the progress visualizer
visualizer = ProgressVisualizer(components, sorted_components)
visualizer.initialize()
# Show dependency statistics
visualizer.show_dependency_stats()
# Process components in order determined by DFS traversal
for component_id in sorted_components:
component = components.get(component_id)
if not component:
logger.warning(f"Component {component_id} not found in parsed components")
continue
# Skip __init__ methods as they don't need docstrings
if component.component_type == "method" and component_id.endswith(".__init__"):
logger.info(f"Skipping {component_id} - __init__ methods don't need docstrings")
visualizer.update(component_id, "completed")
continue
# compute the length of docstring if exists (using white space as delimiter)
docstring_length = len(component.docstring.split()) if component.has_docstring else 0
# Skip components that already have docstrings (unless overwrite_docstrings is True)
if component.has_docstring and not overwrite_docstrings and docstring_length > 10:
logger.info(f"Skipping {component_id} - already has docstring")
visualizer.update(component_id, "completed")
continue
elif component.has_docstring and overwrite_docstrings:
logger.info(f"Overwriting existing docstring for {component_id}")
# Update the visualizer
visualizer.update(component_id, "processing")
# Log the component type
comp_type = component.component_type
logger.info(f"Processing {comp_type}: {component_id}")
# Generate the docstring
logger.info(f"Generating docstring for {component_id}")
docstring = generate_docstring_for_component(component, orchestrator, test_mode, dependency_graph)
# Update the file with the new docstring
file_path = component.file_path
success = set_docstring_in_file(file_path, component, docstring)
if success:
logger.info(f"Successfully updated docstring for {component_id}")
visualizer.update(component_id, "completed")
else:
logger.error(f"Failed to update docstring for {component_id}")
visualizer.update(component_id, "error")
# Re-parse the file in case the line numbers changed due to docstring insertion
# This is only necessary if there are more components from the same file
same_file_components = [
comp_id for comp_id in sorted_components
if comp_id != component_id and components[comp_id].file_path == file_path
]
if same_file_components:
logger.info(f"Re-parsing file {file_path} for updated line numbers")
parser = DependencyParser(repo_path)
updated_components = parser.parse_repository()
# Update the components dictionary with new line numbers
for comp_id, comp in updated_components.items():
if comp_id in components:
components[comp_id] = comp
# Finalize the visualization
visualizer.finalize()
# Create a more descriptive mode message based on the test mode
if test_mode == 'placeholder':
mode_str = "PLACEHOLDER TEST MODE (no LLM calls)"
elif test_mode == 'context_print':
mode_str = "CONTEXT PRINT TEST MODE (with context debugging)"
else:
mode_str = "normal mode"
# Add ordering mode to the completion message
order_mode_str = {
'topo': 'topological ordering',
'random_node': 'random node ordering',
'random_file': 'random file ordering'
}.get(order_mode, 'unknown ordering')
logger.info(f"Docstring generation complete ({mode_str}, {order_mode_str})")
# Print usage statistics for LLM providers if available
if orchestrator:
try:
# Access the rate limiters from agents
rate_limiters = []
for agent_name in ['reader', 'writer', 'verifier']:
agent = getattr(orchestrator, agent_name, None)
if agent and hasattr(agent, 'llm') and hasattr(agent.llm, 'rate_limiter'):
rate_limiters.append(agent.llm.rate_limiter)
# Print statistics for each rate limiter
if rate_limiters:
logger.info("=" * 50)
logger.info("TOKEN USAGE AND COST STATISTICS")
logger.info("=" * 50)
for limiter in rate_limiters:
limiter.print_usage_stats()
# Calculate total cost across all limiters
total_cost = sum(limiter.total_cost for limiter in rate_limiters)
logger.info("=" * 50)
logger.info(f"TOTAL COST: ${total_cost:.6f}")
logger.info("=" * 50)
except Exception as e:
logger.warning(f"Could not print token usage statistics: {e}")
if __name__ == "__main__":
main()
================================================
FILE: output/dependency_graphs/raw_test_repo_dependency_graph.json
================================================
{
"helper.HelperClass": {
"id": "helper.HelperClass",
"component_type": "class",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [
"helper.HelperClass.get_result",
"helper.DataProcessor",
"helper.HelperClass._internal_process",
"helper.HelperClass.process_data"
],
"start_line": 1,
"end_line": 14,
"has_docstring": false,
"docstring": ""
},
"helper.HelperClass.__init__": {
"id": "helper.HelperClass.__init__",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [],
"start_line": 3,
"end_line": 4,
"has_docstring": false,
"docstring": ""
},
"helper.HelperClass.process_data": {
"id": "helper.HelperClass.process_data",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [
"helper.DataProcessor"
],
"start_line": 6,
"end_line": 8,
"has_docstring": false,
"docstring": ""
},
"helper.HelperClass._internal_process": {
"id": "helper.HelperClass._internal_process",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [],
"start_line": 10,
"end_line": 11,
"has_docstring": false,
"docstring": ""
},
"helper.HelperClass.get_result": {
"id": "helper.HelperClass.get_result",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [],
"start_line": 13,
"end_line": 14,
"has_docstring": false,
"docstring": ""
},
"helper.DataProcessor": {
"id": "helper.DataProcessor",
"component_type": "class",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [
"helper.DataProcessor.process",
"helper.DataProcessor._internal_process"
],
"start_line": 16,
"end_line": 72,
"has_docstring": true,
"docstring": "\n \"\"\"Handles basic data processing tasks within a system.\n\n This class is designed to perform simple data processing operations, providing\n utility methods that can be used in various scenarios where basic data manipulation\n is required. It is particularly useful in contexts where a straightforward list of\n integers is needed for further processing or testing.\n\n The `DataProcessor` class fits into the larger system architecture as a utility\n component, offering static and internal methods to handle specific processing tasks.\n It achieves its purpose by providing a static method for general use and an internal\n method for class-specific operations.\n\n Example:\n # Initialize the DataProcessor class\n processor = DataProcessor()\n\n # Use the static method to process data\n result = DataProcessor.process()\n print(result) # Output: [1, 2, 3]\n\n # Use the internal method for internal processing\n internal_result = processor._internal_process()\n print(internal_result) # Output: 'processed'\n \"\"\"\n "
},
"helper.DataProcessor.process": {
"id": "helper.DataProcessor.process",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [],
"start_line": 45,
"end_line": 57,
"has_docstring": true,
"docstring": "\n \"\"\"Processes data and returns a list of integers.\n\n This static method is designed to perform a basic data processing task\n and return a predefined list of integers. It can be used whenever a simple\n list of integers is required for further operations or testing purposes.\n\n Returns:\n list of int: A list containing the integers [1, 2, 3].\n \"\"\"\n "
},
"helper.DataProcessor._internal_process": {
"id": "helper.DataProcessor._internal_process",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/helper.py",
"relative_path": "helper.py",
"depends_on": [],
"start_line": 59,
"end_line": 72,
"has_docstring": true,
"docstring": "\n \"\"\"Processes internal data and returns a status message.\n\n This method is used internally within the `DataProcessor` class to perform\n specific data processing tasks that are not exposed publicly. It is typically\n called by other methods within the class to handle intermediate processing\n steps.\n\n Returns:\n str: A string indicating the processing status, specifically 'processed'.\n \"\"\"\n "
},
"main.main_function": {
"id": "main.main_function",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/main.py",
"relative_path": "main.py",
"depends_on": [
"helper.HelperClass",
"inner.inner_functions.generate_timestamp",
"main.utility_function"
],
"start_line": 5,
"end_line": 10,
"has_docstring": false,
"docstring": ""
},
"main.utility_function": {
"id": "main.utility_function",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/main.py",
"relative_path": "main.py",
"depends_on": [],
"start_line": 13,
"end_line": 14,
"has_docstring": false,
"docstring": ""
},
"processor.AdvancedProcessor": {
"id": "processor.AdvancedProcessor",
"component_type": "class",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/processor.py",
"relative_path": "processor.py",
"depends_on": [
"processor.AdvancedProcessor.run",
"helper.HelperClass",
"processor.AdvancedProcessor.process_result",
"main.utility_function",
"processor.DataProcessor"
],
"start_line": 6,
"end_line": 18,
"has_docstring": false,
"docstring": ""
},
"processor.AdvancedProcessor.__init__": {
"id": "processor.AdvancedProcessor.__init__",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/processor.py",
"relative_path": "processor.py",
"depends_on": [
"helper.HelperClass",
"processor.DataProcessor"
],
"start_line": 8,
"end_line": 10,
"has_docstring": false,
"docstring": ""
},
"processor.AdvancedProcessor.run": {
"id": "processor.AdvancedProcessor.run",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/processor.py",
"relative_path": "processor.py",
"depends_on": [],
"start_line": 12,
"end_line": 15,
"has_docstring": false,
"docstring": ""
},
"processor.AdvancedProcessor.process_result": {
"id": "processor.AdvancedProcessor.process_result",
"component_type": "method",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/processor.py",
"relative_path": "processor.py",
"depends_on": [
"main.utility_function"
],
"start_line": 17,
"end_line": 18,
"has_docstring": false,
"docstring": ""
},
"test_file.test_function": {
"id": "test_file.test_function",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/test_file.py",
"relative_path": "test_file.py",
"depends_on": [],
"start_line": 1,
"end_line": 2,
"has_docstring": false,
"docstring": ""
},
"inner.inner_functions.inner_function": {
"id": "inner.inner_functions.inner_function",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/inner/inner_functions.py",
"relative_path": "inner/inner_functions.py",
"depends_on": [],
"start_line": 1,
"end_line": 15,
"has_docstring": true,
"docstring": "\n Returns a greeting message from an inner function.\n\n 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.\n\n Returns:\n str: A greeting message stating 'Hello from inner function!'\n\n Example:\n >>> message = inner_function()\n >>> print(message)\n 'Hello from inner function!'\n "
},
"inner.inner_functions.get_random_quote": {
"id": "inner.inner_functions.get_random_quote",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/inner/inner_functions.py",
"relative_path": "inner/inner_functions.py",
"depends_on": [],
"start_line": 17,
"end_line": 31,
"has_docstring": true,
"docstring": "\n Fetches a predefined inspirational quote.\n\n 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.\n\n Returns:\n str: A quote string stating 'The best way to predict the future is to create it.'\n\n Example:\n >>> quote = get_random_quote()\n >>> print(quote)\n 'The best way to predict the future is to create it.'\n "
},
"inner.inner_functions.generate_timestamp": {
"id": "inner.inner_functions.generate_timestamp",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/inner/inner_functions.py",
"relative_path": "inner/inner_functions.py",
"depends_on": [],
"start_line": 33,
"end_line": 34,
"has_docstring": false,
"docstring": ""
},
"inner.inner_functions.get_system_status": {
"id": "inner.inner_functions.get_system_status",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/inner/inner_functions.py",
"relative_path": "inner/inner_functions.py",
"depends_on": [],
"start_line": 36,
"end_line": 50,
"has_docstring": true,
"docstring": "\n Provides a static message indicating the operational status of systems.\n\n 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.\n\n Returns:\n str: A status message stating 'All systems operational.'\n\n Example:\n >>> status = get_system_status()\n >>> print(status)\n 'All systems operational'\n "
},
"inner.inner_functions.fetch_user_message": {
"id": "inner.inner_functions.fetch_user_message",
"component_type": "function",
"file_path": "/home/dayuyang/DocAgent/data/raw_test_repo/inner/inner_functions.py",
"relative_path": "inner/inner_functions.py",
"depends_on": [],
"start_line": 52,
"end_line": 67,
"has_docstring": true,
"docstring": "\n \"\"\"Fetches a predefined user message indicating notifications.\n\n 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.\n\n Returns:\n str: A message string stating 'Welcome back! You have 3 notifications.'\n\n Example:\n >>> message = fetch_user_message()\n >>> print(message)\n 'Welcome back! You have 3 notifications.'\n \"\"\"\n "
}
}
================================================
FILE: run_web_ui.py
================================================
#!/usr/bin/env python3
import eventlet
eventlet.monkey_patch()
# Copyright (c) Meta Platforms, Inc. and affiliates
"""
Web UI Launcher for DocAgent Docstring Generator
This script launches the web-based user interface for the docstring generation tool.
The UI provides a more interactive and visual way to use the docstring generator,
with real-time feedback and progress tracking.
Usage:
python run_web_ui.py [--host HOST] [--port PORT] [--debug]
"""
import argparse
import os
import sys
import logging
from pathlib import Path
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.StreamHandler(sys.stdout)
]
)
logger = logging.getLogger("docstring_web")
# Add the current directory to the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
def check_dependencies():
"""Check if all required dependencies are installed."""
try:
import flask
import flask_socketio
import eventlet
import yaml
import tabulate
import colorama
return True
except ImportError as e:
missing_module = str(e).split("'")[1]
logger.error(f"Missing dependency: {missing_module}")
logger.error("Please install all required dependencies with:")
logger.error("pip install -r requirements-web.txt")
return False
def main():
"""Parse command line arguments and start the web UI."""
parser = argparse.ArgumentParser(description='Launch the DocAgent Web UI')
parser.add_argument('--host', default='127.0.0.1', help='Host to bind the server to')
parser.add_argument('--port', type=int, default=5000, help='Port to bind the server to')
parser.add_argument('--debug', action='store_true', help='Run in debug mode')
args = parser.parse_args()
# Check dependencies
if not check_dependencies():
return 1
# Print banner
print("\n" + "=" * 80)
print("DocAgent Web Interface".center(80))
print("=" * 80)
# Import and run the web app
try:
# First try to import eventlet to ensure it's properly initialized
import eventlet
eventlet.monkey_patch()
from src.web.app import create_app
app, socketio = create_app(debug=args.debug)
logger.info(f"Starting DocAgent Web UI at: http://{args.host}:{args.port}")
logger.info("Press Ctrl+C to stop the server")
# Start the server
socketio.run(app, host=args.host, port=args.port, debug=args.debug, allow_unsafe_werkzeug=True)
return 0
except ImportError as e:
logger.error(f"Error importing web application: {e}")
logger.error("Make sure the src/web directory exists and contains the necessary files.")
return 1
except Exception as e:
logger.error(f"Error running web application: {e}")
return 1
if __name__ == '__main__':
try:
sys.exit(main())
except KeyboardInterrupt:
print("\nServer stopped.")
sys.exit(0)
================================================
FILE: setup.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from setuptools import setup, find_packages
# Read the contents of README file
from pathlib import Path
this_directory = Path(__file__).parent
long_description = (this_directory / "README.md").read_text()
# Prepare all extras
dev_requires = [
"pytest>=8.3.4",
"pytest-cov>=2.0",
"black>=22.0",
"flake8>=3.9",
]
web_requires = [
"flask>=3.1.0",
"flask-socketio>=5.5.1",
"eventlet>=0.39.0",
"python-socketio>=5.12.1",
"python-engineio>=4.11.2",
"bidict>=0.23.0",
"dnspython>=2.7.0",
"six>=1.16.0",
]
visualization_requires = [
"matplotlib>=3.10.0",
"pygraphviz>=1.14",
"networkx>=3.4.2",
]
cuda_requires = [
"torch>=2.0.0",
"accelerate>=1.4.0",
]
# Combine all extras for the 'all' option
all_requires = dev_requires + web_requires + visualization_requires + cuda_requires
setup(
name="DocstringGenerator",
version="0.1.0",
author="Dayu Yang",
author_email="dayuyang@meta.com",
description="DocAgent for High-quality docstring generation in Large-scale Python projects",
long_description=long_description,
long_description_content_type="text/markdown",
packages=find_packages(where="src"),
package_dir={"": "src"},
classifiers=[
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
"License :: OSI Approved :: MIT License",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
],
python_requires=">=3.8",
install_requires=[
# Core dependencies
"numpy>=1.23.5",
"pyyaml>=6.0",
"jinja2>=3.1.5",
"requests>=2.32.0",
"urllib3>=2.3.0",
# Code analysis tools
"astor>=0.8.1",
"code2flow>=2.5.1",
"pydeps>=3.0.0",
# AI/LLM related dependencies
"anthropic>=0.45.0",
"openai>=1.60.1",
"langchain-anthropic>=0.3.4",
"langchain-openai>=0.3.2",
"langchain-core>=0.3.31",
"langgraph>=0.2.67",
"tiktoken>=0.8.0",
"transformers>=4.48.0",
"huggingface-hub>=0.28.0",
"google-generativeai>=0.6.0",
# Utility packages
"tqdm>=4.67.1",
"tabulate>=0.9.0",
"colorama>=0.4.6",
"termcolor>=2.5.0",
"pydantic>=2.10.0",
# Web requirements
"flask>=3.1.0",
"flask-socketio>=5.5.1",
"eventlet>=0.39.0",
"python-socketio>=5.12.1",
"python-engineio>=4.11.2",
"bidict>=0.23.0",
"dnspython>=2.7.0",
"six>=1.16.0",
# CUDA requirements
"torch>=2.0.0",
"accelerate>=1.4.0",
],
extras_require={
"dev": dev_requires,
"web": web_requires, # Keep for potential compatibility, now included in core
"visualization": visualization_requires,
"cuda": cuda_requires, # Keep for potential compatibility, now included in core
"all": all_requires,
}
)
================================================
FILE: src/DocstringGenerator.egg-info/PKG-INFO
================================================
Metadata-Version: 2.2
Name: DocstringGenerator
Version: 0.1.0
Summary: DocAgent for High-quality docstring generation in Large-scale Python projects
Author: Dayu Yang
Author-email: dayuyang@meta.com
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.23.5
Requires-Dist: pyyaml>=6.0
Requires-Dist: jinja2>=3.1.5
Requires-Dist: requests>=2.32.0
Requires-Dist: urllib3>=2.3.0
Requires-Dist: astor>=0.8.1
Requires-Dist: code2flow>=2.5.1
Requires-Dist: pydeps>=3.0.0
Requires-Dist: anthropic>=0.45.0
Requires-Dist: openai>=1.60.1
Requires-Dist: langchain-anthropic>=0.3.4
Requires-Dist: langchain-openai>=0.3.2
Requires-Dist: langchain-core>=0.3.31
Requires-Dist: langgraph>=0.2.67
Requires-Dist: tiktoken>=0.8.0
Requires-Dist: transformers>=4.48.0
Requires-Dist: huggingface-hub>=0.28.0
Requires-Dist: google-generativeai>=0.6.0
Requires-Dist: tqdm>=4.67.1
Requires-Dist: tabulate>=0.9.0
Requires-Dist: colorama>=0.4.6
Requires-Dist: termcolor>=2.5.0
Requires-Dist: pydantic>=2.10.0
Requires-Dist: flask>=3.1.0
Requires-Dist: flask-socketio>=5.5.1
Requires-Dist: eventlet>=0.39.0
Requires-Dist: python-socketio>=5.12.1
Requires-Dist: python-engineio>=4.11.2
Requires-Dist: bidict>=0.23.0
Requires-Dist: dnspython>=2.7.0
Requires-Dist: six>=1.16.0
Requires-Dist: torch>=2.0.0
Requires-Dist: accelerate>=1.4.0
Provides-Extra: dev
Requires-Dist: pytest>=8.3.4; extra == "dev"
Requires-Dist: pytest-cov>=2.0; extra == "dev"
Requires-Dist: black>=22.0; extra == "dev"
Requires-Dist: flake8>=3.9; extra == "dev"
Provides-Extra: web
Requires-Dist: flask>=3.1.0; extra == "web"
Requires-Dist: flask-socketio>=5.5.1; extra == "web"
Requires-Dist: eventlet>=0.39.0; extra == "web"
Requires-Dist: python-socketio>=5.12.1; extra == "web"
Requires-Dist: python-engineio>=4.11.2; extra == "web"
Requires-Dist: bidict>=0.23.0; extra == "web"
Requires-Dist: dnspython>=2.7.0; extra == "web"
Requires-Dist: six>=1.16.0; extra == "web"
Provides-Extra: visualization
Requires-Dist: matplotlib>=3.10.0; extra == "visualization"
Requires-Dist: pygraphviz>=1.14; extra == "visualization"
Requires-Dist: networkx>=3.4.2; extra == "visualization"
Provides-Extra: cuda
Requires-Dist: torch>=2.0.0; extra == "cuda"
Requires-Dist: accelerate>=1.4.0; extra == "cuda"
Provides-Extra: all
Requires-Dist: pytest>=8.3.4; extra == "all"
Requires-Dist: pytest-cov>=2.0; extra == "all"
Requires-Dist: black>=22.0; extra == "all"
Requires-Dist: flake8>=3.9; extra == "all"
Requires-Dist: flask>=3.1.0; extra == "all"
Requires-Dist: flask-socketio>=5.5.1; extra == "all"
Requires-Dist: eventlet>=0.39.0; extra == "all"
Requires-Dist: python-socketio>=5.12.1; extra == "all"
Requires-Dist: python-engineio>=4.11.2; extra == "all"
Requires-Dist: bidict>=0.23.0; extra == "all"
Requires-Dist: dnspython>=2.7.0; extra == "all"
Requires-Dist: six>=1.16.0; extra == "all"
Requires-Dist: matplotlib>=3.10.0; extra == "all"
Requires-Dist: pygraphviz>=1.14; extra == "all"
Requires-Dist: networkx>=3.4.2; extra == "all"
Requires-Dist: torch>=2.0.0; extra == "all"
Requires-Dist: accelerate>=1.4.0; extra == "all"
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary
# DocAgent: Agentic Hierarchical Docstring Generation System
<p align="center">
<img src="assets/meta_logo_white.png" width="20%" alt="Meta Logo">
</p>
DocAgent is a system designed to generate high-quality, context-aware docstrings for Python codebases using a multi-agent approach and hierarchical processing.
## Table of Contents
- [Motivation](#motivation)
- [Methodology](#methodology)
- [Installation](#installation)
- [Components](#components)
- [Usage](#usage)
- [Data Handling](#data-handling)
- [Baselines](#baselines)
- [Development Notes](#development-notes)
## Motivation
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.
## Methodology
DocAgent employs two key strategies:
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.
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.
<img src="assets/system.png" width="100%" alt="System Overview">
For more details on the agentic framework, see the [Agent Component README](./src/agent/README.md).
## Installation
Detailed installation instructions using `pip` or `conda`, including optional dependencies and troubleshooting tips, can be found in [INSTALL.md](./INSTALL.md).
## Components
DocAgent is composed of several key parts:
- **[Core Agent Framework](./src/agent/README.md)**: Implements the multi-agent system (Reader, Searcher, Writer, Verifier, Orchestrator) responsible for the generation logic.
- **[Docstring Evaluator](./src/evaluator/README.md)**: Provides tools for evaluating docstring quality, primarily focusing on completeness based on static code analysis (AST).
- **[Generation Web UI](./src/web/README.md)**: A web interface for configuring, running, and monitoring the docstring *generation* process in real-time, visualizing agent activity and repository structure.
- **[Evaluation Web UI](./src/web_eval/README.md)**: A separate web interface for configuring and running docstring *evaluations*, assessing completeness and helpfulness (using LLMs).
## Usage
The primary ways to interact with DocAgent are:
1. **Generation Web UI**: Recommended for visualizing the generation process. Launch via:
```bash
python run_web_ui.py
```
Then access `http://localhost:5000` (or as configured). See the [Generation Web UI README](./src/web/README.md) for details.
2. **Evaluation Web UI**: Recommended for assessing docstring quality. Launch via:
```bash
cd src/web_eval
./start_server.sh # Or python app.py
```
Then access `http://localhost:5000` (or as configured). See the [Evaluation Web UI README](./src/web_eval/README.md) for details.
3. **Command Line (Generation)**: Run the generation process directly:
```bash
# Example: Run on a test repo, removing existing docstrings first
./tool/remove_docstrings.sh data/raw_test_repo
python generate_docstrings.py --repo-path data/raw_test_repo
```
Use `--help` for more options.
## Data Handling
Tools are included for managing datasets for evaluation:
- **GitHub Repository Downloader (`src/data/parse/downloader.py`)**: Finds and downloads GitHub repositories based on configurable criteria (language, stars, size, etc.).
- **Repository Selection (`experiments/select_repos.py`)**: Selects a diverse subset of downloaded repositories based on metrics like code size and complexity.
(See original README sections for detailed usage if needed).
## Baselines
A simple "copy and paste" baseline is implemented (`experiments/generate_docstrings_copy_and_paste.py`) for comparison. It sends isolated code components to an LLM without context.
```bash
# Example: Run baseline on a test repo
./tool/remove_docstrings.sh data/raw_test_repo
python experiments/generate_docstrings_copy_and_paste.py --repo-path data/raw_test_repo
```
## Development Notes
- Remember to activate your chosen environment (`pip` or `conda`).
- Use `pip install -e ".[dev]"` for development dependencies.
- Run tests using `pytest`.
- See [INSTALL.md](./INSTALL.md) for setting up system dependencies like GraphViz if needed for visualizations.
---
*This README provides a high-level overview. Please refer to the linked component READMEs and `INSTALL.md` for specific details.*
# Todo
- repo-level eval script
- argument vs parameter
- "Need more information" seems does not work for codellama34B and gemini.
- some repo depends on "not-install" package(ask you to install autogen after download the repo)
- Query should also search internally
- truncated "called by", especially Class (too long)
- Overkill issue
# For ACL Experiments
## Note
class evaluate:
- really means eval the init function (if has init)
## Data
### GitHub Repository Downloader
The project includes a GitHubRepoDownloader that automates the process of finding and downloading repositories for docstring generation tasks. This tool allows you to specify various criteria to target repositories that match your requirements.
#### Features:
- **Configurable Search Criteria**: Filter repositories by owner, creation date, language, stars, forks, size, and license.
- **Python Language Filtering**: Ensures downloaded repositories contain a minimum percentage of Python code (default: 80%).
- **Repository Metadata**: Automatically saves metadata about each downloaded repository.
- **Rate Limit Handling**: Respects GitHub API rate limits to avoid throttling.
- **Logging**: Comprehensive logging of the download process.
#### Usage:
To download repositories, create a configuration file and run the downloader:
```bash
python -m src.data.parse.downloader
```
#### Configuration:
Create a YAML configuration file with the following structure:
```yaml
# GitHub authentication
GITHUB_TOKEN: "your-github-token"
# Output directory
output_directory: "data/downloaded_repos"
# Repository limits
max_repos: 10
skip_archived: true
skip_forks: true
min_python_percentage: 80 # Minimum percentage of Python code required
# Search criteria
search_criteria:
language: "python"
stars:
min: 100
forks:
min: 10
dates:
created_after: "2020-01-01"
owners:
- "username1"
- "org_name"
```
The downloader will:
1. Search GitHub repositories matching your criteria
2. Check if each repository meets the Python percentage requirement
3. Clone qualifying repositories to the specified output directory
4. Save repository metadata for further analysis
### Repository Selection
After downloading repositories, you may want to select a diverse subset for analysis. The project includes a repository selection tool that helps you choose repositories with varying characteristics:
#### Features:
- **Diversity-Based Selection**: Select repositories based on code size and structural complexity.
- **Code Size Metrics**: Calculates the number of Python files and total lines of code.
- **Topological Complexity**: Measures the depth of the repository directory structure.
- **Visualization**: Generates scatter plots showing the distribution of selected repositories.
#### Usage:
To select repositories from your downloaded collection:
```bash
python -m experiments.select_repos
```
#### Process:
The selection process follows these steps:
1. Analyzes each repository to extract metrics (Python files count, total lines, directory depth)
2. Normalizes the metrics to ensure fair comparison
3. Creates clusters of repositories with similar characteristics
4. Selects representatives from each cluster to ensure diversity
5. Generates a visualization of the selection results
This approach ensures that your analysis includes repositories with varying sizes and complexity levels, providing a more comprehensive evaluation of docstring generation techniques.
## Baseline
### Copy and Paste
We implemented a simple "copy and paste" baseline system that mimics the approach of users copying code components and pasting them directly to an LLM interface. This baseline:
1. Extracts individual code components (functions, classes, methods) from Python files
2. Sends only the component's source code to an LLM without any surrounding context
3. Asks the LLM to generate a docstring based solely on that isolated component
4. Inserts the generated docstring back into the code
This baseline serves as a comparison point to demonstrate the effectiveness of our full agentic hierarchical system, which considers dependency relationships and broader context when generating docstrings.
To run the baseline system:
```bash
clear
./tool/remove_docstrings.sh data/raw_test_repo
python experiments/generate_docstrings_copy_and_paste.py --repo-path data/raw_test_repo
```
The baseline uses the same configuration file (agent_config.yaml) as the main system, so it can work with any supported LLM (Claude, OpenAI, HuggingFace).
To run in placeholder mode (no actual LLM calls):
```bash
python experiments/generate_docstrings_copy_and_paste.py --repo-path data/raw_test_repo --test-mode placeholder
```
To overwrite existing docstrings:
```bash
python experiments/generate_docstrings_copy_and_paste.py --repo-path data/raw_test_repo --overwrite-docstrings
```
### Main Experiments
## Motivation:
In the realm of large-scale software repositories, the presence of high-quality, user-oriented docstrings is crucial for maintaining code readability, usability, and maintainability. A well-crafted docstring should not only provide comprehensive details about parameters, return values, exceptions, and usage examples but also clearly articulate the purpose of the function or class within the broader context of the repository. This includes explaining when and how to use the function or class, as well as its relationship to other components in the codebase.
Despite the importance of such documentation, current large language models (LLMs) often fall short in generating docstrings that meet these expectations. Common issues include the production of redundant or superficial commentary, a failure to highlight the underlying rationale behind implementation choices, and the omission of crucial constraints and assumptions. These shortcomings can lead to misunderstandings and inefficiencies for developers who rely on these docstrings for guidance.
The challenge, therefore, is to develop methods that enable the generation of high-quality docstrings that are both informative and concise, avoiding redundancy by not reiterating information that can be inferred from the code itself, such as parameter types when type hints are present. Addressing these challenges is essential for enhancing the utility of docstrings in large-scale repositories, ultimately contributing to more efficient and effective software development processes.
## Challenges and Limitations of Existing Docstring Generation System:
The task of generating high-quality docstrings in large-scale repositories presents several significant challenges. One of the primary difficulties lies in the evaluation of docstring quality. There is inherent ambiguity in assessing what constitutes a "good" docstring, as gold-standard data is scarce. Even highly-rated repositories often contain docstrings that are either inadequate or only partially effective, complicating the establishment of reliable benchmarks for quality assessment.
Another challenge is the limitation imposed by the context window of large language models (LLMs). It is impractical to include an entire repository in a single prompt, necessitating a focus on selecting and summarizing relevant information. Determining what is "relevant" is crucial for providing the LLM with a comprehensive understanding of the purpose of the focal function or class. This involves discerning which aspects of the codebase should be included to give the LLM a "global sense" of the function's or class's role and significance.
Furthermore, there is a "chicken and egg" problem inherent in this task. Generating high-quality docstrings requires a well-rounded understanding of the context in which the focal function or class operates. However, the context itself often lacks sufficient documentation to clearly convey its purpose and interrelations. This lack of existing high-quality docstrings in the surrounding code complicates the process of generating new ones, as the foundational understanding needed to inform the generation process is itself incomplete.
Addressing these challenges is essential for advancing the capability of LLMs to produce docstrings that are not only informative and concise but also contextually aware and aligned with the broader objectives of the codebase.
## Methodology:
To address the challenges of generating high-quality docstrings in large-scale repositories, we propose a hierarchical traversal approach combined with an agentic system composed of specialized roles: reader, searcher, writer, and verifier. This methodology is designed to systematically and efficiently produce comprehensive and contextually aware docstrings.
Hierarchical Traverse
The hierarchical traverse principle is central to our approach. By prioritizing the generation of docstrings for source code files with fewer dependencies, we aim to build a solid foundation of well-documented base classes and utility functions before tackling more complex implementations. This strategy effectively addresses the "chicken and egg" problem by ensuring that the foundational components of the codebase are well-understood and documented first. Unlike existing systems that generate docstrings in a random order, our method provides a structured and logical progression through the codebase.
Agentic System
Our agentic system is designed to facilitate the docstring generation process through a series of coordinated roles:
- Reader: The reader initiates the process by examining the focal code component and identifying any additional internal or external information needed to understand its purpose and context. If further information is required, the reader sends a request to the searcher.
- Searcher: The searcher traverses the dependency graph to gather relevant information, both from within the codebase and from open-internet sources if necessary. This information is then used to update the context state, providing a more comprehensive understanding of the focal component.
- Writer: Once the context is deemed sufficient, the reader passes the focal code component and its context to the writer. The - writer drafts the docstring, ensuring it adheres to the specified quality and instructional guidelines.
- Verifier: The verifier conducts a final quality check of the drafted docstring. If formatting issues are detected, the docstring is returned to the writer for revision. If additional context is needed to enhance informativeness, the process returns to the reader for further information gathering.
This iterative and collaborative approach ensures that each docstring is not only accurate and informative but also contextually aligned with the broader objectives of the codebase. By leveraging the strengths of each agent, our methodology provides a robust framework for generating high-quality documentation in large-scale repositories.
# For Test
The easiest way to interact with DocAgent is through Web App. Assuming the Web App is hosted on remote server.
## Docstring Generation System (Agentic + Hierarchical)
Without Web UI:
```bash
clear
./tool/remove_docstrings.sh data/raw_test_repo
python generate_docstrings.py --repo-path data/raw_test_repo --test-mode context_print
```
With Web UI:
```bash
clear
./tool/remove_docstrings.sh data/raw_test_repo
python run_web_ui.py --host 0.0.0.0 --port 5000
```
## Docstring Eval system
Without WebUI: Manual run the test files under `test/evaluator`.
With WebUI:
```bash
python src/web_eval/app.py --host 0.0.0.0 --port 5001
```
## Test Hierarchical Generation Only (no LLM call)
For test hierarchical generation only (no LLM call), run the following command: (testing on `data/test_repo_vm` and `data/downloaded_repos/AutoSurvey`)
```bash
./tool/remove_docstrings.sh data/downloaded_repos/AutoSurvey
clear
python generate_docstrings.py --repo-path data/downloaded_repos/AutoSurvey --test-mode
bash tool/visualize.sh output/dependency_graphs/dependency_graph.json output/dependency_graphs/dependency_graph_visualization.png
```
## Depreciated Tests
Test Completeness
```bash
python test/evaluator/test_completeness.py data/downloaded_repos/AutoSurvey/src/agents/judge.py
```
Test reader-searcher communication.
```bash
python test/agent/depreciated_test_orchestrator.py --mode reader-searcher --verbose-context
```
Remove all docstrings from a repository.
```bash
./tool/remove_docstrings.sh <repo_path>
```
Test hierarchical generation.
```bash
python generate_docstrings.py --repo-path data/test_repo_vm --test-mode
```
Visualize dependency graph.
```bash
bash tool/visualize.sh
```
# Installation
## Create Config File
Create a config folder `config/` and a config file `agent_config.yaml`under `config/`. e.g. `config/agent_config.yaml`:
The structure of config is as follows:
```bash
llm:
type: "claude" # Options: openai, claude, huggingface
api_key: "your-anthropic-api-key-here" # Replace with your Anthropic API key
model: "claude-3-5-haiku-latest"
temperature: 0.1
max_output_tokens: 4096
# Flow control parameters
flow_control:
max_reader_search_attempts: 2 # Maximum times reader can call searcher
max_verifier_rejections: 3 # Maximum times verifier can reject a docstring
status_sleep_time: 3 # Time to sleep between status updates (seconds)
# Perplexity API configuration
perplexity:
api_key: "your-perplexity-api-key-here" # Replace with your Perplexity API key
model: "sonar" # Default model
temperature: 0.1
max_output_tokens: 4096
```
## Installation
### Basic Installation
To install the basic package with core dependencies:
```bash
pip install -e .
```
### Install with Additional Features
You can install the package with additional optional dependencies:
```bash
# For development tools (pytest, black, flake8)
pip install -e ".[dev]"
# For web UI components
pip install -e ".[web]"
# For visualization tools
pip install -e ".[visualization]"
# For CUDA support
pip install -e ".[cuda]"
# For all optional dependencies
pip install -e ".[all]"
```
You can also combine multiple optional dependencies:
```bash
pip install -e ".[web,visualization]"
```
## Access Web UI from Local (if running on remote server)
## Running Docstring Generation System
In remote:
```bash
python run_web_ui.py --host 0.0.0.0 --port 5000
```
This tells Flask to listen on all network interfaces, not just the loopback interface.
In local:
```bash
ssh -L 5000:localhost:5000 <remote_host>
```
For example, for devserver, `ssh -L 5000:localhost:5000 dayuyang@devgpu003.rva5.facebook.com`.
This command creates a tunnel from your local port 5000 to port 5000 on the remote server. After running this command, you can open your browser and go to http://localhost:5000 to access the web interface running on the remote server.
kill any program running on port 5000:
```bash
lsof -i :5000 | awk 'NR>1 {print $2}' | sort -u | xargs -r kill
```
## Running Docstring Eval System
In remote:
```bash
python src/web_eval/app.py --host 0.0.0.0 --port 5001
```
In local:
```
ssh -L 5001:localhost:5001 dayuyang@devgpu003.rva5.facebook.com
```
If run both:
in local, run:
```bash
ssh -L 5000:localhost:5000 -L 5001:localhost:5001 -L 5002:localhost:5002 dayuyang@devgpu003.rva5.facebook.com
# 5002 for backup
```
## Serve local LLM
First install vllm
```bash
pip install vllm
```
Run `bash serve_local_llm.sh`
# Concerns
- hierachical generation
- Circular dependencies
- Import from external source?
- assuming "external source" is well-known library and LLM should already know about it?
# TO FIX/ADD/IMPROVE
- If already has docstring, skip. (or al least give an option to skip)
- Improve generation instructions (the LLM will strictly follow the instructions, leading the generation usually too long.)
- high-quality docstring also needs to be concise.
- add time out warning (if stucked...)
- Claude has rate limit: `50,000 input tokens per minute per organization`
- add error handling capability system wise
- add price calculation
- simple code will no need to use this tool.
- before system, a small LLM/ determinstic way to determinate if using the system. (balance between efficiency and effectiveness)
now the logic is file-level, function-level, method-level, class-level.
# Evaluator
## Completeness
## Helpfulness
For Summary, Description, Arguments, Parameters, Attributes, each docstring component is evaluated on a 1-5 scale (POOR to EXCELLENT):
For Example, the docstring is evaluated on Binary scale (0 or 1).
- Evaluates if docstring examples enable users to correctly use the code by comparing predicted usage against ground truth.
# Vulnerability
1. helpfulness description, when class is too long, may need truncate.
2. when evaluating parameters/arguments/attributes, the input context (class/function signatures) should be reasonably sized to avoid LLM token limits.
# Logic Control flow (process function under orchestrator.py)
once searcher is called, reader's memory is refreshed.
once more context is needed by judge from verifier, writer, verifier's memory is refreshed.
# Note
When evaluating examples, the signature must contain decorator `@staticmethod` or `@classmethod`.
when writing docstring for class, first write docstring for __init__ method, then write docstring for other methods, finally write docstring for class. (provide full class code as code component when writing docstring for class)
Method is extremely difficult to handle. Now only support self.method(), instance.method() and ClassName.method(). See `get_child_method` under `CallGraphBuilder` for more details.
Error handle: ask reader: if "XXX is not accessible", do not ask the same code component again. If unsuccessful, callgraphbuilder will return something like "XXX is not accessible".
for LLM generated docstring, no triple quotes (\"\"\") is added originally.
For generate_docstrings.py. Add features:
Multiple Passes ("Category" Approach):
• We split docstring generation for each file into three passes, in this order:
(a) Top-level functions (i.e., "function")
(b) Methods inside classes (i.e., "method")
(c) Classes (i.e., "class")
Each pass visits all .py files in the repo.
Immediate File Rewrite After Each Code Component:
• In each pass, we repeatedly parse a file, gather all code components of the chosen category in ascending line order, pick off the first component, generate a docstring, and immediately rewrite the file. Then we re-parse the updated file before moving on to the next component.
• This ensures that each code component is added in the final version of the file before the next code component's docstring is generated. It also meets the request for "refresh the python file after each generation for a single code component."
However, this approach is more computationally expensive than generating all docstrings in memory and rewriting once per file, but it achieves the desired incremental rewriting and strict function → method → class ordering.
Why the python file is not updated after the docstring is generated for each code component?
- because updating file needs re-parsing the file and rebuild the AST, which is expensive.
Dependency clarification for
# Future Work
human in the loop:
- human can be the judge and can provide more information to the system.
================================================
FILE: src/DocstringGenerator.egg-info/SOURCES.txt
================================================
README.md
setup.py
src/DocstringGenerator.egg-info/PKG-INFO
src/DocstringGenerator.egg-info/SOURCES.txt
src/DocstringGenerator.egg-info/dependency_links.txt
src/DocstringGenerator.egg-info/requires.txt
src/DocstringGenerator.egg-info/top_level.txt
src/agent/__init__.py
src/agent/base.py
src/agent/orchestrator.py
src/agent/reader.py
src/agent/searcher.py
src/agent/verifier.py
src/agent/workflow.py
src/agent/writer.py
src/agent/llm/__init__.py
src/agent/llm/base.py
src/agent/llm/claude_llm.py
src/agent/llm/factory.py
src/agent/llm/gemini_llm.py
src/agent/llm/huggingface_llm.py
src/agent/llm/openai_llm.py
src/agent/llm/rate_limiter.py
src/dependency_analyzer/__init__.py
src/dependency_analyzer/ast_parser.py
src/dependency_analyzer/topo_sort.py
src/evaluator/__init__.py
src/evaluator/base.py
src/evaluator/completeness.py
src/evaluator/evaluation_common.py
src/evaluator/helpfulness_attributes.py
src/evaluator/helpfulness_description.py
src/evaluator/helpfulness_evaluator.py
src/evaluator/helpfulness_evaluator_ablation.py
src/evaluator/helpfulness_examples.py
src/evaluator/helpfulness_parameters.py
src/evaluator/helpfulness_summary.py
src/evaluator/segment.py
src/evaluator/truthfulness.py
src/visualizer/__init__.py
src/visualizer/progress.py
src/visualizer/status.py
src/visualizer/web_bridge.py
src/web/__init__.py
src/web/app.py
src/web/config_handler.py
src/web/process_handler.py
src/web/run.py
src/web/visualization_handler.py
================================================
FILE: src/DocstringGenerator.egg-info/dependency_links.txt
================================================
================================================
FILE: src/DocstringGenerator.egg-info/requires.txt
================================================
numpy>=1.23.5
pyyaml>=6.0
jinja2>=3.1.5
requests>=2.32.0
urllib3>=2.3.0
astor>=0.8.1
code2flow>=2.5.1
pydeps>=3.0.0
anthropic>=0.45.0
openai>=1.60.1
langchain-anthropic>=0.3.4
langchain-openai>=0.3.2
langchain-core>=0.3.31
langgraph>=0.2.67
tiktoken>=0.8.0
transformers>=4.48.0
huggingface-hub>=0.28.0
google-generativeai>=0.6.0
tqdm>=4.67.1
tabulate>=0.9.0
colorama>=0.4.6
termcolor>=2.5.0
pydantic>=2.10.0
flask>=3.1.0
flask-socketio>=5.5.1
eventlet>=0.39.0
python-socketio>=5.12.1
python-engineio>=4.11.2
bidict>=0.23.0
dnspython>=2.7.0
six>=1.16.0
torch>=2.0.0
accelerate>=1.4.0
[all]
pytest>=8.3.4
pytest-cov>=2.0
black>=22.0
flake8>=3.9
flask>=3.1.0
flask-socketio>=5.5.1
eventlet>=0.39.0
python-socketio>=5.12.1
python-engineio>=4.11.2
bidict>=0.23.0
dnspython>=2.7.0
six>=1.16.0
matplotlib>=3.10.0
pygraphviz>=1.14
networkx>=3.4.2
torch>=2.0.0
accelerate>=1.4.0
[cuda]
torch>=2.0.0
accelerate>=1.4.0
[dev]
pytest>=8.3.4
pytest-cov>=2.0
black>=22.0
flake8>=3.9
[visualization]
matplotlib>=3.10.0
pygraphviz>=1.14
networkx>=3.4.2
[web]
flask>=3.1.0
flask-socketio>=5.5.1
eventlet>=0.39.0
python-socketio>=5.12.1
python-engineio>=4.11.2
bidict>=0.23.0
dnspython>=2.7.0
six>=1.16.0
================================================
FILE: src/DocstringGenerator.egg-info/top_level.txt
================================================
agent
dependency_analyzer
evaluator
visualizer
web
================================================
FILE: src/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
================================================
FILE: src/agent/README.md
================================================
# Agent Framework for Docstring Generation
This directory contains the core components of the multi-agent system responsible for generating high-quality docstrings for code components.
## Overview
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.
The main workflow is initiated via the `generate_docstring` function in `workflow.py`.
## Agents
1. **`BaseAgent` (`base.py`)**
* **Role:** Abstract base class for all agents.
* **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.
2. **`Reader` (`reader.py`)**
* **Role:** Contextual Analysis and Information Needs Assessment.
* **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.
3. **`Searcher` (`searcher.py`)**
* **Role:** Information Retrieval.
* **Functionality:** Acts upon the requests generated by the `Reader`. It retrieves the specified information by:
* Querying the internal codebase using AST analysis (`ASTNodeAnalyzer`) and dependency graphs.
* Performing external web searches via APIs (e.g., `PerplexityAPI`).
* Returns the gathered context in a structured format.
4. **`Writer` (`writer.py`)**
* **Role:** Docstring Generation.
* **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 (`<DOCSTRING>`).
5. **`Verifier` (`verifier.py`)**
* **Role:** Quality Assurance.
* **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.
6. **`Orchestrator` (`orchestrator.py`)**
* **Role:** Workflow Management.
* **Functionality:** Coordinates the entire process. It manages the sequence of agent interactions:
* Calls `Reader` to assess context needs.
* Calls `Searcher` iteratively if more context is requested (up to a limit).
* Calls `Writer` to generate the docstring.
* Calls `Verifier` to evaluate the docstring.
* Manages revision loops based on `Verifier` feedback, potentially involving further searches or refinement by the `Writer` (up to a limit).
* Handles context accumulation, token limit constraints, and status visualization.
## Supporting Files
* **`workflow.py`:** Provides the primary entry point function `generate_docstring` to initiate the docstring generation process for a given code component.
* **`__init__.py`:** Makes the `agent` directory a Python package.
* **`llm/`:** Contains LLM-related code, including the `LLMFactory` and base LLM classes.
* **`tool/`:** Contains tools used by agents, such as the `ASTNodeAnalyzer` for internal code traversal and the `PerplexityAPI` wrapper for external search.
================================================
FILE: src/agent/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
# Import only essential components to avoid circular imports
from .reader import CodeComponentType
# Explicitly list what should be accessible, but don't import until needed
# to prevent circular imports
__all__ = ['generate_docstring', 'CodeComponentType']
# Lazy load generate_docstring when it's actually needed
def __getattr__(name):
if name == 'generate_docstring':
from .workflow import generate_docstring
return generate_docstring
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
================================================
FILE: src/agent/base.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, List
import os
from pathlib import Path
from .llm.factory import LLMFactory
from .llm.base import BaseLLM
class BaseAgent(ABC):
"""Base class for all agents in the docstring generation system."""
def __init__(self, name: str, config_path: Optional[str] = None):
"""Initialize the base agent.
Args:
name: The name of the agent
config_path: Optional path to the configuration file
"""
self.name = name
self._memory: list[Dict[str, Any]] = []
# Initialize LLM and parameters from config
self.llm, self.llm_params = self._initialize_llm(name, config_path)
def _initialize_llm(self, agent_name: str, config_path: Optional[str] = None) -> tuple[BaseLLM, Dict[str, Any]]:
"""Initialize the LLM for this agent.
Args:
agent_name: Name of the agent
config_path: Optional path to the configuration file
Returns:
Tuple of (Initialized LLM instance, LLM parameters dictionary)
"""
# Load configuration
if config_path is None:
config_path = "config/agent_config.yaml"
print(f"Using default config from {config_path}")
config = LLMFactory.load_config(config_path)
# Check for agent-specific configuration
agent_config = config.get("agent_llms", {}).get(agent_name.lower())
# Use agent-specific config if available, otherwise use default
llm_config = agent_config if agent_config else config.get("llm", {})
# Verify api_key is provided in config
if ("api_key" not in llm_config or not llm_config["api_key"]) and (llm_config["type"] not in ["huggingface", "local"]):
raise ValueError("API key must be specified directly in the config file")
# Extract LLM parameters
llm_params = {
"max_output_tokens": llm_config.get("max_output_tokens", 4096),
"temperature": llm_config.get("temperature", 0.1),
"model": llm_config.get("model")
}
return LLMFactory.create_llm(llm_config), llm_params
def add_to_memory(self, role: str, content: str) -> None:
"""Add a message to the agent's memory.
Args:
role: The role of the message sender (e.g., 'system', 'user', 'assistant')
content: The content of the message
"""
assert content is not None and content != "", "Content cannot be empty"
self._memory.append(self.llm.format_message(role, content))
def refresh_memory(self, new_memory: list[Dict[str, Any]]) -> None:
"""Replace the current memory with new memory.
Args:
new_memory: The new memory to replace the current memory
"""
self._memory = [
self.llm.format_message(msg["role"], msg["content"])
for msg in new_memory
]
def clear_memory(self) -> None:
"""Clear the agent's memory."""
self._memory = []
@property
def memory(self) -> list[Dict[str, Any]]:
"""Get the agent's memory.
Returns:
The agent's memory as a list of message dictionaries
"""
return self._memory.copy()
def generate_response(self, messages: Optional[List[Dict[str, Any]]] = None) -> str:
"""Generate a response using the agent's LLM and memory.
Args:
messages: Optional list of messages to use instead of memory
Returns:
Generated response text
"""
return self.llm.generate(
messages=messages if messages is not None else self._memory,
temperature=self.llm_params["temperature"],
max_tokens=self.llm_params["max_output_tokens"]
)
@abstractmethod
def process(self, *args, **kwargs) -> Any:
"""Process the input and generate output.
This method should be implemented by each specific agent.
"""
pass
================================================
FILE: src/agent/llm/__init__.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from .base import BaseLLM
from .openai_llm import OpenAILLM
from .claude_llm import ClaudeLLM
from .huggingface_llm import HuggingFaceLLM
from .gemini_llm import GeminiLLM
from .factory import LLMFactory
__all__ = [
'BaseLLM',
'OpenAILLM',
'ClaudeLLM',
'HuggingFaceLLM',
'GeminiLLM',
'LLMFactory'
]
================================================
FILE: src/agent/llm/base.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional
class BaseLLM(ABC):
"""Base class for LLM wrappers."""
@abstractmethod
def generate(
self,
messages: List[Dict[str, str]],
temperature: float = 0.7,
max_output_tokens: Optional[int] = None
) -> str:
"""Generate a response from the LLM.
Args:
messages: List of message dictionaries with 'role' and 'content' keys
temperature: Sampling temperature (0.0 to 1.0)
max_output_tokens: Maximum number of tokens to generate
Returns:
The generated response text
"""
pass
@abstractmethod
def format_message(self, role: str, content: str) -> Dict[str, str]:
"""Format a message for the specific LLM API.
Args:
role: The role of the message sender
content: The content of the message
Returns:
Formatted message dictionary
"""
pass
================================================
FILE: src/agent/llm/claude_llm.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import List, Dict, Any, Optional
import anthropic
from .base import BaseLLM
from .rate_limiter import RateLimiter
import logging
class ClaudeLLM(BaseLLM):
"""Anthropic Claude API wrapper."""
def __init__(
self,
api_key: str,
model: str,
rate_limits: Optional[Dict[str, Any]] = None
):
"""Initialize Claude LLM.
Args:
api_key: Anthropic API key
model: Model identifier (e.g., "claude-3-sonnet-20240229")
rate_limits: Optional dictionary with rate limit settings
"""
self.client = anthropic.Anthropic(api_key=api_key)
self.model = model
# Default rate limits for Claude 3.7 Sonnet
default_limits = {
"requests_per_minute": 50,
"input_tokens_per_minute": 20000,
"output_tokens_per_minute": 8000,
"input_token_price_per_million": 3.0,
"output_token_price_per_million": 15.0
}
# Use provided rate limits or defaults
limits = rate_limits or default_limits
# Initialize rate limiter
self.rate_limiter = RateLimiter(
provider="Claude",
requests_per_minute=limits.get("requests_per_minute", default_limits["requests_per_minute"]),
input_tokens_per_minute=limits.get("input_tokens_per_minute", default_limits["input_tokens_per_minute"]),
output_tokens_per_minute=limits.get("output_tokens_per_minute", default_limits["output_tokens_per_minute"]),
input_token_price_per_million=limits.get("input_token_price_per_million", default_limits["input_token_price_per_million"]),
output_token_price_per_million=limits.get("output_token_price_per_million", default_limits["output_token_price_per_million"])
)
def _count_tokens(self, text: str) -> int:
"""Count tokens in a string using Claude's tokenizer.
Args:
text: Text to count tokens for
Returns:
Token count
"""
if not text:
return 0
try:
# Format text as a message for token counting
count = self.client.beta.messages.count_tokens(
model=self.model,
messages=[
{"role": "user", "content": text}
]
)
return count.input_tokens
except Exception as e:
# Log the error but don't fail
logging.warning(f"Failed to count tokens with Claude tokenizer: {e}")
# Fallback: rough estimate if tokenizer fails
return len(text.split()) * 1.3
def _count_messages_tokens(self, messages: List[Dict[str, str]], system_message: Optional[str] = None) -> int:
"""Count tokens in message list with optional system message.
Args:
messages: List of message dictionaries
system_message: Optional system message
Returns:
Total token count
"""
if not messages:
return 0
# Convert messages to Claude format
claude_messages = [self._convert_to_claude_message(msg) for msg in messages
if msg["role"] != "system"]
# Format system message if provided
system_content = None
if system_message:
system_content = system_message
try:
# Use the API to count tokens for all messages at once
count = self.client.beta.messages.count_tokens(
model=self.model,
messages=claude_messages,
system=system_content
)
return count.input_tokens
except Exception as e:
# Log the error but don't fail
logging.warning(f"Failed to count tokens with Claude tokenizer: {e}")
# Fallback: count tokens individually
total_tokens = 0
for msg in claude_messages:
if "content" in msg and msg["content"]:
total_tokens += self._count_tokens(msg["content"])
# Add system message tokens if provided
if system_message:
total_tokens += self._count_tokens(system_message)
# Add overhead for message formatting
total_tokens += 10 * len(claude_messages) # Add ~10 tokens per message for formatting
return total_tokens
def generate(
self,
messages: List[Dict[str, str]],
temperature: float,
max_tokens: Optional[int]
) -> str:
"""Generate a response using Claude API with rate limiting.
Args:
messages: List of message dictionaries
temperature: Sampling temperature
max_output_tokens: Maximum tokens to generate
Returns:
Generated response text
"""
# Extract system message if present
system_message = None
chat_messages = []
for msg in messages:
if msg["role"] == "system":
system_message = msg["content"]
else:
chat_messages.append(self._convert_to_claude_message(msg))
# Count input tokens
input_tokens = self._count_messages_tokens(messages, system_message)
# Wait if we're approaching rate limits (estimate output tokens as max_output_tokens)
self.rate_limiter.wait_if_needed(input_tokens, max_tokens)
# Make the API call
response = self.client.messages.create(
model=self.model,
messages=chat_messages,
system=system_message,
temperature=temperature,
max_tokens=max_tokens
)
result_text = response.content[0].text
# Count output tokens and record request
output_tokens = self._count_tokens(result_text)
self.rate_limiter.record_request(input_tokens, output_tokens)
return result_text
def format_message(self, role: str, content: str) -> Dict[str, str]:
"""Format message for Claude API.
Args:
role: Message role (system, user, assistant)
content: Message content
Returns:
Formatted message dictionary
"""
# Store in standard format, conversion happens in generate()
return {"role": role, "content": content}
def _convert_to_claude_message(self, message: Dict[str, str]) -> Dict[str, str]:
"""Convert standard message format to Claude's format.
Args:
message: Standard format message
Returns:
Claude format message
"""
role_mapping = {
"user": "user",
"assistant": "assistant"
}
role = role_mapping[message["role"]]
content = message["content"]
return {"role": role, "content": content}
================================================
FILE: src/agent/llm/factory.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import Dict, Any, Optional
from pathlib import Path
import yaml
from .base import BaseLLM
from .openai_llm import OpenAILLM
from .claude_llm import ClaudeLLM
from .huggingface_llm import HuggingFaceLLM
from .gemini_llm import GeminiLLM
class LLMFactory:
"""Factory class for creating LLM instances."""
@staticmethod
def create_llm(config: Dict[str, Any]) -> BaseLLM:
"""Create an LLM instance based on configuration.
Args:
config: Configuration dictionary containing LLM settings
Returns:
An instance of BaseLLM
Raises:
ValueError: If the LLM type is not supported
"""
llm_type = config["type"].lower()
model = config.get("model")
if not model:
raise ValueError("Model must be specified in the config file")
# Extract rate limit settings from config
# First check if there are specific rate limits in the LLM config
rate_limits = config.get("rate_limits", {})
# If not, check if there are global rate limits for this provider type
global_config = LLMFactory.load_config()
if not rate_limits and "rate_limits" in global_config:
# Map LLM types to provider names in rate_limits section
provider_map = {
"openai": "openai",
"claude": "claude",
"gemini": "gemini"
}
provider_key = provider_map.get(llm_type, llm_type)
provider_limits = global_config.get("rate_limits", {}).get(provider_key, {})
if provider_limits:
rate_limits = provider_limits
if llm_type == "openai":
return OpenAILLM(
api_key=config["api_key"],
model=model,
rate_limits=rate_limits
)
elif llm_type == "claude":
return ClaudeLLM(
api_key=config["api_key"],
model=model,
rate_limits=rate_limits
)
elif llm_type == "gemini":
return GeminiLLM(
api_key=config["api_key"],
model=model,
rate_limits=rate_limits
)
elif llm_type == "huggingface":
return HuggingFaceLLM(
model_name=model,
device=config.get("device", "cuda"),
torch_dtype=config.get("torch_dtype", "float16")
)
else:
raise ValueError(f"Unsupported LLM type: {llm_type}")
@staticmethod
def load_config(config_path: Optional[str] = None) -> Dict[str, Any]:
"""Load LLM configuration from file.
Args:
config_path: Path to the configuration file. If None, uses default path.
Returns:
Configuration dictionary
Raises:
FileNotFoundError: If the configuration file doesn't exist
"""
if config_path is None:
config_path = str(Path(__file__).parent.parent.parent.parent / "config" / "agent_config.yaml")
if not Path(config_path).exists():
raise FileNotFoundError(f"Configuration file not found: {config_path}")
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
return config
================================================
FILE: src/agent/llm/gemini_llm.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import List, Dict, Any, Optional
import tiktoken
import google.generativeai as genai
from .base import BaseLLM
from .rate_limiter import RateLimiter
class GeminiLLM(BaseLLM):
"""Google Gemini API wrapper."""
def __init__(
self,
api_key: str,
model: str,
rate_limits: Optional[Dict[str, Any]] = None
):
"""Initialize Gemini LLM.
Args:
api_key: Google API key
model: Model identifier (e.g., "gemini-1.5-flash", "gemini-1.5-pro")
rate_limits: Optional dictionary with rate limit settings
"""
genai.configure(api_key=api_key)
self.model_name = model
self.model = genai.GenerativeModel(model)
try:
# Initialize tokenizer for token counting
# Gemini doesn't have a direct tokenizer in the public API
# Using tiktoken cl100k_base as a reasonable approximation
self.tokenizer = tiktoken.get_encoding("cl100k_base")
except:
# Fallback to basic word counting if tokenizer fails
self.tokenizer = None
# Default rate limits for Gemini (adjust based on actual API limits)
default_limits = {
"requests_per_minute": 60,
"input_tokens_per_minute": 100000,
"output_tokens_per_minute": 50000,
"input_token_price_per_million": 0.125, # Approximate for gemini-1.5-flash
"output_token_price_per_million": 0.375 # Approximate for gemini-1.5-flash
}
# Use provided rate limits or defaults
limits = rate_limits or default_limits
# Initialize rate limiter
self.rate_limiter = RateLimiter(
provider="Gemini",
requests_per_minute=limits.get("requests_per_minute", default_limits["requests_per_minute"]),
input_tokens_per_minute=limits.get("input_tokens_per_minute", default_limits["input_tokens_per_minute"]),
output_tokens_per_minute=limits.get("output_tokens_per_minute", default_limits["output_tokens_per_minute"]),
input_token_price_per_million=limits.get("input_token_price_per_million", default_limits["input_token_price_per_million"]),
output_token_price_per_million=limits.get("output_token_price_per_million", default_limits["output_token_price_per_million"])
)
def _count_tokens(self, text: str) -> int:
"""Count tokens in a string using the model's tokenizer.
Args:
text: Text to count tokens for
Returns:
Token count
"""
if not text:
return 0
try:
if self.tokenizer:
return len(self.tokenizer.encode(text))
else:
# Fallback: rough estimate if tokenizer not available
return len(text.split()) * 1.3
except Exception as e:
# Log the error but don't fail
import logging
logging.warning(f"Failed to count tokens for Gemini: {e}")
# Fallback: rough estimate if tokenizer fails
return len(text.split()) * 1.3
def _count_messages_tokens(self, messages: List[Dict[str, str]]) -> int:
"""Count tokens in all messages.
Args:
messages: List of message dictionaries
Returns:
Total token count
"""
if not messages:
return 0
total_tokens = 0
# Count tokens in each message
for message in messages:
if "content" in message and message["content"]:
total_tokens += self._count_tokens(message["content"])
# Add overhead for message formatting (estimated)
total_tokens += 4 * len(messages)
return total_tokens
def _convert_messages_to_gemini_format(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
"""Convert standard message format to Gemini-specific format.
Args:
messages: List of message dictionaries with 'role' and 'content' keys
Returns:
List of Gemini-formatted messages
"""
gemini_messages = []
# Gemini uses "user" and "model" for roles
role_mapping = {
"user": "user",
"assistant": "model",
"system": "user" # Gemini doesn't have a system role, handle specifically
}
# Check if first message is a system message
if messages and messages[0].get("role") == "system":
# For system message, we'll add it as a user message with a prefix
system_content = messages[0].get("content", "")
if system_content:
# Add the rest of the messages
for message in messages[1:]:
role = role_mapping.get(message.get("role", "user"), "user")
content = message.get("content", "")
gemini_messages.append({"role": role, "parts": content})
else:
# No system message, just convert roles
for message in messages:
role = role_mapping.get(message.get("role", "user"), "user")
content = message.get("content", "")
gemini_messages.append({"role": role, "parts": content})
return gemini_messages
def generate(
self,
messages: List[Dict[str, str]],
temperature: float = 0.7,
max_tokens: Optional[int] = None
) -> str:
"""Generate a response using Gemini API with rate limiting.
Args:
messages: List of message dictionaries
temperature: Sampling temperature
max_output_tokens: Maximum tokens to generate
Returns:
Generated response text
"""
# Count input tokens
input_tokens = self._count_messages_tokens(messages)
# Wait if we're approaching rate limits
self.rate_limiter.wait_if_needed(input_tokens, max_tokens if max_tokens else 1000)
# Format messages for Gemini API
gemini_messages = self._convert_messages_to_gemini_format(messages)
# Check if we need to start a chat or just generate
if len(gemini_messages) > 1:
# Start a chat with history
history = gemini_messages[:-1] # All but the last message
last_message = gemini_messages[-1] # The last message to send
chat = self.model.start_chat(
history=history,
)
# Send the last message to get a response
response = chat.send_message(last_message.get("parts", ""))
result_text = response.text
else:
# Single message, use generate_content
content = gemini_messages[0].get("parts", "") if gemini_messages else ""
response = self.model.generate_content(
content,
generation_config={
"temperature": temperature,
"max_tokens": max_tokens if max_tokens else None
}
)
result_text = response.text
# Estimate output tokens (Gemini API doesn't provide usage stats)
output_tokens = self._count_tokens(result_text)
# Record the request
self.rate_limiter.record_request(input_tokens, output_tokens)
return result_text
def format_message(self, role: str, content: str) -> Dict[str, str]:
"""Format message for standard API.
Args:
role: Message role (system, user, assistant)
content: Message content
Returns:
Formatted message dictionary
"""
# Standard format - conversion to Gemini format happens in generate method
return {"role": role, "content": content}
================================================
FILE: src/agent/llm/huggingface_llm.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import List, Dict, Any, Optional
from openai import OpenAI
import torch
import tiktoken
from .base import BaseLLM
class HuggingFaceLLM(BaseLLM):
"""HuggingFace model wrapper using vLLM's OpenAI-compatible API."""
def __init__(
self,
model_name: str,
api_base: str = "http://localhost:8000/v1",
api_key: str = "EMPTY",
device: str = None, # Kept for backward compatibility
torch_dtype: torch.dtype = None, # Kept for backward compatibility
max_input_tokens: int = 10000 # Maximum input tokens allowed
):
"""Initialize HuggingFace LLM via vLLM API.
Args:
model_name: Name of the model
api_base: Base URL for the vLLM API endpoint
api_key: API key (typically "EMPTY" for local vLLM deployments)
device: Ignored (handled by vLLM server)
torch_dtype: Ignored (handled by vLLM server)
max_input_tokens: Maximum number of input tokens allowed
"""
self.model_name = model_name
self.client = OpenAI(
api_key=api_key,
base_url=api_base,
)
self.max_input_tokens = max_input_tokens
# Initialize tokenizer based on model
try:
self.tokenizer = tiktoken.encoding_for_model(model_name)
except KeyError:
# Fall back to cl100k_base for unknown models (used by GPT-4, GPT-3.5-turbo)
self.tokenizer = tiktoken.get_encoding("cl100k_base")
def _count_tokens(self, messages: List[Dict[str, str]]) -> int:
"""Count the number of tokens in a list of messages.
Args:
messages: List of message dictionaries
Returns:
Total token count
"""
token_count = 0
for message in messages:
# Count tokens in content
token_count += len(self.tokenizer.encode(message["content"]))
# Add overhead for message format (role, etc.)
token_count += 4 # Approximate tokens for message formatting
# Add tokens for the formatting between messages
token_count += 2 # Final assistant message tokens
return token_count
def _truncate_messages(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
"""Truncate messages to stay within the token limit.
Args:
messages: List of message dictionaries
Returns:
Truncated list of message dictionaries
"""
if not messages:
return []
system_messages = [m for m in messages if m["role"].lower() == "system"]
non_system_messages = [m for m in messages if m["role"].lower() != "system"]
# Always keep system messages intact
result = system_messages.copy()
token_budget = self.max_input_tokens - self._count_tokens(result)
# Process non-system messages from newest to oldest
for message in reversed(non_system_messages):
message_tokens = self._count_tokens([message])
if message_tokens <= token_budget:
# We can include the entire message
result.insert(len(system_messages), message)
token_budget -= message_tokens
elif message["role"].lower() == "user" and token_budget > 20:
# For user messages, we can truncate content if needed
# Keep enough tokens for comprehension (at least some portion)
content = message["content"]
# Estimate how much content to keep
keep_ratio = token_budget / message_tokens
# Truncate from beginning to keep most recent content
if keep_ratio < 0.5:
# If we need to cut more than half, add indicator of truncation
truncated_content = f"[...truncated...] {content[int(len(content) * (1 - keep_ratio + 0.1)):].strip()}"
else:
truncated_content = content[int(len(content) * (1 - keep_ratio)):].strip()
truncated_message = {
"role": message["role"],
"content": truncated_content
}
# Verify the truncated message fits
truncated_tokens = self._count_tokens([truncated_message])
if truncated_tokens <= token_budget:
result.insert(len(system_messages), truncated_message)
token_budget -= truncated_tokens
# If we can't fit any more messages, stop
if token_budget <= 20: # Keep some buffer
break
# Ensure the messages are in the correct order (system first, then chronological)
result.sort(key=lambda m: 0 if m["role"].lower() == "system" else 1)
return result
def generate(
self,
messages: List[Dict[str, str]],
temperature: float,
max_tokens: Optional[int]
) -> str:
"""Generate a response using the vLLM API.
Args:
messages: List of message dictionaries
temperature: Sampling temperature
max_output_tokens: Maximum tokens to generate
Returns:
Generated response text
"""
max_output_tokens = max_tokens if max_tokens is not None else self.max_output_tokens
# Check token count and truncate if needed
total_tokens = self._count_tokens(messages)
if total_tokens > self.max_input_tokens:
messages = self._truncate_messages(messages)
# vLLM expects strictly alternating user/assistant roles with an optional system message at the beginning
# Prepare the messages with the proper format
formatted_messages = []
# First, check for a system message to include at the beginning
system_messages = [m for m in messages if m["role"].lower() == "system"]
if system_messages:
# Use the last system message if multiple exist
formatted_messages.append({
"role": "system",
"content": system_messages[-1]["content"]
})
# Filter out system messages and process the rest
user_assistant_messages = [m for m in messages if m["role"].lower() != "system"]
# Ensure messages alternate between user and assistant
current_role = "user" # Start with user message
for message in user_assistant_messages:
role = message["role"].lower()
# Map roles to either user or assistant
if role in ["user", "human"]:
mapped_role = "user"
else:
mapped_role = "assistant"
# If this message would create consecutive messages with the same role,
# skip adding it to avoid the alternating pattern error
if formatted_messages and mapped_role == formatted_messages[-1]["role"]:
continue
# Add the properly mapped message
formatted_messages.append({
"role": mapped_role,
"content": message["content"]
})
# Make sure the last message is from the user, so the model will respond as assistant
if not formatted_messages or formatted_messages[-1]["role"] != "user":
# If we don't have any messages or the last one isn't from user, we need to add a user message
# Use an empty message or the last assistant message as context
formatted_messages.append({
"role": "user",
"content": "Please continue." if not formatted_messages else
f"Based on your last response: '{formatted_messages[-1]['content']}', please continue."
})
# Call the API
response = self.client.chat.completions.create(
model=self.model_name,
messages=formatted_messages,
temperature=temperature,
max_tokens=max_output_tokens
)
# Extract the generated text
return response.choices[0].message.content
def format_message(self, role: str, content: str) -> Dict[str, str]:
"""Format message for OpenAI API compatible format.
Args:
role: Message role (system, user, assistant)
content: Message content
Returns:
Formatted message dictionary
"""
# Map to standard OpenAI roles if needed
if role.lower() not in ["system", "user", "assistant"]:
if role.lower() in ["human"]:
role = "user"
elif role.lower() in ["ai", "assistant"]:
role = "assistant"
else:
# Default unexpected roles to user
role = "user"
return {"role": role, "content": content}
def _messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
"""Convert messages to a single prompt string.
This method is kept for backward compatibility but is not used
in the API-based implementation.
Args:
messages: List of message dictionaries
Returns:
Formatted prompt string
"""
prompt_parts = []
for message in messages:
role = message["role"]
content = message["content"]
if role == "system":
prompt_parts.append(f"System: {content}")
elif role == "user":
prompt_parts.append(f"Human: {content}")
elif role == "assistant":
prompt_parts.append(f"Assistant: {content}")
prompt_parts.append("Assistant: ") # Add final prompt for generation
return "\n".join(prompt_parts)
================================================
FILE: src/agent/llm/openai_llm.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import List, Dict, Any, Optional
import openai
import tiktoken
from .base import BaseLLM
from .rate_limiter import RateLimiter
class OpenAILLM(BaseLLM):
"""OpenAI API wrapper."""
def __init__(
self,
api_key: str,
model: str,
rate_limits: Optional[Dict[str, Any]] = None
):
"""Initialize OpenAI LLM.
Args:
api_key: OpenAI API key
model: Model identifier (e.g., "gpt-4", "gpt-3.5-turbo")
rate_limits: Optional dictionary with rate limit settings
"""
self.client = openai.OpenAI(api_key=api_key)
self.model = model
try:
# Initialize tokenizer for the model
self.tokenizer = tiktoken.encoding_for_model(model)
except:
# Fallback to cl100k_base for new models
self.tokenizer = tiktoken.get_encoding("cl100k_base")
# Default rate limits for GPT-4o-mini
default_limits = {
"requests_per_minute": 500,
"input_tokens_per_minute": 200000,
"output_tokens_per_minute": 100000,
"input_token_price_per_million": 0.15,
"output_token_price_per_million": 0.60
}
# Use provided rate limits or defaults
limits = rate_limits or default_limits
# Initialize rate limiter
self.rate_limiter = RateLimiter(
provider="OpenAI",
requests_per_minute=limits.get("requests_per_minute", default_limits["requests_per_minute"]),
input_tokens_per_minute=limits.get("input_tokens_per_minute", default_limits["input_tokens_per_minute"]),
output_tokens_per_minute=limits.get("output_tokens_per_minute", default_limits["output_tokens_per_minute"]),
input_token_price_per_million=limits.get("input_token_price_per_million", default_limits["input_token_price_per_million"]),
output_token_price_per_million=limits.get("output_token_price_per_million", default_limits["output_token_price_per_million"])
)
def _count_tokens(self, text: str) -> int:
"""Count tokens in a string using the model's tokenizer.
Args:
text: Text to count tokens for
Returns:
Token count
"""
if not text:
return 0
try:
return len(self.tokenizer.encode(text))
except Exception as e:
# Log the error but don't fail
import logging
logging.warning(f"Failed to count tokens with OpenAI tokenizer: {e}")
# Fallback: rough estimate if tokenizer fails
return len(text.split()) * 1.3
def _count_messages_tokens(self, messages: List[Dict[str, str]]) -> int:
"""Count tokens in all messages.
Args:
messages: List of message dictionaries
Returns:
Total token count
"""
if not messages:
return 0
total_tokens = 0
# Count tokens in each message
for message in messages:
if "content" in message and message["content"]:
total_tokens += self._count_tokens(message["content"])
# Add overhead for message formatting (varies by model, but ~4 tokens per message)
total_tokens += 4 * len(messages)
# Add tokens for model overhead (varies by model)
total_tokens += 3 # Every reply is primed with <|start|>assistant<|message|>
return total_tokens
def generate(
self,
messages: List[Dict[str, str]],
temperature: float,
max_tokens: Optional[int]
) -> str:
"""Generate a response using OpenAI API with rate limiting.
Args:
messages: List of message dictionaries
temperature: Sampling temperature
max_output_tokens: Maximum tokens to generate
Returns:
Generated response text
"""
# Count input tokens
input_tokens = self._count_messages_tokens(messages)
# Wait if we're approaching rate limits (estimate output tokens as max_output_tokens)
self.rate_limiter.wait_if_needed(input_tokens, max_tokens)
# Make the API call
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens if max_tokens else None
)
result_text = response.choices[0].message.content
# Count output tokens and record request
output_tokens = response.usage.completion_tokens if hasattr(response, 'usage') else self._count_tokens(result_text)
input_tokens = response.usage.prompt_tokens if hasattr(response, 'usage') else input_tokens
self.rate_limiter.record_request(input_tokens, output_tokens)
return result_text
def format_message(self, role: str, content: str) -> Dict[str, str]:
"""Format message for OpenAI API.
Args:
role: Message role (system, user, assistant)
content: Message content
Returns:
Formatted message dictionary
"""
# OpenAI uses standard role names
return {"role": role, "content": content}
================================================
FILE: src/agent/llm/rate_limiter.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
import time
from typing import Dict, List, Optional
from collections import deque
import threading
import logging
# Configure logging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("RateLimiter")
class RateLimiter:
"""
Rate limiter for LLM API calls.
Tracks requests, input tokens, and output tokens per minute.
Also tracks cost based on token pricing.
"""
def __init__(
self,
provider: str,
requests_per_minute: int,
input_tokens_per_minute: int,
output_tokens_per_minute: int,
input_token_price_per_million: float,
output_token_price_per_million: float,
buffer_percentage: float = 0.1 # Buffer to avoid hitting exact limits
):
"""
Initialize the rate limiter.
Args:
provider: LLM provider name ("openai" or "claude")
requests_per_minute: Maximum requests per minute
input_tokens_per_minute: Maximum input tokens per minute
output_tokens_per_minute: Maximum output tokens per minute
input_token_price_per_million: Price per million input tokens
output_token_price_per_million: Price per million output tokens
buffer_percentage: Percentage buffer to avoid hitting exact limits
"""
self.provider = provider
self.requests_per_minute = requests_per_minute * (1 - buffer_percentage)
self.input_tokens_per_minute = input_tokens_per_minute * (1 - buffer_percentage)
self.output_tokens_per_minute = output_tokens_per_minute * (1 - buffer_percentage)
# Pricing
self.input_token_price = input_token_price_per_million / 1_000_000
self.output_token_price = output_token_price_per_million / 1_000_000
# Track usage within a sliding window (1 minute)
self.request_timestamps = deque()
self.input_token_usage = deque() # Tuples of (timestamp, token_count)
self.output_token_usage = deque() # Tuples of (timestamp, token_count)
# Total usage stats
self.total_requests = 0
self.total_input_tokens = 0
self.total_output_tokens = 0
self.total_cost = 0.0
# Thread lock for thread safety
self.lock = threading.Lock()
def _clean_old_entries(self, usage_queue: deque, current_time: float):
"""Remove entries older than 1 minute from the queue."""
one_minute_ago = current_time - 60
# Handle different queue formats (timestamps vs. (timestamp, value) tuples)
if usage_queue and isinstance(usage_queue[0], tuple):
# For token usage queues that store (timestamp, count) tuples
while usage_queue and usage_queue[0][0] < one_minute_ago:
usage_queue.popleft()
else:
# For request_timestamps queue that stores timestamp floats directly
while usage_queue and usage_queue[0] < one_minute_ago:
usage_queue.popleft()
def _get_usage_count(self, usage_queue: deque):
"""Get the total count from a usage queue."""
return sum(count for _, count in usage_queue)
def wait_if_needed(self, input_tokens: int, estimated_output_tokens: Optional[int] = None):
"""
Check if we're about to exceed rate limits and wait if necessary.
This improved version uses a while loop instead of recursion to
avoid potential infinite waiting scenarios.
Args:
input_tokens: Number of input tokens for the upcoming request
estimated_output_tokens: Estimated number of output tokens
"""
with self.lock:
if estimated_output_tokens is None:
estimated_output_tokens = input_tokens // 2 # Rough fallback estimate
# If this single request is bigger than the entire capacity, warn or handle
if input_tokens > self.input_tokens_per_minute or estimated_output_tokens > self.output_tokens_per_minute:
logger.warning(
f"Request uses more tokens ({input_tokens} in / {estimated_output_tokens} out) "
f"than the configured per-minute capacity. This request may never succeed."
)
while True:
current_time = time.time()
# Clean up old entries
self._clean_old_entries(self.request_timestamps, current_time)
self._clean_old_entries(self.input_token_usage, current_time)
self._clean_old_entries(self.output_token_usage, current_time)
# Calculate current usage
current_requests = len(self.request_timestamps)
current_input_tokens = self._get_usage_count(self.input_token_usage)
current_output_tokens = self._get_usage_count(self.output_token_usage)
# Check if adding this request would exceed limits
if ((current_requests + 1) <= self.requests_per_minute and
(current_input_tokens + input_tokens) <= self.input_tokens_per_minute and
(current_output_tokens + estimated_output_tokens) <= self.output_tokens_per_minute):
# We can proceed now
break
# Otherwise, compute how long to wait
wait_time = 0
if self.request_timestamps:
wait_time = max(wait_time, 60 - (current_time - self.request_timestamps[0]))
if self.input_token_usage:
wait_time = max(wait_time, 60 - (current_time - self.input_token_usage[0][0]))
if self.output_token_usage:
wait_time = max(wait_time, 60 - (current_time - self.output_token_usage[0][0]))
# If wait_time is still <= 0, we won't fix usage by waiting
if wait_time <= 0:
logger.warning(
"Waiting cannot reduce usage enough to allow this request; "
"request exceeds per-minute capacity or usage remains too high."
)
break
logger.info(f"Rate limit approaching for {self.provider}. Waiting {wait_time:.2f} seconds...")
time.sleep(wait_time)
def record_request(self, input_tokens: int, output_tokens: int):
"""
Record an API request and its token usage.
Args:
input_tokens: Number of input tokens used
output_tokens: Number of output tokens generated
"""
with self.lock:
current_time = time.time()
# Record request and token usage
self.request_timestamps.append(current_time)
self.input_token_usage.append((current_time, input_tokens))
self.output_token_usage.append((current_time, output_tokens))
# Update total stats
self.total_requests += 1
self.total_input_tokens += input_tokens
self.total_output_tokens += output_tokens
# Calculate cost
input_cost = input_tokens * self.input_token_price
output_cost = output_tokens * self.output_token_price
total_cost = input_cost + output_cost
self.total_cost += total_cost
# Log usage and cost
logger.info(
f"{self.provider} Request: {self.total_requests} | "
f"Tokens: {input_tokens}in/{output_tokens}out | "
f"Cost: ${total_cost:.6f} | "
f"Total Cost: ${self.total_cost:.6f}"
)
def print_usage_stats(self):
"""Print current usage statistics."""
with self.lock:
logger.info(f"{self.provider} Usage Statistics:")
logger.info(f" Total Requests: {self.total_requests}")
logger.info(f" Total Input Tokens: {self.total_input_tokens}")
logger.info(f" Total Output Tokens: {self.total_output_tokens}")
logger.info(f" Total Cost: ${self.total_cost:.6f}")
================================================
FILE: src/agent/orchestrator.py
================================================
# Copyright (c) Meta Platforms, Inc. and affiliates
from typing import Dict, Any, Optional, List
import time
from .base import BaseAgent
from .reader import Reader
from .searcher import Searcher
from .writer import Writer
from .verifier import Verifier
from visualizer import StatusVisualizer
import re
import yaml
import ast
import tiktoken
# Dummy visualizer class that mimics StatusVisualizer but does nothing
class DummyVisualizer:
"""A no-op visualizer that implements the same interface as StatusVisualizer but does nothing."""
def reset(self):
"""Do nothing."""
pass
def set_current_component(self, component, file_path):
"""Do nothing."""
pass
def update(self, agent_name, status):
"""Do nothing."""
pass
class Orchestrator(BaseAgent):
"""Agent responsible for managing the workflow between all other agents."""
def __init__(self, repo_path: str, config_path: Optional[str] = None, test_mode: Optional[str] = None):
"""Initialize the Orchestrator agent and its sub-agents.
Args:
repo_path: Path to the repository being analyzed
config_path: Optional path to the configuration file
test_mode: Optional test mode to run only specific components. Values: "reader_searcher", "context_print" or None
"""
super().__init__("Orchestrator")
self.repo_path = repo_path
self.context = ""
self.test_mode = test_mode
# Load configuration
self.config = {}
if config_path:
with open(config_path, 'r') as f:
self.config = yaml.safe_load(f)
# Get flow control parameters with defaults
flow_config = self.config.get('flow_control', {})
self.max_reader_search_attempts = flow_config.get('max_reader_search_attempts', 4)
self.max_verifier_rejections = flow_config.get('max_verifier_rejections', 3)
self.status_sleep_time = flow_config.get('status_sleep_time', 3)
# Check model type for context constraints
llm_config = self.config.get('llm', {})
self.model_type = llm_config.get('type', 'openai')
# Add max_input_tokens to config for context length constraint
if 'max_input_tokens' not in self.config:
self.config['max_input_tokens'] = llm_config.get('max_input_tokens', 10000)
# Initialize visualization - use dummy visualizer for "context_print" test mode
if test_mode == "context_print":
self.visualizer = DummyVisualizer()
else:
self.visualizer = StatusVisualizer()
# Initialize all sub-agents
self.reader = Reader(config_path=config_path)
self.searcher = Searcher(repo_path, config_path=config_path)
# Only initialize writer and verifier if not in reader_searcher test mode
if test_mode != "reader_searcher":
self.writer = Writer(config_path=config_path)
self.verifier = Verifier(config_path=config_path)
def _parse_verifier_response(self, response: str) -> Dict[str, Any]:
"""Parse the verifier's XML response into a structured format.
Args:
response: The XML response from the verifier
Returns:
Dictionary containing parsed verification results with structure:
{
'needs_revision': bool,
'needs_context': bool,
'suggestion': str,
'context_suggestion': str
}
"""
result = {
'needs_revision': False,
'needs_cont
gitextract_774lw2f5/
├── .gitignore
├── CHANGELOG.md
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── INSTALL.md
├── LICENSE
├── README.md
├── 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
│ ├── agent/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── llm/
│ │ │ ├── __init__.py
│ │ │ ├── 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
│ │ │ ├── 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
│ │ ├── ast_parser.py
│ │ └── topo_sort.py
│ ├── evaluate_helpfulness.py
│ ├── evaluator/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── completeness.py
│ │ ├── evaluation_common.py
│ │ ├── helper/
│ │ │ └── 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
│ │ ├── progress.py
│ │ ├── status.py
│ │ └── web_bridge.py
│ ├── web/
│ │ ├── README.md
│ │ ├── __init__.py
│ │ ├── app.py
│ │ ├── config_handler.py
│ │ ├── process_handler.py
│ │ ├── run.py
│ │ ├── static/
│ │ │ ├── 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
│ ├── app.py
│ ├── helpers.py
│ ├── requirements.txt
│ ├── start_server.sh
│ ├── static/
│ │ └── css/
│ │ └── style.css
│ ├── templates/
│ │ ├── index.html
│ │ └── results.html
│ └── test_docstring_parser.py
└── tool/
├── remove_docstrings.py
├── remove_docstrings.sh
└── serve_local_llm.sh
SYMBOL INDEX (502 symbols across 69 files)
FILE: data/raw_test_repo/example.py
function main (line 8) | def main():
FILE: data/raw_test_repo/inventory/inventory_manager.py
class Store (line 6) | class Store:
method __init__ (line 8) | def __init__(self, cap: int=20):
method put (line 13) | def put(self, obj: Item, pos: Optional[int]=None) ->bool:
method rm (line 34) | def rm(self, code: str) ->bool:
method get (line 43) | def get(self, code: str) ->Optional[Item]:
method get_at (line 46) | def get_at(self, pos: int) ->Optional[Item]:
method ls (line 52) | def ls(self) ->List[Item]:
method find (line 55) | def find(self, code: str) ->Optional[int]:
FILE: data/raw_test_repo/models/product.py
class Item (line 7) | class Item:
method check (line 56) | def check(self) -> bool:
method mod (line 79) | def mod(self, n: int=1) -> bool:
FILE: data/raw_test_repo/payment/payment_processor.py
class TxStatus (line 9) | class TxStatus(Enum):
class Tx (line 17) | class Tx:
class Handler (line 25) | class Handler(ABC):
method proc (line 28) | def proc(self, amt: Decimal) ->Tx:
method rev (line 32) | def rev(self, tx: Tx) ->bool:
class Cash (line 36) | class Cash(Handler):
method __init__ (line 38) | def __init__(self):
method add (line 41) | def add(self, amt: Decimal) ->None:
method proc (line 44) | def proc(self, amt: Decimal) ->Tx:
method rev (line 52) | def rev(self, tx: Tx) ->bool:
method ret (line 59) | def ret(self) ->Decimal:
FILE: data/raw_test_repo/vending_machine.py
class SysErr (line 9) | class SysErr(Exception):
class Sys (line 13) | class Sys:
method __init__ (line 15) | def __init__(self, h: Optional[Handler]=None):
method ls (line 20) | def ls(self) ->List[Tuple[int, Item]]:
method pick (line 28) | def pick(self, pos: int) ->Optional[Item]:
method add_money (line 36) | def add_money(self, amt: Decimal) ->None:
method buy (line 41) | def buy(self, pos: int) ->Tuple[Item, Optional[Decimal]]:
method cancel (line 55) | def cancel(self) ->Optional[Decimal]:
FILE: data/raw_test_repo_simple/helper.py
class HelperClass (line 2) | class HelperClass:
method __init__ (line 25) | def __init__(self):
method process_data (line 28) | def process_data(self):
method _internal_process (line 40) | def _internal_process(self):
method get_result (line 46) | def get_result(self):
class DataProcessor (line 52) | class DataProcessor:
method process (line 81) | def process():
method _internal_process (line 95) | def _internal_process(self):
FILE: data/raw_test_repo_simple/inner/inner_functions.py
function inner_function (line 2) | def inner_function():
function get_random_quote (line 18) | def get_random_quote():
function generate_timestamp (line 34) | def generate_timestamp():
function get_system_status (line 45) | def get_system_status():
function fetch_user_message (line 61) | def fetch_user_message():
FILE: data/raw_test_repo_simple/main.py
function main_function (line 5) | def main_function():
function utility_function (line 25) | def utility_function():
FILE: data/raw_test_repo_simple/processor.py
class AdvancedProcessor (line 6) | class AdvancedProcessor:
method __init__ (line 27) | def __init__(self):
method run (line 31) | def run(self):
method process_result (line 52) | def process_result(self):
FILE: data/raw_test_repo_simple/test_file.py
function test_function (line 2) | def test_function():
FILE: eval_completeness.py
function run_docstring_tests (line 9) | def run_docstring_tests(source_file: str) -> Dict[str, Any]:
function process_directory (line 101) | def process_directory(directory_path: str) -> Dict[str, Any]:
function print_evaluation_results (line 190) | def print_evaluation_results(results: Dict[str, Any]) -> None:
FILE: generate_docstrings.py
function generate_test_docstring (line 56) | def generate_test_docstring(component: CodeComponent) -> str:
function generate_docstring_for_component (line 117) | def generate_docstring_for_component(component: CodeComponent, orchestra...
function set_docstring_in_file (line 205) | def set_docstring_in_file(file_path: str, component: CodeComponent, docs...
function set_node_docstring (line 285) | def set_node_docstring(node: ast.AST, docstring: str):
function main (line 344) | def main():
FILE: run_web_ui.py
function check_dependencies (line 35) | def check_dependencies():
function main (line 52) | def main():
FILE: src/agent/__init__.py
function __getattr__ (line 10) | def __getattr__(name):
FILE: src/agent/base.py
class BaseAgent (line 10) | class BaseAgent(ABC):
method __init__ (line 13) | def __init__(self, name: str, config_path: Optional[str] = None):
method _initialize_llm (line 27) | def _initialize_llm(self, agent_name: str, config_path: Optional[str] ...
method add_to_memory (line 63) | def add_to_memory(self, role: str, content: str) -> None:
method refresh_memory (line 73) | def refresh_memory(self, new_memory: list[Dict[str, Any]]) -> None:
method clear_memory (line 84) | def clear_memory(self) -> None:
method memory (line 89) | def memory(self) -> list[Dict[str, Any]]:
method generate_response (line 97) | def generate_response(self, messages: Optional[List[Dict[str, Any]]] =...
method process (line 113) | def process(self, *args, **kwargs) -> Any:
FILE: src/agent/llm/base.py
class BaseLLM (line 5) | class BaseLLM(ABC):
method generate (line 9) | def generate(
method format_message (line 28) | def format_message(self, role: str, content: str) -> Dict[str, str]:
FILE: src/agent/llm/claude_llm.py
class ClaudeLLM (line 8) | class ClaudeLLM(BaseLLM):
method __init__ (line 11) | def __init__(
method _count_tokens (line 49) | def _count_tokens(self, text: str) -> int:
method _count_messages_tokens (line 76) | def _count_messages_tokens(self, messages: List[Dict[str, str]], syste...
method generate (line 125) | def generate(
method format_message (line 174) | def format_message(self, role: str, content: str) -> Dict[str, str]:
method _convert_to_claude_message (line 187) | def _convert_to_claude_message(self, message: Dict[str, str]) -> Dict[...
FILE: src/agent/llm/factory.py
class LLMFactory (line 12) | class LLMFactory:
method create_llm (line 16) | def create_llm(config: Dict[str, Any]) -> BaseLLM:
method load_config (line 80) | def load_config(config_path: Optional[str] = None) -> Dict[str, Any]:
FILE: src/agent/llm/gemini_llm.py
class GeminiLLM (line 8) | class GeminiLLM(BaseLLM):
method __init__ (line 11) | def __init__(
method _count_tokens (line 59) | def _count_tokens(self, text: str) -> int:
method _count_messages_tokens (line 84) | def _count_messages_tokens(self, messages: List[Dict[str, str]]) -> int:
method _convert_messages_to_gemini_format (line 108) | def _convert_messages_to_gemini_format(self, messages: List[Dict[str, ...
method generate (line 145) | def generate(
method format_message (line 205) | def format_message(self, role: str, content: str) -> Dict[str, str]:
FILE: src/agent/llm/huggingface_llm.py
class HuggingFaceLLM (line 8) | class HuggingFaceLLM(BaseLLM):
method __init__ (line 11) | def __init__(
method _count_tokens (line 43) | def _count_tokens(self, messages: List[Dict[str, str]]) -> int:
method _truncate_messages (line 65) | def _truncate_messages(self, messages: List[Dict[str, str]]) -> List[D...
method generate (line 125) | def generate(
method format_message (line 207) | def format_message(self, role: str, content: str) -> Dict[str, str]:
method _messages_to_prompt (line 229) | def _messages_to_prompt(self, messages: List[Dict[str, str]]) -> str:
FILE: src/agent/llm/openai_llm.py
class OpenAILLM (line 8) | class OpenAILLM(BaseLLM):
method __init__ (line 11) | def __init__(
method _count_tokens (line 56) | def _count_tokens(self, text: str) -> int:
method _count_messages_tokens (line 77) | def _count_messages_tokens(self, messages: List[Dict[str, str]]) -> int:
method generate (line 104) | def generate(
method format_message (line 144) | def format_message(self, role: str, content: str) -> Dict[str, str]:
FILE: src/agent/llm/rate_limiter.py
class RateLimiter (line 13) | class RateLimiter:
method __init__ (line 20) | def __init__(
method _clean_old_entries (line 65) | def _clean_old_entries(self, usage_queue: deque, current_time: float):
method _get_usage_count (line 79) | def _get_usage_count(self, usage_queue: deque):
method wait_if_needed (line 83) | def wait_if_needed(self, input_tokens: int, estimated_output_tokens: O...
method record_request (line 144) | def record_request(self, input_tokens: int, output_tokens: int):
method print_usage_stats (line 179) | def print_usage_stats(self):
FILE: src/agent/orchestrator.py
class DummyVisualizer (line 16) | class DummyVisualizer:
method reset (line 19) | def reset(self):
method set_current_component (line 23) | def set_current_component(self, component, file_path):
method update (line 27) | def update(self, agent_name, status):
class Orchestrator (line 31) | class Orchestrator(BaseAgent):
method __init__ (line 34) | def __init__(self, repo_path: str, config_path: Optional[str] = None, ...
method _parse_verifier_response (line 82) | def _parse_verifier_response(self, response: str) -> Dict[str, Any]:
method process (line 128) | def process(
method _update_context (line 272) | def _update_context(self, search_results: Dict[str, Any], token_consum...
method _constrain_context_length (line 367) | def _constrain_context_length(self, max_input_tokens: int = 10000, tok...
FILE: src/agent/reader.py
class CodeComponentType (line 9) | class CodeComponentType(Enum):
class InformationRequest (line 18) | class InformationRequest:
class Reader (line 25) | class Reader(BaseAgent):
method __init__ (line 28) | def __init__(self, config_path: Optional[str] = None):
method process (line 129) | def process(self, focal_component: str, context: str = "") -> str:
FILE: src/agent/searcher.py
class ParsedInfoRequest (line 14) | class ParsedInfoRequest:
class Searcher (line 34) | class Searcher(BaseAgent):
method __init__ (line 37) | def __init__(self, repo_path: str, config_path: Optional[str] = None):
method process (line 48) | def process(
method _parse_reader_response (line 103) | def _parse_reader_response(self, reader_response: str) -> ParsedInfoRe...
method _parse_comma_list (line 151) | def _parse_comma_list(self, text: str | None) -> List[str]:
method _gather_internal_info (line 164) | def _gather_internal_info(
method _gather_external_info (line 312) | def _gather_external_info(self, queries: List[str]) -> Dict[str, str]:
FILE: src/agent/tool/ast.py
class ASTUtility (line 8) | class ASTUtility(ABC):
method _get_component_name_from_code (line 12) | def _get_component_name_from_code(self, code_snippet: str) -> Optional...
method _is_code_similar (line 30) | def _is_code_similar(self, code1: str, code2: str, threshold: float = ...
function _get_component_name_from_code (line 65) | def _get_component_name_from_code(code_snippet: str) -> Optional[str]:
class ParentNodeTransformer (line 99) | class ParentNodeTransformer(ast.NodeTransformer):
method visit (line 101) | def visit(self, node):
class CallGraphBuilder (line 106) | class CallGraphBuilder(ASTUtility):
method __init__ (line 113) | def __init__(self, repo_path: str):
method _parse_file (line 127) | def _parse_file(self, file_path: str) -> ast.AST:
method _get_signature_from_code (line 148) | def _get_signature_from_code(self, code: str, is_class: bool = False) ...
method _get_node_code (line 181) | def _get_node_code(self, file_path: str, node: ast.AST) -> str:
method _is_method (line 193) | def _is_method(self, node: ast.FunctionDef) -> bool:
method _build_call_graph (line 202) | def _build_call_graph(self):
method _get_component_name_from_code (line 232) | def _get_component_name_from_code(self, code_snippet: str) -> Optional...
method get_child_function (line 243) | def get_child_function(self, code_component: str, file_path: str, chil...
method _resolve_instance_type (line 294) | def _resolve_instance_type(self, node: ast.AST, instance_name: str) ->...
method _get_class_node (line 330) | def _get_class_node(self, method_node: ast.FunctionDef) -> Optional[as...
method get_child_method (line 339) | def get_child_method(self, code_component: str, file_path: str,
method get_child_class (line 465) | def get_child_class(self, code_component: str, file_path: str, child_c...
method get_child_class_init (line 524) | def get_child_class_init(self, code_component: str, file_path: str, ch...
method _get_class_of_method (line 580) | def _get_class_of_method(self, method_node: ast.FunctionDef) -> Option...
method get_parent (line 589) | def get_parent(self, code_component: str, file_path: str, class_name: ...
class ASTNodeAnalyzer (line 707) | class ASTNodeAnalyzer:
method __init__ (line 715) | def __init__(self, repo_path: str):
method get_child_function (line 725) | def get_child_function(self, focal_node: ast.AST, file_tree: ast.AST,
method get_child_method (line 749) | def get_child_method(self, focal_node: ast.AST, file_tree: ast.AST,
method get_child_class_init (line 846) | def get_child_class_init(self, focal_node: ast.AST, file_tree: ast.AST,
method get_parent_components (line 880) | def get_parent_components(self, focal_node: ast.AST, file_tree: ast.AST,
FILE: src/agent/tool/internal_traverse.py
class ASTNodeAnalyzer (line 7) | class ASTNodeAnalyzer:
method __init__ (line 13) | def __init__(self, repo_path: str):
method get_component_by_path (line 22) | def get_component_by_path(
method _get_class_component (line 63) | def _get_class_component(self, ast_node: ast.AST, ast_tree: ast.AST, d...
method _get_function_component (line 114) | def _get_function_component(self, ast_node: ast.AST, ast_tree: ast.AST...
method _get_method_component (line 164) | def _get_method_component(self, ast_node: ast.AST, ast_tree: ast.AST, ...
method get_child_class_init (line 226) | def get_child_class_init(
method get_child_function (line 276) | def get_child_function(
method get_child_method (line 295) | def get_child_method(
method get_parent_components (line 314) | def get_parent_components(
method _find_class_init_in_node (line 377) | def _find_class_init_in_node(self, ast_node: ast.AST, class_name: str)...
method _find_function_call_in_node (line 393) | def _find_function_call_in_node(self, ast_node: ast.AST, function_name...
method _find_method_call_in_node (line 411) | def _find_method_call_in_node(
method _find_class_for_prefix (line 437) | def _find_class_for_prefix(self, ast_tree: ast.AST, prefix: Optional[s...
method _get_component_name (line 474) | def _get_component_name(self, ast_node: ast.AST) -> Optional[str]:
method _contains_call_to (line 488) | def _contains_call_to(self, ast_node: ast.AST, component_name: str) ->...
method _get_call_name (line 506) | def _get_call_name(self, call_node: ast.Call) -> Optional[str]:
method _format_call_node (line 522) | def _format_call_node(self, call_node: ast.Call) -> str:
method _get_node_source (line 535) | def _get_node_source(self, file_path: str, node: ast.AST) -> str:
method _get_end_line (line 569) | def _get_end_line(self, node: ast.AST, file_content: str) -> int:
FILE: src/agent/tool/perplexity_api.py
class PerplexityResponse (line 9) | class PerplexityResponse:
class PerplexityAPI (line 14) | class PerplexityAPI:
method __init__ (line 17) | def __init__(self, api_key: str | None = None, config_path: str = "con...
method _load_config (line 35) | def _load_config(self, config_path: str) -> Dict[str, Any]:
method query (line 45) | def query(self,
method batch_query (line 99) | def batch_query(self,
FILE: src/agent/verifier.py
class Verifier (line 7) | class Verifier(BaseAgent):
method __init__ (line 10) | def __init__(self, config_path: Optional[str] = None):
method process (line 62) | def process(
FILE: src/agent/workflow.py
function generate_docstring (line 7) | def generate_docstring(
FILE: src/agent/writer.py
class Writer (line 7) | class Writer(BaseAgent):
method __init__ (line 10) | def __init__(self, config_path: Optional[str] = None):
method is_class_component (line 119) | def is_class_component(code: str) -> bool:
method get_custom_prompt (line 130) | def get_custom_prompt(self, code: str) -> str:
method extract_docstring (line 143) | def extract_docstring(self, response: str) -> str:
method process (line 165) | def process(
FILE: src/analyze_helpfulness_significance.py
function load_results (line 21) | def load_results(filepath: str) -> Dict[str, Any]:
function get_system_scores (line 27) | def get_system_scores(results: Dict[str, Any], system: str) -> Dict[str,...
function get_paired_scores (line 49) | def get_paired_scores(results: Dict[str, Any], system1: str, system2: st...
function run_significance_tests (line 79) | def run_significance_tests(results: Dict[str, Any]) -> Dict[str, Any]:
function format_significance_markdown (line 166) | def format_significance_markdown(significance_results: Dict[str, Any]) -...
function update_markdown_report (line 193) | def update_markdown_report(stats_path: str, significance_md: str):
function main (line 205) | def main():
FILE: src/data/parse/data_process.py
function is_english (line 10) | def is_english(text):
function is_high_quality_file_docstring (line 17) | def is_high_quality_file_docstring(docstring):
function is_high_quality_class_docstring (line 27) | def is_high_quality_class_docstring(docstring):
function is_high_quality_function_docstring (line 44) | def is_high_quality_function_docstring(docstring):
function is_high_quality_docstring (line 62) | def is_high_quality_docstring(docstring, doc_type):
function get_repo_name_from_path (line 80) | def get_repo_name_from_path(path):
function extract_docstrings_from_file (line 93) | def extract_docstrings_from_file(file_path):
function add_parent_references (line 183) | def add_parent_references(tree):
function gather_python_files (line 189) | def gather_python_files(top_dir):
function process_all_repos (line 197) | def process_all_repos(top_dir, output_file):
function main (line 294) | def main():
FILE: src/data/parse/downloader.py
class GitHubRepoDownloader (line 14) | class GitHubRepoDownloader:
method __init__ (line 15) | def __init__(self, config_path: str):
method _load_config (line 23) | def _load_config(self, config_path: str) -> Dict[str, Any]:
method setup_logging (line 37) | def setup_logging(self):
method build_query (line 48) | def build_query(self) -> str:
method clone_repository (line 121) | def clone_repository(self, repo, output_dir: Path) -> bool:
method run (line 160) | def run(self):
FILE: src/data/parse/repo_tree.py
class ProjectStructureGenerator (line 9) | class ProjectStructureGenerator:
method __init__ (line 10) | def __init__(self, ignore_patterns: List[str] = None):
method should_ignore (line 17) | def should_ignore(self, path: str) -> bool:
method generate_structure (line 26) | def generate_structure(self, root_path: str, max_depth: Optional[int] ...
method format_structure (line 67) | def format_structure(self, structure: Dict, indent: int = 0) -> str:
function main (line 89) | def main():
FILE: src/dependency_analyzer/ast_parser.py
class CodeComponent (line 31) | class CodeComponent:
method to_dict (line 68) | def to_dict(self) -> Dict[str, Any]:
method from_dict (line 83) | def from_dict(data: Dict[str, Any]) -> 'CodeComponent':
class ImportCollector (line 100) | class ImportCollector(ast.NodeVisitor):
method __init__ (line 103) | def __init__(self):
method visit_Import (line 107) | def visit_Import(self, node: ast.Import):
method visit_ImportFrom (line 113) | def visit_ImportFrom(self, node: ast.ImportFrom):
class MethodDependencyCollector (line 127) | class MethodDependencyCollector(ast.NodeVisitor):
method __init__ (line 133) | def __init__(self, class_id: str, method_id: str, class_methods: Dict[...
method visit_Attribute (line 139) | def visit_Attribute(self, node: ast.Attribute):
method get_method_dependencies (line 151) | def get_method_dependencies(self) -> Set[str]:
class DependencyCollector (line 169) | class DependencyCollector(ast.NodeVisitor):
method __init__ (line 175) | def __init__(self, imports, from_imports, current_module, repo_modules):
method visit_ClassDef (line 185) | def visit_ClassDef(self, node: ast.ClassDef):
method visit_Assign (line 202) | def visit_Assign(self, node: ast.Assign):
method visit_Call (line 210) | def visit_Call(self, node: ast.Call):
method visit_Name (line 221) | def visit_Name(self, node: ast.Name):
method visit_Attribute (line 227) | def visit_Attribute(self, node: ast.Attribute):
method _process_attribute (line 232) | def _process_attribute(self, node: ast.Attribute):
method _add_dependency (line 276) | def _add_dependency(self, name):
function add_parent_to_nodes (line 305) | def add_parent_to_nodes(tree: ast.AST) -> None:
class DependencyParser (line 317) | class DependencyParser:
method __init__ (line 322) | def __init__(self, repo_path: str):
method parse_repository (line 328) | def parse_repository(self):
method _file_to_module_path (line 359) | def _file_to_module_path(self, file_path: str) -> str:
method _parse_file (line 365) | def _parse_file(self, file_path: str, relative_path: str, module_path:...
method _collect_components (line 386) | def _collect_components(self, tree: ast.AST, file_path: str, relative_...
method _resolve_dependencies (line 482) | def _resolve_dependencies(self):
method _add_class_method_dependencies (line 566) | def _add_class_method_dependencies(self):
method _get_source_segment (line 595) | def _get_source_segment(self, source: str, node: ast.AST) -> str:
method _get_docstring (line 613) | def _get_docstring(self, source: str, node: ast.AST) -> str:
method save_dependency_graph (line 631) | def save_dependency_graph(self, output_path: str):
method load_dependency_graph (line 647) | def load_dependency_graph(self, input_path: str):
FILE: src/dependency_analyzer/topo_sort.py
function detect_cycles (line 15) | def detect_cycles(graph: Dict[str, Set[str]]) -> List[List[str]]:
function resolve_cycles (line 75) | def resolve_cycles(graph: Dict[str, Set[str]]) -> Dict[str, Set[str]]:
function topological_sort (line 118) | def topological_sort(graph: Dict[str, Set[str]]) -> List[str]:
function dependency_first_dfs (line 168) | def dependency_first_dfs(graph: Dict[str, Set[str]]) -> List[str]:
function build_graph_from_components (line 236) | def build_graph_from_components(components: Dict[str, Any]) -> Dict[str,...
FILE: src/evaluate_helpfulness.py
function main (line 23) | def main():
FILE: src/evaluator/base.py
class BaseEvaluator (line 6) | class BaseEvaluator(ABC):
method __init__ (line 20) | def __init__(self, name: str, description: str):
method score (line 26) | def score(self) -> float:
method score (line 36) | def score(self, value: float) -> None:
method evaluate (line 51) | def evaluate(self, node: ast.AST) -> float:
FILE: src/evaluator/completeness.py
class CompletenessEvaluator (line 9) | class CompletenessEvaluator(BaseEvaluator):
method __init__ (line 23) | def __init__(self, name: str, description: str):
method evaluate (line 29) | def evaluate(self, node: ast.AST) -> float:
class ClassCompletenessEvaluator (line 57) | class ClassCompletenessEvaluator(CompletenessEvaluator):
method __init__ (line 89) | def __init__(self):
method evaluate_summary (line 111) | def evaluate_summary(docstring: str) -> bool:
method evaluate_description (line 125) | def evaluate_description(docstring: str) -> bool:
method evaluate_attributes (line 171) | def evaluate_attributes(docstring: str) -> bool:
method evaluate_parameters (line 188) | def evaluate_parameters(docstring: str) -> bool:
method evaluate_examples (line 205) | def evaluate_examples(docstring: str) -> bool:
method _has_attributes (line 221) | def _has_attributes(self, node: ast.ClassDef) -> bool:
method _get_required_sections (line 260) | def _get_required_sections(self, node: ast.ClassDef) -> List[str]:
method _has_init_parameters (line 285) | def _has_init_parameters(self, node: ast.ClassDef) -> bool:
method evaluate (line 301) | def evaluate(self, node: ast.ClassDef) -> float:
method evaluate_using_string (line 359) | def evaluate_using_string(self, docstring: str, element_required: Dict...
class FunctionCompletenessEvaluator (line 386) | class FunctionCompletenessEvaluator(CompletenessEvaluator):
method __init__ (line 421) | def __init__(self):
method evaluate_summary (line 443) | def evaluate_summary(docstring: str) -> bool:
method evaluate_description (line 457) | def evaluate_description(docstring: str) -> bool:
method evaluate_args (line 504) | def evaluate_args(docstring: str) -> bool:
method evaluate_returns (line 521) | def evaluate_returns(docstring: str) -> bool:
method evaluate_raises (line 538) | def evaluate_raises(docstring: str) -> bool:
method evaluate_examples (line 555) | def evaluate_examples(docstring: str) -> bool:
method evaluate (line 571) | def evaluate(self, node: ast.FunctionDef) -> float:
method evaluate_using_string (line 637) | def evaluate_using_string(self, docstring: str, element_required: Dict...
method _get_required_sections (line 664) | def _get_required_sections(self, node: ast.FunctionDef) -> List[str]:
method _has_return_statement (line 695) | def _has_return_statement(self, node: ast.FunctionDef) -> bool:
method _has_raise_statement (line 728) | def _has_raise_statement(self, node: ast.FunctionDef) -> bool:
FILE: src/evaluator/evaluation_common.py
class ScoreLevel (line 8) | class ScoreLevel(Enum):
class SummaryEvaluationExample (line 17) | class SummaryEvaluationExample:
class DescriptionEvaluationExample (line 24) | class DescriptionEvaluationExample:
class ParameterEvaluationExample (line 31) | class ParameterEvaluationExample:
FILE: src/evaluator/helper/context_finder.py
class UsageLocation (line 9) | class UsageLocation:
method __init__ (line 11) | def __init__(self, file_path: str, line_number: int, usage_type: str):
method to_dict (line 16) | def to_dict(self) -> Dict:
method from_dict (line 27) | def from_dict(cls, data: Dict) -> 'UsageLocation':
class ContextSearcher (line 31) | class ContextSearcher:
method __init__ (line 37) | def __init__(self, repo_path: str):
method _get_cache_key (line 48) | def _get_cache_key(self, file_path: str, signature: str) -> str:
method _load_from_cache (line 55) | def _load_from_cache(self, cache_key: str) -> Optional[List[UsageLocat...
method _save_to_cache (line 64) | def _save_to_cache(self, cache_key: str, locations: List[UsageLocation]):
method find_usages (line 70) | def find_usages(self, target_file: str, signature: str) -> List[UsageL...
method _parse_signature (line 130) | def _parse_signature(self, signature: str) -> Tuple[str, str]:
method _find_usages_in_file (line 152) | def _find_usages_in_file(self, content: str, file_path: str, name: str,
class ContextPreparer (line 190) | class ContextPreparer:
method __init__ (line 196) | def __init__(self, repo_path: str):
method prepare_contexts (line 206) | def prepare_contexts(self, target_file: str, signature: str) -> List[T...
method _prepare_single_context (line 229) | def _prepare_single_context(self, location: UsageLocation) -> Tuple[Op...
FILE: src/evaluator/helpfulness_attributes.py
class ScoreLevel (line 7) | class ScoreLevel(Enum):
class EvaluationExample (line 16) | class EvaluationExample:
class DocstringAttributeEvaluator (line 24) | class DocstringAttributeEvaluator:
method __init__ (line 33) | def __init__(self):
method _initialize_criteria (line 38) | def _initialize_criteria(self) -> Dict[str, Any]:
method _initialize_examples (line 86) | def _initialize_examples(self) -> List[EvaluationExample]:
method get_evaluation_prompt (line 164) | def get_evaluation_prompt(self, class_signature: str, init_function: str,
method parse_llm_response (line 255) | def parse_llm_response(self, response: str) -> Tuple[int, str]:
method get_criteria_description (line 288) | def get_criteria_description(self) -> str:
method get_score_criteria (line 292) | def get_score_criteria(self, level: ScoreLevel) -> str:
method get_examples (line 296) | def get_examples(self) -> List[EvaluationExample]:
FILE: src/evaluator/helpfulness_description.py
class DescriptionAspect (line 9) | class DescriptionAspect(Enum):
class AspectCriteria (line 17) | class AspectCriteria:
class DocstringDescriptionEvaluator (line 24) | class DocstringDescriptionEvaluator:
method __init__ (line 38) | def __init__(self):
method _initialize_criteria (line 42) | def _initialize_criteria(self) -> Dict[DescriptionAspect, AspectCriter...
method get_evaluation_prompt (line 119) | def get_evaluation_prompt(self, code_implementation: str, docstring: s...
method parse_llm_response (line 210) | def parse_llm_response(self, response: str) -> Tuple[int, str]:
method _extract_description (line 259) | def _extract_description(self, docstring: str) -> str:
FILE: src/evaluator/helpfulness_evaluator.py
class EvaluationResult (line 20) | class EvaluationResult:
class DocstringHelpfulnessEvaluator (line 28) | class DocstringHelpfulnessEvaluator:
method __init__ (line 41) | def __init__(self, data_path: str, output_dir: str, api_key: str, mode...
method sample_components (line 68) | def sample_components(self, n: int = 50, seed: int = 42) -> List[str]:
method evaluate_component (line 106) | def evaluate_component(self, component_id: str) -> List[EvaluationResu...
method run_evaluation (line 178) | def run_evaluation(self, n_samples: int = 50, seed: int = 42) -> Dict[...
method calculate_statistics (line 234) | def calculate_statistics(self, results: Dict[str, Any]) -> Dict[str, A...
method format_statistics_markdown (line 284) | def format_statistics_markdown(self, stats: Dict[str, Any]) -> str:
function main (line 326) | def main():
FILE: src/evaluator/helpfulness_evaluator_ablation.py
class EvaluationResult (line 20) | class EvaluationResult:
class DocstringHelpfulnessEvaluatorAblation (line 28) | class DocstringHelpfulnessEvaluatorAblation:
method __init__ (line 42) | def __init__(self, data_path: str, output_dir: str, api_key: str, mode...
method sample_components (line 69) | def sample_components(self, n: Optional[int] = 50, seed: int = 42) -> ...
method evaluate_component (line 112) | def evaluate_component(self, component_id: str) -> List[EvaluationResu...
method run_evaluation (line 184) | def run_evaluation(self, n_samples: int = 50, seed: int = 42) -> Dict[...
method calculate_statistics (line 240) | def calculate_statistics(self, results: Dict[str, Any]) -> Dict[str, A...
method format_statistics_markdown (line 290) | def format_statistics_markdown(self, stats: Dict[str, Any]) -> str:
function main (line 332) | def main():
FILE: src/evaluator/helpfulness_examples.py
function get_callable_name (line 10) | def get_callable_name(node: Union[ast.Name, ast.Attribute]) -> str:
class FunctionCallExample (line 26) | class FunctionCallExample:
class ClassCallExample (line 34) | class ClassCallExample:
class MethodCallExample (line 43) | class MethodCallExample:
class BaseExampleEvaluator (line 50) | class BaseExampleEvaluator(ABC):
method get_evaluation_prompt (line 59) | def get_evaluation_prompt(self, context_code: str, signature: str, exa...
method evaluate_prediction (line 74) | def evaluate_prediction(self, prediction: str, ground_truth: str) -> T...
class FunctionExampleEvaluator (line 89) | class FunctionExampleEvaluator(BaseExampleEvaluator):
method get_evaluation_prompt (line 95) | def get_evaluation_prompt(self, context_code: str, signature: str, exa...
method evaluate_prediction (line 132) | def evaluate_prediction(self, prediction: str, ground_truth: str) -> T...
method _compare_ast_nodes (line 194) | def _compare_ast_nodes(self, node1: ast.AST, node2: ast.AST) -> bool:
class ClassExampleEvaluator (line 225) | class ClassExampleEvaluator(BaseExampleEvaluator):
method get_evaluation_prompt (line 231) | def get_evaluation_prompt(self, context_code: str, signature: str, exa...
method _compare_ast_nodes (line 268) | def _compare_ast_nodes(self, node1: ast.AST, node2: ast.AST) -> bool:
method _get_func_name (line 279) | def _get_func_name(self, node: Union[ast.Name, ast.Attribute]) -> str:
method evaluate_prediction (line 293) | def evaluate_prediction(self, prediction: str, ground_truth: str) -> T...
method _compare_ast_nodes (line 352) | def _compare_ast_nodes(self, node1: ast.AST, node2: ast.AST) -> bool:
class MethodExampleEvaluator (line 357) | class MethodExampleEvaluator(BaseExampleEvaluator):
method get_evaluation_prompt (line 363) | def get_evaluation_prompt(self, context_code: str, signature: str, exa...
method evaluate_prediction (line 406) | def evaluate_prediction(self, prediction: str, ground_truth: str) -> T...
method _compare_ast_nodes (line 473) | def _compare_ast_nodes(self, node1: ast.AST, node2: ast.AST) -> bool:
FILE: src/evaluator/helpfulness_parameters.py
class DocstringParametersEvaluator (line 9) | class DocstringParametersEvaluator:
method __init__ (line 18) | def __init__(self):
method _initialize_criteria (line 23) | def _initialize_criteria(self) -> Dict[str, Any]:
method _initialize_examples (line 70) | def _initialize_examples(self) -> List[ParameterEvaluationExample]:
method get_evaluation_prompt (line 136) | def get_evaluation_prompt(self, code_component: str, docstring: str, e...
method parse_llm_response (line 232) | def parse_llm_response(self, response: str) -> Tuple[int, str]:
method get_criteria_description (line 291) | def get_criteria_description(self) -> str:
method get_score_criteria (line 295) | def get_score_criteria(self, level: ScoreLevel) -> str:
method get_examples (line 299) | def get_examples(self) -> List[ParameterEvaluationExample]:
FILE: src/evaluator/helpfulness_summary.py
class DocstringSummaryEvaluator (line 9) | class DocstringSummaryEvaluator:
method __init__ (line 18) | def __init__(self):
method _initialize_criteria (line 23) | def _initialize_criteria(self) -> Dict[str, Any]:
method _initialize_examples (line 76) | def _initialize_examples(self) -> List[SummaryEvaluationExample]:
method get_evaluation_prompt (line 159) | def get_evaluation_prompt(self, code_component: str, docstring: str, e...
method parse_llm_response (line 250) | def parse_llm_response(self, response: str) -> Tuple[int, str]:
method get_criteria_description (line 314) | def get_criteria_description(self) -> str:
method get_score_criteria (line 318) | def get_score_criteria(self, level: ScoreLevel) -> str:
method get_examples (line 322) | def get_examples(self) -> List[SummaryEvaluationExample]:
FILE: src/evaluator/segment.py
function parse_google_style_docstring (line 4) | def parse_google_style_docstring(docstring):
FILE: src/evaluator/truthfulness.py
function extract_components_from_docstring (line 29) | def extract_components_from_docstring(docstring: str) -> List[str]:
function load_dependency_graph (line 103) | def load_dependency_graph(repo_name: str) -> Dict[str, Any]:
function check_component_existence (line 121) | def check_component_existence(
function main (line 163) | def main():
function generate_summary_report (line 229) | def generate_summary_report(results: Dict[str, Dict[str, Dict[str, Any]]]):
FILE: src/visualizer/progress.py
class ProgressVisualizer (line 16) | class ProgressVisualizer:
method __init__ (line 19) | def __init__(self, components: Dict[str, any], sorted_order: List[str]):
method initialize (line 35) | def initialize(self):
method update (line 50) | def update(self, component_id: str = None, status: str = "processing"):
method finalize (line 68) | def finalize(self):
method _clear_screen (line 98) | def _clear_screen(self):
method _print_header (line 103) | def _print_header(self):
method _print_component_status (line 111) | def _print_component_status(self):
method show_dependency_stats (line 167) | def show_dependency_stats(self):
FILE: src/visualizer/status.py
class StatusVisualizer (line 8) | class StatusVisualizer:
method __init__ (line 11) | def __init__(self):
method _clear_screen (line 41) | def _clear_screen(self):
method _get_agent_color (line 46) | def _get_agent_color(self, agent: str) -> str:
method set_current_component (line 50) | def set_current_component(self, focal_component: str, file_path: str):
method _display_component_info (line 67) | def _display_component_info(self):
method update (line 73) | def update(self, active_agent: str, status_message: str = ""):
method reset (line 141) | def reset(self):
FILE: src/visualizer/web_bridge.py
class WebSocketManager (line 16) | class WebSocketManager:
method __new__ (line 23) | def __new__(cls):
method set_socket (line 29) | def set_socket(cls, socket):
method is_enabled (line 35) | def is_enabled(cls):
method emit (line 40) | def emit(cls, event, data):
method disable (line 49) | def disable(cls):
class WebStatusAdapter (line 54) | class WebStatusAdapter:
method __init__ (line 57) | def __init__(self, original_visualizer):
method set_active_agent (line 72) | def set_active_agent(self, agent_name):
method set_status_message (line 95) | def set_status_message(self, message):
method set_current_component (line 118) | def set_current_component(self, focal_component, file_path):
class WebProgressAdapter (line 146) | class WebProgressAdapter:
method __init__ (line 149) | def __init__(self, original_visualizer):
method update (line 164) | def update(self, component_id=None, status="processing"):
method mark_complete (line 201) | def mark_complete(self, component_id):
function patch_visualizers (line 224) | def patch_visualizers():
FILE: src/web/app.py
function create_app (line 25) | def create_app(debug=True):
FILE: src/web/config_handler.py
function get_default_config (line 15) | def get_default_config():
function validate_config (line 48) | def validate_config(config):
function save_config (line 72) | def save_config(config):
FILE: src/web/process_handler.py
class OutputHandler (line 26) | class OutputHandler(threading.Thread):
method __init__ (line 29) | def __init__(self, process, socketio):
method run (line 42) | def run(self):
function start_generation_process (line 206) | def start_generation_process(socketio, repo_path: str, config_path: str):
function stop_generation_process (line 294) | def stop_generation_process():
FILE: src/web/run.py
function main (line 16) | def main():
FILE: src/web/static/js/completeness.js
function updateCompletenessView (line 14) | function updateCompletenessView(completenessData) {
FILE: src/web/static/js/config.js
function loadDefaultConfig (line 12) | function loadDefaultConfig() {
function applyConfigToForm (line 31) | function applyConfigToForm(config) {
function buildConfigFromForm (line 59) | function buildConfigFromForm() {
FILE: src/web/static/js/log-handler.js
constant MAX_LOG_LINES (line 10) | const MAX_LOG_LINES = 5000;
function addLogMessage (line 18) | function addLogMessage(level, message) {
FILE: src/web/static/js/main.js
function testApiConnection (line 50) | function testApiConnection() {
function checkProcessStatus (line 117) | function checkProcessStatus() {
function setupSocketHandlers (line 151) | function setupSocketHandlers() {
function startGeneration (line 225) | function startGeneration() {
function stopGeneration (line 292) | function stopGeneration() {
function showConfigView (line 329) | function showConfigView() {
function showRunningView (line 339) | function showRunningView() {
function showMessage (line 371) | function showMessage(type, message) {
function startTimer (line 399) | function startTimer() {
function stopTimer (line 425) | function stopTimer() {
function loadCompletenessData (line 435) | function loadCompletenessData() {
FILE: src/web/static/js/repo-structure.js
function updateRepoStructure (line 27) | function updateRepoStructure(repoStructure) {
function zoomToNode (line 191) | function zoomToNode(svg, node, width, height) {
function updateFileStatus (line 224) | function updateFileStatus(file_path, status) {
FILE: src/web/static/js/status-visualizer.js
function initAgentWorkflow (line 36) | function initAgentWorkflow() {
function updateStatusVisualizer (line 179) | function updateStatusVisualizer(status) {
function updateAgentWorkflow (line 219) | function updateAgentWorkflow(activeAgent) {
function updateProgress (line 282) | function updateProgress(progress) {
FILE: src/web/visualization_handler.py
class VisualizationState (line 18) | class VisualizationState:
method __new__ (line 23) | def __new__(cls):
function get_current_status (line 48) | def get_current_status():
function update_agent_status (line 61) | def update_agent_status(active_agent: str, status_message: str):
function update_component_focus (line 72) | def update_component_focus(component_path: str, file_path: str):
function update_progress (line 84) | def update_progress(total: int, processed: int, current: str, components...
function add_log_message (line 99) | def add_log_message(message: str):
function get_repo_structure (line 111) | def get_repo_structure(repo_path: str) -> Dict[str, Any]:
function update_file_status (line 156) | def update_file_status(file_path: str, status: str):
function get_completeness_data (line 179) | def get_completeness_data(repo_path: str) -> Dict[str, Any]:
FILE: src/web_eval/app.py
function extract_component_filter (line 31) | def extract_component_filter(docstring, component):
function index (line 50) | def index():
function test_api (line 59) | def test_api():
function evaluate (line 109) | def evaluate():
function results (line 131) | def results():
function evaluate_helpfulness (line 138) | def evaluate_helpfulness():
function refresh_evaluation (line 234) | def refresh_evaluation():
function run_docstring_tests (line 252) | def run_docstring_tests(source_file: str) -> Dict[str, Any]:
function process_directory (line 353) | def process_directory(directory_path: str) -> Dict[str, Any]:
FILE: src/web_eval/helpers.py
function parse_llm_score_from_text (line 9) | def parse_llm_score_from_text(text: str) -> Tuple[int, str]:
function parse_google_style_docstring (line 58) | def parse_google_style_docstring(docstring: str) -> Dict[str, str]:
function extract_docstring_component (line 163) | def extract_docstring_component(docstring: str, component: str) -> Optio...
FILE: src/web_eval/test_docstring_parser.py
function test_and_print_result (line 11) | def test_and_print_result(test_name: str, docstring: str) -> Dict[str, A...
function test_extract_component (line 46) | def test_extract_component(docstring: str) -> None:
function main (line 72) | def main():
FILE: tool/remove_docstrings.py
class DocstringRemover (line 14) | class DocstringRemover(ast.NodeTransformer):
method visit_ClassDef (line 19) | def visit_ClassDef(self, node):
method visit_FunctionDef (line 31) | def visit_FunctionDef(self, node):
method visit_AsyncFunctionDef (line 43) | def visit_AsyncFunctionDef(self, node):
function find_python_files (line 56) | def find_python_files(directory: str) -> List[str]:
function remove_docstrings_from_file (line 68) | def remove_docstrings_from_file(file_path: str, dry_run: bool = False) -...
function main (line 105) | def main():
Condensed preview — 109 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (837K chars).
[
{
"path": ".gitignore",
"chars": 43,
"preview": "config/agent_config.yaml\ntool/add_header.sh"
},
{
"path": "CHANGELOG.md",
"chars": 76,
"preview": "0.0.1 (April 17, 2025)\n\n### First Version\n\nInclude web UI, CLI for DocAgent."
},
{
"path": "CODE_OF_CONDUCT.md",
"chars": 3536,
"preview": "# Code of Conduct\n\n## Our Pledge\n\nIn the interest of fostering an open and welcoming environment, we as\ncontributors and"
},
{
"path": "CONTRIBUTING.md",
"chars": 1979,
"preview": "# Contributing to DocAgent\nWe want to make contributing to this project as easy and transparent as\npossible.\n\n## Pull Re"
},
{
"path": "INSTALL.md",
"chars": 951,
"preview": "# Installation Guide\n\nThis guide details how to set up the environment for DocAgent.\n\n## Option 1: Installation with pip"
},
{
"path": "LICENSE",
"chars": 1087,
"preview": "MIT License\n\nCopyright (c) Meta Platforms, Inc. and affiliates.\n\nPermission is hereby granted, free of charge, to any pe"
},
{
"path": "README.md",
"chars": 7685,
"preview": "# DocAgent: Agentic Hierarchical Docstring Generation System\n\n<p align=\"center\">\n <img src=\"assets/meta_logo_white.png\""
},
{
"path": "config/example_config.yaml",
"chars": 2762,
"preview": "# Example configuration file for DocAgent\n# Copy this file to agent_config.yaml and add your own API keys\n\n# LLM configu"
},
{
"path": "data/raw_test_repo/README.md",
"chars": 3832,
"preview": "# Vending Machine Test Repository\n\nA comprehensive vending machine implementation in Python that demonstrates various pr"
},
{
"path": "data/raw_test_repo/__init__.py",
"chars": 224,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nVending Machine Package\n\nA comprehensive vending machine impleme"
},
{
"path": "data/raw_test_repo/example.py",
"chars": 1269,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom decimal import Decimal\nfrom datetime import datetime, timedelta"
},
{
"path": "data/raw_test_repo/inventory/__init__.py",
"chars": 115,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"Inventory management package for product stock tracking.\"\"\"\n"
},
{
"path": "data/raw_test_repo/inventory/inventory_manager.py",
"chars": 1707,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, List, Optional\nfrom ..models.product import"
},
{
"path": "data/raw_test_repo/models/__init__.py",
"chars": 122,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"Models package for data structures used in the vending machine.\"\""
},
{
"path": "data/raw_test_repo/models/product.py",
"chars": 5484,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom dataclasses import dataclass\nfrom typing import Optional\nfrom d"
},
{
"path": "data/raw_test_repo/payment/__init__.py",
"chars": 125,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"Payment processing package for handling different payment methods"
},
{
"path": "data/raw_test_repo/payment/payment_processor.py",
"chars": 1328,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclas"
},
{
"path": "data/raw_test_repo/vending_machine.py",
"chars": 1909,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom decimal import Decimal\nfrom typing import Optional, List, Tuple"
},
{
"path": "data/raw_test_repo_simple/helper.py",
"chars": 4327,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nclass HelperClass:\n \"\"\"\n Represents a utility for managing and"
},
{
"path": "data/raw_test_repo_simple/inner/inner_functions.py",
"chars": 3063,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\ndef inner_function():\n \"\"\"\n Returns a greeting message from an"
},
{
"path": "data/raw_test_repo_simple/main.py",
"chars": 1605,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom helper import HelperClass\nfrom inner.inner_functions import inn"
},
{
"path": "data/raw_test_repo_simple/processor.py",
"chars": 2879,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom helper import HelperClass\nfrom processor import DataProcessor\nf"
},
{
"path": "data/raw_test_repo_simple/test_file.py",
"chars": 544,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\ndef test_function():\n \"\"\"\n Returns a boolean value indicating "
},
{
"path": "eval_completeness.py",
"chars": 15394,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport ast\nimport os\nfrom typing import Dict, Any, List, Union\nfrom "
},
{
"path": "generate_docstrings.py",
"chars": 25126,
"preview": "#!/usr/bin/env python3\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nDocstring Generator with Dependency-Based"
},
{
"path": "output/dependency_graphs/raw_test_repo_dependency_graph.json",
"chars": 12003,
"preview": "{\n \"helper.HelperClass\": {\n \"id\": \"helper.HelperClass\",\n \"component_type\": \"class\",\n \"file_path\": \"/home/dayuy"
},
{
"path": "run_web_ui.py",
"chars": 3146,
"preview": "#!/usr/bin/env python3\nimport eventlet\neventlet.monkey_patch() \n# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\""
},
{
"path": "setup.py",
"chars": 3156,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom setuptools import setup, find_packages\n\n# Read the contents of "
},
{
"path": "src/DocstringGenerator.egg-info/PKG-INFO",
"chars": 28111,
"preview": "Metadata-Version: 2.2\nName: DocstringGenerator\nVersion: 0.1.0\nSummary: DocAgent for High-quality docstring generation in"
},
{
"path": "src/DocstringGenerator.egg-info/SOURCES.txt",
"chars": 1445,
"preview": "README.md\nsetup.py\nsrc/DocstringGenerator.egg-info/PKG-INFO\nsrc/DocstringGenerator.egg-info/SOURCES.txt\nsrc/DocstringGen"
},
{
"path": "src/DocstringGenerator.egg-info/dependency_links.txt",
"chars": 1,
"preview": "\n"
},
{
"path": "src/DocstringGenerator.egg-info/requires.txt",
"chars": 1191,
"preview": "numpy>=1.23.5\npyyaml>=6.0\njinja2>=3.1.5\nrequests>=2.32.0\nurllib3>=2.3.0\nastor>=0.8.1\ncode2flow>=2.5.1\npydeps>=3.0.0\nanth"
},
{
"path": "src/DocstringGenerator.egg-info/top_level.txt",
"chars": 51,
"preview": "agent\ndependency_analyzer\nevaluator\nvisualizer\nweb\n"
},
{
"path": "src/__init__.py",
"chars": 52,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n"
},
{
"path": "src/agent/README.md",
"chars": 3976,
"preview": "# Agent Framework for Docstring Generation\n\nThis directory contains the core components of the multi-agent system respon"
},
{
"path": "src/agent/__init__.py",
"chars": 587,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n# Import only essential components to avoid circular imports\nfrom .r"
},
{
"path": "src/agent/base.py",
"chars": 4262,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Dict, Op"
},
{
"path": "src/agent/llm/__init__.py",
"chars": 376,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom .base import BaseLLM\nfrom .openai_llm import OpenAILLM\nfrom .cl"
},
{
"path": "src/agent/llm/base.py",
"chars": 1126,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom abc import ABC, abstractmethod\nfrom typing import List, Dict, A"
},
{
"path": "src/agent/llm/claude_llm.py",
"chars": 7284,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import List, Dict, Any, Optional\nimport anthropic\nfrom ."
},
{
"path": "src/agent/llm/factory.py",
"chars": 3508,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, Optional\nfrom pathlib import Path\nimpo"
},
{
"path": "src/agent/llm/gemini_llm.py",
"chars": 8231,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import List, Dict, Any, Optional\nimport tiktoken\nimport "
},
{
"path": "src/agent/llm/huggingface_llm.py",
"chars": 10293,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import List, Dict, Any, Optional\nfrom openai import Open"
},
{
"path": "src/agent/llm/openai_llm.py",
"chars": 5583,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import List, Dict, Any, Optional\nimport openai\nimport ti"
},
{
"path": "src/agent/llm/rate_limiter.py",
"chars": 8478,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport time\nfrom typing import Dict, List, Optional\nfrom collections"
},
{
"path": "src/agent/orchestrator.py",
"chars": 21019,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, Optional, List\nimport time\nfrom .base "
},
{
"path": "src/agent/reader.py",
"chars": 7077,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom dataclasses import dataclass\nfrom enum import Enum\nfrom typing "
},
{
"path": "src/agent/searcher.py",
"chars": 14895,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, List, Any, Optional\nfrom .base import BaseA"
},
{
"path": "src/agent/tool/README.md",
"chars": 3211,
"preview": "# AST Call Graph Analysis Tool\n\nThis tool provides functionality to analyze Python codebases by building and querying ca"
},
{
"path": "src/agent/tool/ast.py",
"chars": 44468,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport ast\nfrom typing import Dict, List, Optional, Set, Tuple, Unio"
},
{
"path": "src/agent/tool/internal_traverse.py",
"chars": 23577,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport ast\nimport os\nfrom typing import List, Optional, Dict, Any, T"
},
{
"path": "src/agent/tool/perplexity_api.py",
"chars": 5108,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport os\nimport requests\nfrom typing import List, Dict, Any\nfrom da"
},
{
"path": "src/agent/verifier.py",
"chars": 3936,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\nfrom typing import Optional, List\nfrom .base import BaseAgent\n\n\ncla"
},
{
"path": "src/agent/workflow.py",
"chars": 2116,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Optional\nfrom pathlib import Path\nfrom .orchestra"
},
{
"path": "src/agent/writer.py",
"chars": 7906,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, Optional\nfrom abc import abstractmetho"
},
{
"path": "src/analyze_helpfulness_significance.py",
"chars": 10049,
"preview": "#!/usr/bin/env python\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nScript to analyze statistical significance"
},
{
"path": "src/data/parse/data_process.py",
"chars": 12458,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport os\nimport ast\nimport json\nfrom tqdm import tqdm\nimport argpar"
},
{
"path": "src/data/parse/downloader.py",
"chars": 12062,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport yaml\nimport os\nimport logging\nfrom github import Github\nfrom "
},
{
"path": "src/data/parse/repo_tree.py",
"chars": 4704,
"preview": "#!/usr/bin/env python3\n# Copyright (c) Meta Platforms, Inc. and affiliates\nimport os\nimport argparse\nfrom pathlib import"
},
{
"path": "src/dependency_analyzer/__init__.py",
"chars": 505,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nDependency analyzer module for building and processing import de"
},
{
"path": "src/dependency_analyzer/ast_parser.py",
"chars": 26773,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nAST-based Python code parser that extracts dependency informatio"
},
{
"path": "src/dependency_analyzer/topo_sort.py",
"chars": 9385,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nTopological sorting utilities for dependency graphs with cycle h"
},
{
"path": "src/evaluate_helpfulness.py",
"chars": 3032,
"preview": "#!/usr/bin/env python\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nScript to evaluate the helpfulness of docs"
},
{
"path": "src/evaluator/README.md",
"chars": 9121,
"preview": "# Docstring Quality Evaluator\n\nprovides a robust framework for evaluating the quality of Python docstrings. It uses stat"
},
{
"path": "src/evaluator/__init__.py",
"chars": 379,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom .base import BaseEvaluator\nfrom .completeness import ( # Remov"
},
{
"path": "src/evaluator/base.py",
"chars": 1866,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom abc import ABC, abstractmethod\nimport ast\nfrom typing import Op"
},
{
"path": "src/evaluator/completeness.py",
"chars": 27219,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport ast\nimport re\nfrom typing import Dict, List, Optional\n\nfrom e"
},
{
"path": "src/evaluator/evaluation_common.py",
"chars": 1148,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"Common utilities and classes for docstring evaluation.\"\"\"\n\nfrom t"
},
{
"path": "src/evaluator/helper/context_finder.py",
"chars": 10027,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import List, Dict, Optional, Tuple\nimport os\nimport ast\n"
},
{
"path": "src/evaluator/helpfulness_attributes.py",
"chars": 14995,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, List, Optional, Tuple\nimport re\nfrom d"
},
{
"path": "src/evaluator/helpfulness_description.py",
"chars": 12363,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, List, Tuple\nfrom dataclasses import da"
},
{
"path": "src/evaluator/helpfulness_evaluator.py",
"chars": 13029,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport json\nimport random\nimport os\nimport sys\nfrom pathlib import P"
},
{
"path": "src/evaluator/helpfulness_evaluator_ablation.py",
"chars": 13392,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport json\nimport random\nimport os\nimport sys\nfrom pathlib import P"
},
{
"path": "src/evaluator/helpfulness_examples.py",
"chars": 19789,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, List, Optional, Tuple, Union\nfrom data"
},
{
"path": "src/evaluator/helpfulness_parameters.py",
"chars": 15399,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, List, Optional, Tuple\nimport re\nfrom d"
},
{
"path": "src/evaluator/helpfulness_summary.py",
"chars": 15458,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Any, List, Optional, Tuple\nimport re\nfrom d"
},
{
"path": "src/evaluator/segment.py",
"chars": 6342,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport re\n\ndef parse_google_style_docstring(docstring):\n \"\"\"\n "
},
{
"path": "src/evaluator/truthfulness.py",
"chars": 11767,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport json\nimport os\nimport re\nimport sys\nfrom typing import List, "
},
{
"path": "src/visualizer/__init__.py",
"chars": 184,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom .status import StatusVisualizer\nfrom .progress import ProgressV"
},
{
"path": "src/visualizer/progress.py",
"chars": 7716,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nTerminal-based progress visualization for docstring generation.\n"
},
{
"path": "src/visualizer/status.py",
"chars": 5339,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nfrom typing import Dict, Set\nfrom colorama import Fore, Back, Style,"
},
{
"path": "src/visualizer/web_bridge.py",
"chars": 9687,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nWeb bridge for the docstring generation visualizers.\n\nThis modul"
},
{
"path": "src/web/README.md",
"chars": 5280,
"preview": "# DocAgent Web Interface\n\nA real-time web visualization system for the DocAgent docstring generation tool.\n\n## Overview\n"
},
{
"path": "src/web/__init__.py",
"chars": 306,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nWeb application for docstring generation visualization.\n\nThis mo"
},
{
"path": "src/web/app.py",
"chars": 8530,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nMain Flask application for the docstring generation visualizatio"
},
{
"path": "src/web/config_handler.py",
"chars": 2672,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nConfiguration handler for the docstring generation web interface"
},
{
"path": "src/web/process_handler.py",
"chars": 14457,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nProcess handler for running the docstring generation.\n\nThis modu"
},
{
"path": "src/web/run.py",
"chars": 1409,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nEntry point for running the docstring generation visualization w"
},
{
"path": "src/web/static/css/style.css",
"chars": 3379,
"preview": "/* Copyright (c) Meta Platforms, Inc. and affiliates */\n/* Main layout styles */\nbody {\n overflow-x: hidden;\n}\n\n.side"
},
{
"path": "src/web/static/js/completeness.js",
"chars": 5864,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Completeness visualization for the docstring generation web "
},
{
"path": "src/web/static/js/config.js",
"chars": 2661,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Configuration handling for the docstring generation web appl"
},
{
"path": "src/web/static/js/log-handler.js",
"chars": 1560,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Log message handler for the docstring generation web applica"
},
{
"path": "src/web/static/js/main.js",
"chars": 14528,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Main JavaScript for the docstring generation web application"
},
{
"path": "src/web/static/js/repo-structure.js",
"chars": 9828,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Repository structure visualization for the docstring generat"
},
{
"path": "src/web/static/js/status-visualizer.js",
"chars": 9920,
"preview": "// Copyright (c) Meta Platforms, Inc. and affiliates\n/**\n * Status visualizer for the docstring generation web applicati"
},
{
"path": "src/web/templates/index.html",
"chars": 11876,
"preview": "<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width"
},
{
"path": "src/web/visualization_handler.py",
"chars": 10207,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nVisualization handler for the docstring generation web interface"
},
{
"path": "src/web_eval/README.md",
"chars": 3046,
"preview": "# DocAgent - Docstring Evaluation System\n\nA web application for evaluating the quality of Python docstrings in your code"
},
{
"path": "src/web_eval/app.py",
"chars": 17863,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nimport os\nimport sys\nimport ast\nimport json\nimport argparse\nfrom fla"
},
{
"path": "src/web_eval/helpers.py",
"chars": 7592,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nHelper functions for the DocAgent web application\n\"\"\"\n\nimport re"
},
{
"path": "src/web_eval/requirements.txt",
"chars": 60,
"preview": "flask>=2.0.0\nopenai>=1.0.0\nanthropic>=0.5.0\ntabulate>=0.8.0 "
},
{
"path": "src/web_eval/start_server.sh",
"chars": 1512,
"preview": "#!/bin/bash\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\n# Default values\nHOST=\"0.0.0.0\"\nPORT=\"8080\"\nDEBUG=\"\"\n\n#"
},
{
"path": "src/web_eval/static/css/style.css",
"chars": 1561,
"preview": "/* Copyright (c) Meta Platforms, Inc. and affiliates */\n/* DocAgent - Docstring Evaluation System Styles */\n\n/* General "
},
{
"path": "src/web_eval/templates/index.html",
"chars": 10385,
"preview": "<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width"
},
{
"path": "src/web_eval/templates/results.html",
"chars": 37502,
"preview": "<!DOCTYPE html>\n<html lang=\"en\">\n<head>\n <meta charset=\"UTF-8\">\n <meta name=\"viewport\" content=\"width=device-width"
},
{
"path": "src/web_eval/test_docstring_parser.py",
"chars": 14261,
"preview": "#!/usr/bin/env python\n# Copyright (c) Meta Platforms, Inc. and affiliates\n# -*- coding: utf-8 -*-\n\"\"\"Test script for the"
},
{
"path": "tool/remove_docstrings.py",
"chars": 4082,
"preview": "#!/usr/bin/env python3\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\"\"\"\nTool to remove docstrings from Python fil"
},
{
"path": "tool/remove_docstrings.sh",
"chars": 1813,
"preview": "#!/bin/bash\n# Copyright (c) Meta Platforms, Inc. and affiliates\n\n# Shell script wrapper for the remove_docstrings.py too"
},
{
"path": "tool/serve_local_llm.sh",
"chars": 288,
"preview": "# Copyright (c) Meta Platforms, Inc. and affiliates\nCUDA_VISIBLE_DEVICES=0 python -m vllm.entrypoints.openai.api_server "
}
]
About this extraction
This page contains the full source code of the facebookresearch/DocAgent GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 109 files (781.2 KB), approximately 163.5k tokens, and a symbol index with 502 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.