Create an AI code generator for your design system

AI Recipes

In this series of blog articles I'm going to be creating a code generator for a design system. We will first start with a basic specification for the component in a YAML format and then expand upon it there. This recipe was briefly touched upon in my smashing magazine talk which you can find here.

You can find the code for this project here on github

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.

The schema that we're going to be using is as follows:

name: Button
description: A button component
props:
    - name: text
      type: string
      description: The text to display on the button
    - name: variant
      type: string
      description: The variant of the button
      options:
          - primary
          - secondary
          - tertiary

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 code-generator
cd code-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 the specification ready

Let's create a specification for our component. Create a new file called button.yaml in a new directory called specs and add the following code:

name: Button
description: A button component
props:
    - name: text
      type: string
      description: The text to display on the button
    - name: variant
      type: string
      description: The variant of the button
      options:
          - primary
          - secondary
          - tertiary

We're going to need one more specification though and that is for your target platform! Create a new file called platform.yaml in the specs directory and add the following code:

platform: React
dependencies:
    - react
    - typescript
    - scss
    - BEM
    - classname library
platformNotes:
    - You must use the classnames library to help with the BEM syntax
    - You must use the typescript types for the props
    - You must use scss for the styles
    - Use BEM for the class names
    - Use CSS custom properties for the design tokens
    - Use Typescript for the component

This eloquently tells the AI what platform we're targeting and what tools we're going to be using. It also tells the AI what it needs to do to make the component. This is preferable than just giving it a long list of instructions in prose.

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 specification in YAML format for a design system component.
  You are to generate a complete {{ platform }} component with styles and brief usage instructions which describes each property and example markup. You are to use the following tools: {{ platform }}. Please take a look at these specific details for that platform: {{ platformSpec }}
  You are to use the following design tokens, for example brand-blue would transform into var(--brand-blue) or platform specific equivalent
  {{ tokens }}. Use composition when it is specified as a dependency. Here are the other components you can use: {{ otherComponents }}.
  Please respond in the following JSON schema componentTemplate is optional if the platform requires it:
  { componentCode: string,
  componentStyles: string,
  componentTemplate: string,
  componentName: string,
  usageInstructions: string
  }

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 specification = fs.readFileSync('./specs/button.yaml', 'utf8')
const platformSpecification = fs.readFileSync('./specs/platform.yaml', 'utf8')
const promptTemplate = fs.readFileSync('./prompt.twig', 'utf8')

We don't really need to process the YAML at this point however we can do in the future. For now we're just going to pass the raw string to the AI.

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({
    platform: 'React',
    platformSpec: platformSpecification,
    tokens: 'brand-primary, brand-red, brand-green, brand-yellow',
})

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:

interface Component {
    code: string
    styles: string
    name: string
}

async function generateComponent(): Promise<Component> {
    const response = await openai.chat.completions.create({
        messages: [
            {
                role: 'system',
                content: systemPrompt,
            },
            {
                role: 'user',
                content: specification,
            },
        ],
        model: 'gpt-4-0125-preview',
        response_format: {
            type: 'json_object',
        },
    })
    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.

This is a good start. We have a component, styles, a name and usage instructions. We can use this to generate the component in our codebase. Let's now write the code to their respective files with another function.

Step 7: Write the component to a file

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

function writeComponentToFile(args: Component) {
    // Create the components directory if it doesn't exist
    if (!fs.existsSync('./components')) {
        fs.mkdirSync('./components')
    }
    fs.writeFileSync(`./components/${args.name}.tsx`, args.code)
    fs.writeFileSync(`./components/${args.name}.scss`, args.styles)
    console.log(`Component ${args.name} 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 component = await generateComponent()
    writeComponentToFile(component)
}
main()

Now when you run your code you should see a new components directory with a Button.tsx and Button.scss file.

Conclusion

You've now created a simple CLI that can generate components from a specification. You can now expand on this to add more features such json schema validation, more complex components and more. You can also use this as a starting point to create a more complex CLI that can generate entire projects. The possibilities are endless. Keep cookin'.

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!