Blog

Build an AI-Powered Code Documentation Bot with the Gemini API

A hands-on tutorial showing how to use the Gemini API and Python to automate the tedious task of writing code documentation.

Posted on: 2026-03-11 by AI Assistant


We’ve all been there. You write a brilliant piece of code, it works perfectly, and then comes the part nobody loves: writing the documentation. It’s tedious, time-consuming, but absolutely essential for maintainable code. What if you could automate it?

In this tutorial, you’ll learn how to build a simple but powerful “docstring bot” using Python and the Google Gemini API. This bot will read a Python file, identify functions missing docstrings, and generate them automatically.

The “Why”: The Problem with Documentation

An AI-powered bot can solve these problems by providing instant, consistently formatted documentation.

Prerequisites: What You Need

The “How”: Building the Bot

Step 1: Set Up Your Environment

First, let’s set up our Python environment and install the necessary library.

# Create and activate a virtual environment
python -m venv venv
source venv/bin/activate

# Install the Google Generative AI library
pip install google-genai

It’s crucial to protect your API key. Set it as an environment variable instead of hardcoding it.

export GOOGLE_API_KEY='YOUR_API_KEY_HERE'

Step 2: Crafting the Perfect Prompt

The magic of any LLM application is in the prompt. We need to tell the AI exactly what we want. Our prompt should instruct the model to act as a senior Python developer and generate a Google-style docstring for a given function.

PROMPT_TEMPLATE = """
As a senior Python developer, your task is to write a concise, professional, Google-style docstring for the following Python function.

Do not include the function signature in your response, only the docstring itself, enclosed in triple quotes.

Function:
```python
{function_code}

"""


#### Step 3: Writing the Python Script

Now, let's write the core logic. We'll use Python's built-in `ast` (Abstract Syntax Tree) module to safely parse the code and find functions. This is much more reliable than using regular expressions.

```python
import os
import ast
from google import genai

# Initialize the client - it automatically looks for GOOGLE_API_KEY in env
client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])

PROMPT_TEMPLATE = "Generate a Python docstring for the following function:\n{function_code}"

def get_functions_without_docstrings(filepath):
    """
    Parses a Python file and yields functions that lack a docstring.
    """
    with open(filepath, 'r') as file:
        source = file.read()
    
    tree = ast.parse(source)
    for node in ast.walk(tree):
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            # Check if the function has a docstring
            if not ast.get_docstring(node):
                print(f"Found function without docstring: {node.name}")
                yield ast.get_source_segment(source, node)

def generate_docstring(function_code):
    """
    Generates a docstring for a given function using the new GenAI SDK.
    """
    prompt = PROMPT_TEMPLATE.format(function_code=function_code)
    try:
        # Using the updated models.generate_content syntax
        response = client.models.generate_content(
            model='gemini-2.5-flash', 
            contents=prompt
        )
        return response.text.strip().replace("```", "").strip()
    except Exception as e:
        print(f"Error generating docstring: {e}")
        return None

if __name__ == "__main__":
    target_file = "sample.py"
    # Create a sample file if it doesn't exist
    if not os.path.exists(target_file):
        with open(target_file, "w") as f:
            f.write("def add(a, b):\n    return a + b\n\n")
            f.write("def subtract(a, b):\n    \"\"\"Subtracts b from a.\"\"\"\n    return a - b\n")

    print(f"Scanning {target_file} for functions without docstrings...")
    for func_code in get_functions_without_docstrings(target_file):
        print("\nOriginal Function:")
        print(func_code)
        
        print("\nGenerating Docstring...")
        docstring = generate_docstring(func_code)
        
        if docstring:
            print("\nGenerated Docstring:")
            print(docstring)
            print("-" * 20)

Step 4: Putting It All Together

Create a file named sample.py with a function that’s missing a docstring. Run the script python doc_bot.py. The bot will identify the function, send it to the Gemini API, and print the generated docstring!

What’s Next?

This is a powerful starting point. You can extend this bot to:

Automating routine tasks like documentation frees you up to focus on what really matters: solving complex problems and building great software.