Agent Code Generation Part 1: Augment Code

Leave a Comment / AI / By

Automate a Causal Inference Python Repo from Scratch with Augment Code

If you’ve ever wanted to build a clean Python repository for implementing methods from Causal Inference for the Brave and True but didn’t want to do all the setup manually, you might love using Augment Code. It is a platform that connects AI to your code repo. Since it works with your entire code base, it can see things from the helicopter perspective and perhaps better help steer a messy repo back on track. Let’s give it a spin.

In this post, I’ll walk you through how to use Augment Code to bootstrap a scientific Python project for causal inference, including:

  • Setting up a Python repo structure
  • Implementing Difference-in-Differences (DiD)
  • Adding a basic Synthetic Control implementation
  • Letting Augment do the heavy lifting on boilerplate and structure

I’ll be using VS Code for this demo.

Step 1: Think of an Idea and Create a Github Repo

I will create a new SaaS website where people can upload CSV files and analyse them using difference-in-difference and synthetic control.
The name will be “Causelle”, which is a mnemonic combination of “Causal Inference” an “Gazelle”. Imagine a scientific-looking gazelle wearing a monocle.

Let’s go to github.com and create a new empty GitHub repo:

causelle

For now, the repo contains just a README.md and a Python .gitignore file using the Github defaults. Augment Code will fill in the rest.

Step 2: Create an Account and install VS Code Extension

  1. Head to augmentcode.com, create an account and log in
  2. Note: The free community edition is sufficient for this demo.
  3. In VS Code, install the Augment extension from Augment Computing. You get a new tab in the editor.

Last step before we generate code, is to clone the empty repo to you local machine and open it up in VS Code:

  1. Clone your Github repo to your local computer
  2. Open the repo in VS Code
  3. Select the newly-added Augment tab in VS Code.
  4. Click the “Index Codebase” button. After some processing, it will display a summary message of the repo.

We are now ready to to generate the project of our dreams.

Step 3: Generate Initial Code in the Repo

We will now build out the code using a series of prompts through the Augment Tab in VS Code.

Set up a Python package for "causelle" with a `src/` folder and implement two modules:

1. `difference_in_differences.py` that provides a function to run DiD as in Causal Inference for the Brave and True.
2. `synthetic_control.py` that sets up a simple synthetic control estimator using NumPy and pandas.

Also add basic tests and example notebooks for both methods.

We will add the web application later. For now, we asked for a Python package with the main functions included.

Step 4: Review the code

Using the tree command, we can see what was added:

├── README.md
├── examples
│   ├── difference_in_differences_example.ipynb
│   └── synthetic_control_example.ipynb
├── pyproject.toml
├── src
│   ├── causelle
│   │   ├── __init__.py
│   │   ├── __pycache__
│   │   │   ├── __init__.cpython-312.pyc
│   │   │   ├── difference_in_differences.cpython-312.pyc
│   │   │   └── synthetic_control.cpython-312.pyc
│   │   ├── difference_in_differences.py
│   │   └── synthetic_control.py
│   └── causelle.egg-info
│       ├── PKG-INFO
│       ├── SOURCES.txt
│       ├── dependency_links.txt
│       ├── requires.txt
│       └── top_level.txt
└── tests
    ├── __init__.py
    ├── __pycache__
    │   └── __init__.cpython-312.pyc
    └── causelle
        ├── __init__.py
        ├── __pycache__
        │   ├── __init__.cpython-312.pyc
        │   ├── test_difference_in_differences.cpython-312-pytest-8.3.5.pyc
        │   └── test_synthetic_control.cpython-312-pytest-8.3.5.pyc
        ├── test_difference_in_differences.py
        └── test_synthetic_control.py

Augment needed a few attempts to get all the unit tests to pass, but it succeeded in the end. The good news is that it drove the entire process without any further prompting. All I had to to do was click continue a few times. A bit boring actually.

After all the unit tests passed, I ran the notebooks. One of them had an error, so I asked Augment to fix it.

Fix the error in the synthetic control notebook. SyntheticControlResult.__init__() missing 1 required positional argument: 'donor_units'

It got it done, and all the unit tests still passed. Finally, I asked it to add some docstrings and put some useful content in the README file.

Add docstrings to functions and add content to the readme.

This was done without issue.

Now, to really make it useful we need a UI. First, I created a new branch, for the UI:

git checkout -b saas-website

I chose to not be too specific about the tech-stack for the frontend, just to see what it would come up with.

Create a frontend that can be used to call the functions in the browser.

It opted for a Flask backend and a Bootstrap frontend. Fine with me. It asked for my permission to create a directory upload and then it was done. It worked like a charm.

The web application now looks like this.

Take my word for it. It works. You can upload CSV files or use example files provided, which I didn’t even ask for, it will analyse them and show the appropriate visualisations. It can be improved, but is pretty neat already.

Step 5: Deploy

Let’s refactor a bit and create deployment configuration for Digital Ocean. You can of course choose a different cloud provider. I entered the following prompt:


Move the UI code into a frontend folder, to separate it from the rest. Also add configuration that makes it deploy to digital ocean app platform.

After this step, I could begin to deploy on Digital Ocean’s App Platform. First, I pushed all the changes to Github. Then I added the new repo in Digital Ocean’s interface. After this point, the code will deploy with each change to the main branch.

There was a small error that Augment need to fix.


When I deploy the app to digital ocean, I get the following error during deployment.:
[2025-04-25 15:53:48] [2025-04-25 15:53:48 +0000] [2] [ERROR] Exception in worker process
[2025-04-25 15:53:48] Traceback (most recent call last):
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/arbiter.py", line 608, in spawn_worker
[2025-04-25 15:53:48]     worker.init_process()
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/workers/base.py", line 135, in init_process
[2025-04-25 15:53:48]     self.load_wsgi()
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/workers/base.py", line 147, in load_wsgi
[2025-04-25 15:53:48]     self.wsgi = self.app.wsgi()
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/app/base.py", line 66, in wsgi
[2025-04-25 15:53:48]     self.callable = self.load()
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/app/wsgiapp.py", line 57, in load
[2025-04-25 15:53:48]     return self.load_wsgiapp()
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/app/wsgiapp.py", line 47, in load_wsgiapp
[2025-04-25 15:53:48]     return util.import_app(self.app_uri)
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/site-packages/gunicorn/util.py", line 370, in import_app
[2025-04-25 15:53:48]     mod = importlib.import_module(module)
[2025-04-25 15:53:48]   File "/usr/local/lib/python3.9/importlib/__init__.py", line 127, in import_module
[2025-04-25 15:53:48]     return _bootstrap._gcd_import(name[level:], package, level)
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap_external>", line 850, in exec_module
[2025-04-25 15:53:48]   File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
[2025-04-25 15:53:48]   File "/app/app.py", line 8, in <module>
[2025-04-25 15:53:48]     from causelle.synthetic_control import synthetic_control
[2025-04-25 15:53:48] ModuleNotFoundError: No module named 'causelle'
Fix it.

Augment fixed the error in the first attempt and the web application came online.

Summary

Clearly, the users task is to read and review the generated code. It’s a good idea to generate code on a feature branch, create a PR, merge it into main if OK, and start a new feature branch.

In this article, we have seen that with Augment Code, you can:
– Set up clean scientific codebases from a simple prompt
– Avoid boilerplate and focus on features
– Auto-generate PRs that are easy to review and tweak

Next, you could add more models from Causal Inference for the Brave and True, like matching or instrumental variables. Or maybe create a SaaS frontend?

Leave a Comment

This site uses Akismet to reduce spam. Learn how your comment data is processed.