Skip to main content

Build a unit test agent

In this tutorial, you will learn how to leverage aiXplain to automate the generation of Python unit tests and GitHub repository README documentation. We will build two specialized agents- a test agent for generating unit tests and a document agent for creating a project README file. Agents use Large Language Models (LLMs) to break down complex tasks, select appropriate tools, and generate actionable results. Finally, we’ll combine these agents into a team of agents to streamline software testing and documentation workflows. Let’s get started!

Step 1: Set Up File Content Reader​

We’ll create a utility function to extract the content of all .py files in a given directory.

import os
def get_files_content(code_directory):
    python_files = []
    for r, d, f in os.walk(code_directory):
        for file in f:
            if file.endswith(".py"):
                python_files.append(os.path.join(r, file))

    output = []
    for file in python_files:
        with open(file, "r", encoding="utf8") as f:
            content = f.read()
        output.append((file, content))
    return output

Step 2: Create Individual Agents​

Choose LLM for Agents (Optional)​

For agents in this tutorial, we will use the GPT-4o model. You can select an LLM of your choice from the marketplace. By default, agents use the GPT 4o model.

from aixplain.factories.agent_factory import AgentFactory
from aixplain.factories.team_agent_factory import TeamAgentFactory

llm_id = "6646261c6eb563165658bbb1"

Create a Test Agent​

This agent generates unit tests for Python files provided in the folder.

description = "You are an AI with advanced code understanding and planning capabilities. You are specialized in generating Python unit tests using pytest package. Search on the web for any information about the code that is missing. Also, you MUST mock any dependencies so it's not calling external dependencies."
test_agent = AgentFactory.create(name="TestAgent", description=description, llm_id=llm_id)

Create a Documentation Agent​

This agent generates README markdown documentation for a project repository.

description = "You are an AI with advanced documentation skills. You are specialized in generating markdown documentation given all the code repository. Search on the web for any package to get more information about it."
doc_agent = AgentFactory.create(name="DocumentationAgent", description=description, llm_id=llm_id)

Step 3: Create a Team Agent​

Next, we combine the test and documentation agents into a team of agents to help streamline the workflow.

team_agent = TeamAgentFactory.create(name="DocumentationTeamAgent", agents=[test_agent, doc_agent], llm_id=llm_id)
team_agent.__dict__
Show output

Step 4: Build a Prompt and Run the Team Agent​

Prepare the input prompt for the team agent by including the content of all Python files in a directory.

Folder Structure​

Before proceeding, ensure your project is structured as follows:

  • src/: A folder containing your Python files (.py) that you want the agent to analyze.
  • prompt.txt: A text file defining the format of input and output for the agent.

Example prompt.txt​

This file serves as a template for the input prompt. Here's a simple example:

INPUT:
File path: [Insert file path here]
Content: [Insert file content here]

OUTPUT:
[Describe the output structure here]

Save this file in the same directory as your script or notebook.

# Get file content (add your files to folder_path)
folder_path = "<YOUR_FOLDER_PATH>"
file_content = get_files_content(folder_path)
if not file_content:
  print("Folder must contain .py files")
  response = None
else:
  # Following part of the code assumes you have a text file with the prompt you want with the formattiing of the input and output
  with open("prompt.txt", "r") as f:
      prompt = f.read().strip()
  for file_path, content in file_content:
      prompt += f"\nFile path: {file_path}\nContent: {content}"
  prompt += "OUTPUT:\n"

  # Run team agent
  response = team_agent.run(prompt)
response

Step 5: Parse and Save the Output​

Once the team agent completes execution, parse and save the outputs into appropriate files.

if not os.path.exists("tests"):
  os.mkdir("tests")
output_text = response["data"]["output"]
with open(f"tests/test.md", "w") as f:
  f.write(output_text)
print(output_text)

Step 6: Deploy your Agents​

Once you are happy with your agents, you can deploy it with the code below.

team_agent.deploy()

This tutorial demonstrated how aiXplain’s platform can simplify software testing and documentation workflows by automating repetitive tasks. With minimal setup, you’ve created agents capable of writing unit tests and generating project documentation.

Start building agents today with aiXplain SDK and streamline your development process!