Skip to main content

A module for generating AI-based code documentation and data flow diagrams.

Project description

FastWrite

Python Module for AI-Assisted Documentation

Current Statistics:

Overview

This module provides functionality to:

  • Process Code Files: Extract and list Python files from a ZIP archive.
  • Generate Data Flow Diagrams: Create a data flow chart (in Graphviz format) by analyzing Python code using the AST module.
  • Generate Documentation: Produce detailed documentation for Python code using multiple AI models:
    • Groq-based models (remote)
    • Gemini-based models (remote)
    • OpenAI-based models (remote)
    • OpenRouter-based models (remote)
    • Ollama-based models (local)
  • Evaluate Documentation Quality: Compute BLEU/ROUGE scores to compare generated documentation against a reference document.

Installation

Requirements

Install Dependencies

pip install groq google-genai requests nltk python-dotenv openai rouge_metric

Usage

Processing Files:

from FastWrite import extract_zip, list_python_files, read_file
import tempfile
import os

# Specify the path to your ZIP file containing Python code
zip_file_path = "path/to/your/code.zip"

with tempfile.TemporaryDirectory() as tmp_dir:
    # Extract the ZIP file
    extract_zip(zip_file_path, tmp_dir)
    
    # List Python files in the extracted directory
    py_files = list_python_files(tmp_dir)
    
    if py_files:
        # For example, choose the first Python file as the main file
        main_file_path = os.path.join(tmp_dir, py_files[0])
        code_content = read_file(main_file_path)

Generating Data Flow Diagrams:

from FastWrite import generate_data_flow

# Generate Graphviz code for the data flow diagram
graphviz_code = generate_data_flow(code_content)
print(graphviz_code)

Generating Documentation (Express Mode):

# Normal Mode:
py -m FastWrite --LLM_NAME

# With custom code file
py -m FastWrite code_filename.py --LLM_NAME

# Simplified (Novice Friendly) Mode:
py -m FastWrite code_filename.py --LLM_NAME --Simplify

# Formal (Precise and to the Point) Mode:
py -m FastWrite code_filename.py --LLM_NAME --Formal

# Research (In-Depth and Elaborate Documentation):
py -m FastWrite code_filename.py --LLM_NAME --Research

# Custom (Allows you to feed a custom prompt of your own):
py -m FastWrite code_filename.py --LLM_NAME --Custom-Prompt "Enter Custom Prompt here"

Generating Documentation (Groq):

from FastWrite import generate_documentation_groq

custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""

groq_api_key = "your_groq_api_key"
groq_model = "deepseek-r1-distill-llama-70b"  # Replace with your desired model

doc_groq = generate_documentation_groq(code_content, custom_prompt, groq_api_key, groq_model)
print(doc_groq)

Generating Documentation (Gemini):

from FastWrite import generate_documentation_gemini

custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""

gemini_api_key = "your_gemini_api_key"
gemini_model = "gemini-2.0-flash"  # Replace with your desired model

doc_gemini = generate_documentation_gemini(code_content, custom_prompt, gemini_api_key, gemini_model)
print(doc_gemini)

Generating Documentation (OpenAI):

from FastWrite import generate_documentation_openai

custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
doc_openai = generate_documentation_openai(code_content, custom_prompt)
print(doc_openai)

Generating Documentation (Ollama):

from FastWrite import generate_documentation_ollama

custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""

# Replace with your local Ollama model name (e.g., "ollama-llama-70b")
ollama_model = "ollama-llama-70b"

doc_ollama = generate_documentation_ollama(code_content, custom_prompt, ollama_model)
print(doc_ollama)

Generating Documentation (OpenRouter):

from FastWrite import generate_documentation_openrouter

custom_prompt = """
Objective:
Generate high-quality, developer-friendly documentation for the following Python code Ensure you include Detailed function-level and file-level documentation and a high level slightly less technical documentation at the start to make it friendly. Do not print full code snippets of existing code, just explain them:
"""
doc_openrouter = generate_documentation_openrouter(code_content, custom_prompt)
print(doc_openrouter)

Calculating Bleu Score:

from FastWrite import calculate_bleu

# Provide a reference documentation string for comparison
reference_doc = "Your reference documentation text here..."

# Basic BLEU score calculation
bleu_score = calculate_bleu(doc_llm_host, reference_doc) ##LLM host may include Groq,Gemini,OpenAI or Ollama
print("BLEU Score:", bleu_score)

# BLEU score with smoothing
# Available smoothing methods: 'method0', 'method1', 'method2', 'method3', 'method4', 'method5', 'method6', 'method7'
bleu_score_smoothed = calculate_bleu(doc_llm_host, reference_doc, smoothing_method='method1')
print("BLEU Score with Smoothing:", bleu_score_smoothed)

# Calculate BLEU score with multiple reference documents
from FastWrite import calculate_bleu_multi_reference
reference_docs = ["Reference doc 1...", "Reference doc 2..."]
bleu_score_multi = calculate_bleu_multi_reference(doc_llm_host, reference_docs, smoothing_method='method1')
print("BLEU Score with Multiple References:", bleu_score_multi)

# Calculate BLEU score directly from files
from FastWrite import calculate_bleu_from_files
bleu_score_from_files = calculate_bleu_from_files("candidate_doc.txt", "reference_doc.txt", smoothing_method='method1')
print("BLEU Score from Files:", bleu_score_from_files)

Calculating Rouge Score:

from FastWrite.rouge import calculate_rouge

# Provide a reference documentation string for comparison
reference_doc = "Your reference documentation text here..."

# Calculate Rouge
rouge = calculate_rouge(doc_llm_host,reference_doc)

Generating README File:

from FastWrite.print import readmegen

readmegen(doc_llm,llm_used)

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

fastwrite-1.2.1.tar.gz (11.1 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

fastwrite-1.2.1-py3-none-any.whl (13.9 kB view details)

Uploaded Python 3

File details

Details for the file fastwrite-1.2.1.tar.gz.

File metadata

  • Download URL: fastwrite-1.2.1.tar.gz
  • Upload date:
  • Size: 11.1 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.13.3

File hashes

Hashes for fastwrite-1.2.1.tar.gz
Algorithm Hash digest
SHA256 77dfcae96904afc921fba77f4dab80671981ff0a18b98324ce127c6119e48f7f
MD5 a5a5b7f86ce26706fc3058446a2c7998
BLAKE2b-256 5abdc3639d5efa7590eac61215c655b14d26e8a10d8b2a05e228f291ccd8386f

See more details on using hashes here.

File details

Details for the file fastwrite-1.2.1-py3-none-any.whl.

File metadata

  • Download URL: fastwrite-1.2.1-py3-none-any.whl
  • Upload date:
  • Size: 13.9 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.13.3

File hashes

Hashes for fastwrite-1.2.1-py3-none-any.whl
Algorithm Hash digest
SHA256 772ca2e820b8235e94352269a4f97cbb848e8c081d426b36fe630c99b52d2bc9
MD5 d8c6bb31d3e92cd34d446c7cf6af1c97
BLAKE2b-256 f85a58bc50fdd6496c4cc7e97e26518f0bb1bd2227f28d00f7e0fc5f7c47e019

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page