Code documentation

Introduction

This guide demonstrates how to create a code documentation bot using Not Diamond. This bot can automatically generate documentation for various programming languages by routing to the most appropriate Large Language Model (LLM) based on the code's complexity and language.

Why Use Not Diamond for Code Documentation?

Not Diamond offers several advantages for code documentation:

Intelligent Model Selection: Not Diamond routes requests to the best LLMs for each specific code file, ensuring accurate and relevant documentation.
Language Adaptability: Different programming languages may require varying levels of expertise. Not Diamond selects the most suitable model based on the code's language and complexity.
Cost Efficiency: By choosing the most appropriate model for each task, Not Diamond optimizes LLM usage, potentially reducing costs for large-scale documentation projects.
Scalability: As new, specialized coding models become available, Not Diamond can easily incorporate them into its routing decisions, keeping your documentation bot up-to-date.

What You'll Learn

By the end of this guide, you will be able to:

Set up a Not Diamond client for code documentation.
Create a function to read and process code files.
Implement language detection for tailored documentation generation.
Generate comprehensive code documentation using Not Diamond.
Handle different programming languages and code complexities.

Prerequisites

Python 3.10 or later
Basic knowledge of Python programming
API key for Not Diamond
Access to code files for documentation (ensure compliance with all relevant intellectual property regulations)

Step 1: Installation

First, install Not Diamond using pip:

pip install notdiamond[create]

Step 2: Setting Up the Not Diamond Client

Before using Not Diamond, you need to set up your API keys. We'll use a .env file for secure storage of your API keys.

Create a .env file in your project root:

NOTDIAMOND_API_KEY='your-notdiamond-api-key'
OPENAI_API_KEY='your-openai-api-key'
ANTHROPIC_API_KEY='your-anthropic-api-key'

Install the python-dotenv package:

pip install python-dotenv

Set up the Not Diamond client in your Python script:

import os
from dotenv import load_dotenv
from notdiamond import NotDiamond

# Load environment variables
load_dotenv()

# Initialize the Not Diamond client
client = NotDiamond()

Note: Replace the placeholder API keys in the .env file with your actual API keys. Always follow best practices for handling API keys, such as using environment variables or a secure secrets manager.

Step 3: Creating a Function to Read and Process Code Files

Let's create a function that reads a code file and prepares it for documentation:

def read_code_file(file_path):
    with open(file_path, 'r') as file:
        code_content = file.read()
    return code_content

# Example usage
code_file_path = 'path/to/your/code/file.py'
code_content = read_code_file(code_file_path)

Step 4: Implementing Language Detection

To provide more accurate documentation, let's add language detection:

import re

def detect_language(code_content):
    # This is a simple detection method. For production use, consider using a more robust library.
    if re.search(r'\bdef\b.*:|\bclass\b', code_content):
        return 'python'
    elif re.search(r'\bfunction\b.*{|\bclass\b.*{', code_content):
        return 'javascript'
    elif re.search(r'\bpublic\b.*\bclass\b|\bimport\b.*;', code_content):
        return 'java'
    else:
        return 'unknown'

Step 5: Generating Code Documentation

Now, let's create a function that generates documentation using Not Diamond:

def generate_documentation(code_content):
    language = detect_language(code_content)
    
    system_messages = {
        'python': "You are an AI assistant specialized in documenting Python code. Provide comprehensive documentation including function descriptions, parameter explanations, and usage examples.",
        'javascript': "You are an AI assistant specialized in documenting JavaScript code. Provide comprehensive documentation including function descriptions, parameter explanations, and usage examples.",
        'java': "You are an AI assistant specialized in documenting Java code. Provide comprehensive documentation including class and method descriptions, parameter explanations, and usage examples.",
        'unknown': "You are an AI assistant specialized in documenting code. Analyze the given code and provide comprehensive documentation including function/method descriptions, parameter explanations, and usage examples."
    }
    
    messages = [
        {"role": "system", "content": system_messages.get(language, system_messages['unknown'])},
        {"role": "user", "content": f"Please provide comprehensive documentation for the following {language} code:\n\n{code_content}"}
    ]
    
    result = client.chat.completions.create(
        messages=messages,
        model=['openai/gpt-4', 'anthropic/claude-3-5-sonnet-20240620']
    )
    
    return result.content, result.session_id, result.provider.model

# Example usage
code_file_path = 'example_code.py'
code_content = read_code_file(code_file_path)
documentation, session_id, model_used = generate_documentation(code_content)

print(f"Documentation:\n{documentation}")
print(f"Session ID: {session_id}")
print(f"Model used: {model_used}")

Example: Small Python File

To help you run through this cookbook, here's an example of a small Python file you can use. Save this content in a file named example_code.py:

def fibonacci(n):
    if n <= 1:
        return n
    else:
        return fibonacci(n-1) + fibonacci(n-2)

def main():
    num = 10
    print(f"The {num}th Fibonacci number is: {fibonacci(num)}")

if __name__ == "__main__":
    main()

This file contains a simple implementation of the Fibonacci sequence and a main function to demonstrate its usage.

Running the Code Documentation Bot

Now that you have all the components in place, you can run your code documentation bot. Here's a complete script that puts everything together:

import os
from dotenv import load_dotenv
from notdiamond import NotDiamond
import re

# Load environment variables
load_dotenv()

# Initialize the Not Diamond client
client = NotDiamond(api_key=os.getenv('NOTDIAMOND_API_KEY'))

# Set up API keys for other providers
os.environ['OPENAI_API_KEY'] = os.getenv('OPENAI_API_KEY')
os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY')

def read_code_file(file_path):
    with open(file_path, 'r') as file:
        code_content = file.read()
    return code_content

def detect_language(code_content):
    if re.search(r'\bdef\b.*:|\bclass\b', code_content):
        return 'python'
    elif re.search(r'\bfunction\b.*{|\bclass\b.*{', code_content):
        return 'javascript'
    elif re.search(r'\bpublic\b.*\bclass\b|\bimport\b.*;', code_content):
        return 'java'
    else:
        return 'unknown'

def generate_documentation(code_content):
    language = detect_language(code_content)
    
    system_messages = {
        'python': "You are an AI assistant specialized in documenting Python code. Provide comprehensive documentation including function descriptions, parameter explanations, and usage examples.",
        'javascript': "You are an AI assistant specialized in documenting JavaScript code. Provide comprehensive documentation including function descriptions, parameter explanations, and usage examples.",
        'java': "You are an AI assistant specialized in documenting Java code. Provide comprehensive documentation including class and method descriptions, parameter explanations, and usage examples.",
        'unknown': "You are an AI assistant specialized in documenting code. Analyze the given code and provide comprehensive documentation including function/method descriptions, parameter explanations, and usage examples."
    }
    
    messages = [
        {"role": "system", "content": system_messages.get(language, system_messages['unknown'])},
        {"role": "user", "content": f"Please provide comprehensive documentation for the following {language} code:\n\n{code_content}"}
    ]
    
    result = client.chat.completions.create(
        messages=messages,
        model=['openai/gpt-4', 'anthropic/claude-3-5-sonnet-20240620']
    )
    
    return result.content, result.session_id, result.provider.model

# Main execution
if __name__ == "__main__":
    code_file_path = 'example_code.py'
    code_content = read_code_file(code_file_path)
    documentation, session_id, model_used = generate_documentation(code_content)

    print(f"Documentation:\n{documentation}")
    print(f"Session ID: {session_id}")
    print(f"Model used: {model_used}")

To run this script:

Ensure you have created the example_code.py file with the Fibonacci sequence code provided earlier.
Save this script in the same directory as example_code.py.
Run the script using Python:

python your_script_name.py

This will generate documentation for the example Fibonacci code using Not Diamond.

Conclusion

In this guide, we've built a versatile code documentation bot using Not Diamond. This bot can:

Read and process code files from various programming languages.
Detect the programming language automatically.
Generate comprehensive documentation tailored to the specific language and complexity of the code.
Leverage Not Diamond's intelligent routing to select the best-suited model for each documentation task.

By using Not Diamond's capabilities, this documentation bot provides high-quality, context-aware documentation while optimizing for both performance and cost. As you continue to develop and refine your code documentation bot, consider expanding on this foundation to include features such as:

Integration with version control systems (e.g., Git).
Batch processing for documenting entire codebases.
Custom documentation templates for different project types or company standards.
Incremental documentation updates based on code changes.

This code documentation bot can significantly enhance your development workflow by automating the creation of accurate and comprehensive documentation, saving time and improving code maintainability across your projects.