Agent Framework for Docstring Generation

April 17, 2025 ยท View on GitHub

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.