Document your codebase with AI

In this series of blog articles I'm going to be creating a documentation generator for a design system. We will first start with a basic component that we'll write ourselves here.

This recipe was briefly touched upon in my smashing magazine talk which you can find here.

This is the second in a series of articles. The first article can be found here.

Some of the tasks are the same as they were in the first recipe but like a recipe book, I'm going to repeat them here for completeness.

For this recipe you will need:

• Node.js version 18 or higher
• An OpenAI API key

It should run on Windows, Mac or Linux however I've only tested it on Mac and Linux.

For this project I'm going to be using NodeJS. However it can all be done with python or even PHP if you're feeling brave.

You're also going to need an OpenAI API key. You can get one by signing up here.

Step 1: Create a new project

Firstly, create a new directory for your project and then run npm init to create a new package.json file.

mkdir docs-generator
cd docs-generator
npm init -y

I've skipped the questions by using the -y flag. If you want to answer the questions then you can leave it off.

Step 2: Set up your project

We're going to need a few packages to get started. Run the following command to install them:

npm install --save-dev openai ts-node @types/node dotenv twig @types/twig

Inside package.json let's add a start function that will run ts-node on index.ts

    "scripts": {
        "start": "ts-node index.ts"
    }

Step 3: Create a new file

Create a new file called index.ts and add the following code:

import { OpenAI } from 'openai'
import dotenv from 'dotenv'
import fs from 'fs'

// Configure dotenv to get our API key from the .env file which we'll create in a moment
dotenv.config()

Ok next we will need to create a .env file in the root of our project. This will contain our OpenAI API key. Add the following to the file:

OPENAI_API_KEY=your-api-key-here

The reason why we're using dotenv is so that we can keep our API key out of our code and also out of our git repository.

Step 4: Get your component code prepared

For the sake of brevitry let's quickly knock up a Lit button component that has a couple of variants and can change to an anchor tag inside it if it has an href:

import { html, css, LitElement } from 'lit'
import { customElement, property } from 'lit/decorators.js'

@customElement('my-button')
export class Button extends LitElement {
    static styles = css`
        button {
            padding: 10px 20px;
            border: none;
            border-radius: 5px;
            cursor: pointer;
        }
        .primary {
            background-color: var(--brand-primary);
            color: white;
        }
        .secondary {
            background-color: var(--brand-secondary);
            color: white;
        }
    `

    @property({ type: String })
    href?: string

    @property({ type: String })
    variant: 'primary' | 'secondary' = 'primary'

    render() {
        if (this.href) {
            return html`<a href="${this.href}"><slot></slot></a>`
        }
        return html`<button><slot></slot></button>`
    }
}

Save this file as button.ts in the root of your project.

Step 5: Set up your prompt

Inside our index.ts file we need to create a new instance of the OpenAI class. Add the following code:

const openai = new OpenAI()

Next, let's create a new twig file with our system prompt in it. I prefer to keep my prompts in a separate file so that I can easily change them without having to change my code. Create a new file called prompt.twig and add the following code:

You are to be given a lit component  for a design system component.
You are to write documentation for the component in the following format:

Description: A description of the component
Usage: How to use the component
Properties: A list of properties
Events: A list of events

What is a system prompt?
A system prompt is like the first inpression of the AI. It's the first thing that the AI sees and it's what it uses to generate the rest of the conversation. It's important to get this right as it can affect the quality of the response.

Save this file as system_prompt.twig in the root of your project.

Next lets load the specifications and the prompt into our code. Add the following code to your index.ts file:

const code = fs.readFileSync('button.ts', 'utf-8')

Let's now process the twig with the specifications. Add the following code to your index.ts file:

At the top of the file add the following import:

import Twig from 'twig'

Then add the following code to the bottom of the file:

const twig = Twig.twig
const template = twig({ data: promptTemplate })
const systemPrompt = template.render()

This will render out a system prompt that we can use to talk to the AI.

Step 6: Talk to the AI

Now that we have our system prompt we can talk to the AI. First lets define an interface for the component that we're going to get back from the AI. Add the following code to the bottom of your index.ts file:

async function generateDocs() {
    const response = await openai.chat.completions.create({
        messages: [
            {
                role: 'system',
                content: systemPrompt,
            },
            {
                role: 'user',
                content: code,
            },
        ],
        model: 'gpt-4-0125-preview',
    })
    const component = JSON.parse(response.choices[0].message.content || '')
    return component
}

This will send the system prompt and the specification to the AI and then return the component that it generates. We're using the gpt-4-0125-preview model which is the latest model at the time of writing. This model is in preview so it may not be available to everyone.

Step 7: Write the documentation to a file

We need to write the component to a file. Add the following code to your index.ts file:

function writeDocsToFile(docs: string) {
    // Create the components directory if it doesn't exist
    fs.writeFileSync(`button.md`, docs)
    console.log(`Docs has been written to file`)
}

Let's put this all together now in a main function. Add the following code to your index.ts file:

async function main() {
    const docs = await generateDocs()
    writeDocsToFile(component)
}
main()

Conclusion

You've now created a simple CLI that can generate documentation from an existing component. You can now expand on this to add more features such as:

• Generating it from a schema
• Reading in the storybook stories and generating the documentation from that
• Generating the documentation for the entire design system
• And more!

Tweak the system prompt to make it more dynamic so you can generate documentation for different types of components.

Interested in more?

Does your organization need help integrating AI into your design system operation? We can help! Big Medium helps complex organizations adopt new technologies and workflows while also building and shipping great digital products. Feel free to get in touch!