How I Automated API Documentation with AI – Lessons Learned

I inherited a legacy Django codebase a few months ago. The code worked, but there was zero documentation – no docstrings, no README, nothing. My first instinct was to manually write it, but after 3 days and only covering 10% of the routes, I knew there had to be a better way.

So I started looking into using LLMs to generate documentation automatically. What I found surprised me: it's not as straightforward as “throw the code at GPT and get docs back.” Here's what I tried, what failed, and what eventually worked.

The Problem: Legacy Code, No Comments

The project had about 50 API endpoints scattered across several files. Business logic was mixed with views, and some functions were 500 lines long. I needed docstrings for every function, plus a high-level overview of each endpoint.

I could have used tools like Doxygen or Sphinx, but they only extract existing comments – they don't write new ones. So I turned to LLMs.

What I Tried First (and Why It Didn't Work)

Attempt 1: Copy-paste to ChatGPT

I literally copied a 200-line view function into ChatGPT with the prompt “Write a docstring for this function.” It worked for the first one, but then:

• Token limits forced me to split code manually.
• The output was inconsistent – sometimes it guessed parameters wrong.
• Doing this for 50 endpoints was too tedious.

Attempt 2: A simple Python script using OpenAI's API

I wrote a quick script that read each file, sent chunks to GPT-3.5, and inserted docstrings. But I hit rate limits, got messy formatting, and the LLM often hallucinated parameter types.

import openai

openai.api_key = "sk-..."

def generate_docstring(code):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system", "content": "You are a Python expert. Write a clear docstring."},
            {"role": "user", "content": f"```
{% endraw %}
python\n{code}\n
{% raw %}
```"}
        ]
    )
    return response.choices[0].message.content

This worked for small snippets, but for large files I had to split by function, losing context. Plus, the cost added up.

Attempt 3: Using a hosted AI documentation service

I looked into a few hosted solutions (including the one at https://ai.interwestinfo.com/). Some were promising, but they required uploading my entire codebase to their servers, which wasn't an option for this client project due to data privacy. I needed something local.

What Eventually Worked: A RAG-Inspired Pipeline

I realized the real problem was context. An LLM can't read your entire codebase at once. I needed a retrieval-augmented generation (RAG) approach: index the code, then for each function, retrieve related code (imports, other functions it calls) and feed that into the LLM along with the function itself.

I used LangChain to orchestrate this. Here's the simplified pipeline:

1. Load and chunk the codebase (one chunk per function).
2. Embed each chunk using a local embedding model (sentence-transformers).
3. Store embeddings in a vector store (Chroma).
4. For each target function, query the vector store for the top 3 similar code chunks (e.g., functions it imports or depends on).
5. Build a prompt containing the target function + context + instructions.
6. Send to GPT-4 (I found GPT-4 much better than 3.5 for code understanding).
7. Parse the response and write the docstring back into the file.

Here's the core code (simplified for clarity):

import os
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA

# 1. Load all Python files
from pathlib import Path

code_dir = Path("./myproject")
all_docs = []
for py_file in code_dir.rglob("*.py"):
    loader = TextLoader(str(py_file))
    docs = loader.load()
    all_docs.extend(docs)

# 2. Split into function-level chunks (custom splitter)
# In reality, I used a function that uses regex to split by `def `
splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
chunks = splitter.split_documents(all_docs)

# 3. Embed and store
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectordb = Chroma.from_documents(chunks, embeddings, persist_directory="./chroma_db")

# 4. Set up QA chain
llm = OpenAI(model="gpt-4", temperature=0)
qa = RetrievalQA.from_chain_type(llm, retriever=vectordb.as_retriever(k=3))

# 5. For each function, query and get docstring
def generate_docstring_for(code_snippet):
    prompt = f"""You are writing docstrings for a Python codebase. 
Given the function below and some context from the same codebase (imports, other functions it uses), 
write a detailed docstring explaining what the function does, its parameters, and return value.

Context:
{{context}}

Code:
{code_snippet}

Docstring:"""
    return qa.run(prompt)

# Iterate over all functions...

This was bliss. The LLM had enough context to generate accurate docstrings. I tweaked the prompts to match the team's style (Google-style docstrings). The whole thing took about 2 hours to run for 200 functions, costing ~$15 in GPT-4 API fees.

Lessons Learned & Trade-offs

• GPT-4 is worth the cost for code tasks. GPT-3.5 often missed context or invented parameters.
• Local embeddings saved money on the retrieval step. I used all-MiniLM-L6-v2 – cheap and fast.
• Context window matters: For very long functions, I had to truncate context. In those cases, I fell back to a smaller prompt.
• Not all languages work well: Python was fine; I tried it on a Ruby project and results were worse. Your mileage may vary.
• You still need human review: The generated docs were about 80% accurate. I had to fix parameter names and edge cases.
• Data privacy: Running locally (or with a private LLM) avoids sending sensitive code to third parties. That's why I didn't use the hosted service I mentioned earlier, though it might work for non-sensitive projects.

What I'd Do Differently Next Time

• I'd use a code-specific embedding model like codebert or graphcodebert instead of a general one. That might improve similarity matches.
• I'd add a validation step: after generating a docstring, run a simple regex to check that all parameters mentioned actually exist in the function signature. That caught hallucinations.
• I'd consider fine-tuning a small model for the docstring task if I had enough examples. But for a one-off project, GPT-4 was fine.

The Bigger Picture

Automating documentation with AI is absolutely doable, but it's not a magic button. You need to think about context, cost, and privacy. The approach I described – RAG with an LLM – can be adapted to many code understanding tasks: generating tests, summarizing changes, or even explaining complex logic.

If you're considering a similar project, start small. Pick one module, prototype the pipeline, and check if the output quality meets your bar. Then scale up.

What's your experience been with AI-generated code documentation? Have you found a better way to handle the context problem? I'd love to hear what worked – or failed – for you.