Build an AI Agent — Chapter 2: Functions

1️⃣ Calculator

So, we know we’re building an AI Agent — but here’s the catch: every agent needs a project to work on.

To make things concrete (and fun), I’ve created a simple command-line calculator app. This will be our test project — something the AI can read, update, and even run as we go along. Think of it as the playground where our agent starts learning how to act like… well, an agent.

Assignment

Create a new directory called calculator in the root of your project.
Copy and paste the main.py and tests.py files from below into the calculator directory.

# main.py
import sys
from pkg.calculator import Calculator
from pkg.render import render

def main():
    calculator = Calculator()
    if len(sys.argv) <= 1:
        print("Calculator App")
        print('Usage: python main.py "<expression>"')
        print('Example: python main.py "3 + 5"')
        return
    expression = " ".join(sys.argv[1:])
    try:
        result = calculator.evaluate(expression)
        to_print = render(expression, result)
        print(to_print)
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    main()

# tests.py
import unittest
from pkg.calculator import Calculator

class TestCalculator(unittest.TestCase):
    def setUp(self):
        self.calculator = Calculator()
    def test_addition(self):
        result = self.calculator.evaluate("3 + 5")
        self.assertEqual(result, 8)
    def test_subtraction(self):
        result = self.calculator.evaluate("10 - 4")
        self.assertEqual(result, 6)
    def test_multiplication(self):
        result = self.calculator.evaluate("3 * 4")
        self.assertEqual(result, 12)
    def test_division(self):
        result = self.calculator.evaluate("10 / 2")
        self.assertEqual(result, 5)
    def test_nested_expression(self):
        result = self.calculator.evaluate("3 * 4 + 5")
        self.assertEqual(result, 17)
    def test_complex_expression(self):
        result = self.calculator.evaluate("2 * 3 - 8 / 2 + 5")
        self.assertEqual(result, 7)
    def test_empty_expression(self):
        result = self.calculator.evaluate("")
        self.assertIsNone(result)
    def test_invalid_operator(self):
        with self.assertRaises(ValueError):
            self.calculator.evaluate("$ 3 5")
    def test_not_enough_operands(self):
        with self.assertRaises(ValueError):
            self.calculator.evaluate("+ 3")

if __name__ == "__main__":
    unittest.main()

Create a new directory in calculator called pkg.
Copy and paste the calculator.py and render.py files from below into the pkg directory.

# calculator.py
class Calculator:
    def __init__(self):
        self.operators = {
            "+": lambda a, b: a + b,
            "-": lambda a, b: a - b,
            "*": lambda a, b: a * b,
            "/": lambda a, b: a / b,
        }
        self.precedence = {
            "+": 1,
            "-": 1,
            "*": 2,
            "/": 2,
        }
    def evaluate(self, expression):
        if not expression or expression.isspace():
            return None
        tokens = expression.strip().split()
        return self._evaluate_infix(tokens)
    def _evaluate_infix(self, tokens):
        values = []
        operators = []
        for token in tokens:
            if token in self.operators:
                while (
                    operators
                    and operators[-1] in self.operators
                    and self.precedence[operators[-1]] >= self.precedence[token]
                ):
                    self._apply_operator(operators, values)
                operators.append(token)
            else:
                try:
                    values.append(float(token))
                except ValueError:
                    raise ValueError(f"invalid token: {token}")
        while operators:
            self._apply_operator(operators, values)
        if len(values) != 1:
            raise ValueError("invalid expression")
        return values[0]
    def _apply_operator(self, operators, values):
        if not operators:
            return
        operator = operators.pop()
        if len(values) < 2:
            raise ValueError(f"not enough operands for operator {operator}")
        b = values.pop()
        a = values.pop()
        values.append(self.operators[operator](a, b))

# render.py
def render(expression, result):
    if isinstance(result, float) and result.is_integer():
        result_str = str(int(result))
    else:
        result_str = str(result)
    box_width = max(len(expression), len(result_str)) + 4
    box = []
    box.append("┌" + "─" * box_width + "┐")
    box.append(
        "│" + " " * 2 + expression + " " * (box_width - len(expression) - 2) + "│"
    )
    box.append("│" + " " * box_width + "│")
    box.append("│" + " " * 2 + "=" + " " * (box_width - 3) + "│")
    box.append("│" + " " * box_width + "│")
    box.append(
        "│" + " " * 2 + result_str + " " * (box_width - len(result_str) - 2) + "│"
    )
    box.append("└" + "─" * box_width + "┘")
    return "\n".join(box)

This is the final structure:

├── calculator
│   ├── main.py
│   ├── pkg
│   │   ├── calculator.py
│   │   └── render.py
│   └── tests.py
├── main.py
├── pyproject.toml
├── README.md
└── uv.lock

Run the calculator tests:

uv run calculator/tests.py

Hopefully the tests all pass!

Now, run the calculator app:

uv run calculator/main.py "3 + 5"

Hopefully you get 8!

2️⃣ Get Files

Now that our agent has a playground, it’s time to teach it how to actually do things. Let’s start simple: we’ll give it the ability to list the contents of a directory and check each file’s metadata (like its name and size).

But before we plug this into our LLM-powered agent, we need to focus on building the function itself. Remember, LLMs deal with text, so our function’s job will be straightforward:

👉 Take in a directory path as input, and return a clean, human-readable string that describes what’s inside that directory.

This will be our agent’s very first “skill.”

Assignment

Create a new directory called functions in the root of your project (not inside the calculator directory). Inside, create a new file called get_files_info.py. Inside, write this function:

def get_files_info(working_directory, directory="."):

Here is my project structure so far:

project_root/
 ├── calculator/
 │   ├── main.py
 │   ├── pkg/
 │   │   ├── calculator.py
 │   │   └── render.py
 │   └── tests.py
 └── functions/
     └── get_files_info.py

The directory parameter should be treated as a relative path within the working_directory. Use os.path.join(working_directory, directory) to create the full path, then validate it stays within the working directory boundaries.

If the absolute path to the directory is outside the working_directory, return a string error message:

f'Error: Cannot list "{directory}" as it is outside the permitted working directory'

This will give our LLM some guardrails: we never want it to be able to perform any work outside the “working_directory” we give it.

Without this restriction, the LLM might go running amok anywhere on the machine, reading sensitive files or overwriting important data. This is a very important step that we’ll bake into every function the LLM can call.

If the directory argument is not a directory, again, return an error string:

f'Error: "{directory}" is not a directory'

All of our “tool call” functions, including get_files_info, should always return a string. If errors can be raised inside them, we need to catch those errors and return a string describing the error instead. This will allow the LLM to handle the errors gracefully.

Build and return a string representing the contents of the directory. It should use this format:

- README.md: file_size=1032 bytes, is_dir=False
- src: file_size=128 bytes, is_dir=True
- package.json: file_size=1234 bytes, is_dir=False

I’ve listed useful standard library functions in the tips section.

The exact file sizes and even the order of files may vary depending on your operating system and file system. Your output doesn’t need to match the example byte-for-byte, just the overall format

If any errors are raised by the standard library functions, catch them and instead return a string describing the error. Always prefix error strings with “Error:”.

To import from a subdirectory, use this syntax: from DIRNAME.FILENAME import FUNCTION_NAME

Where DIRNAME is the name of the subdirectory, FILENAME is the name of the file without the .py extension, and FUNCTION_NAME is the name of the function you want to import.

We need a way to manually debug our new get_files_info function! Create a new tests.py file in the root of your project. When executed directly (uv run tests.py) it should run the following function calls and output the results matching the formatting below (not necessarily the exact numbers).:
get_files_info("calculator", "."):
Result for current directory: - main.py: file_size=576 bytes, is_dir=False - tests.py: file_size=1343 bytes, is_dir=False - pkg: file_size=92 bytes, is_dir=True
get_files_info("calculator", "pkg"):
Result for 'pkg' directory: - calculator.py: file_size=1739 bytes, is_dir=False - render.py: file_size=768 bytes, is_dir=False
get_files_info("calculator", "/bin"):
Result for '/bin' directory: Error: Cannot list "/bin" as it is outside the permitted working directory
get_files_info("calculator", "../"):
Result for '../' directory: Error: Cannot list "../" as it is outside the permitted working directory
Run uv run tests.py, and ensure your function works as expected.

Tips

Here are some standard library functions you’ll find helpful:

os.path.abspath(): Get an absolute path from a relative path
os.path.join(): Join two paths together safely (handles slashes)
.startswith(): Check if a string starts with a substring
os.path.isdir(): Check if a path is a directory
os.listdir(): List the contents of a directory
os.path.getsize(): Get the size of a file
os.path.isfile(): Check if a path is a file
.join(): Join a list of strings together with a separator

3️⃣ Get File Content

Listing files is great — but what if our agent wants to actually look inside one of them? That’s our next step.

We’ll build a function that can read the contents of a file and return it as a string. If something goes wrong (say, the file doesn’t exist or can’t be opened), the function will instead return a clear error message.

And just like before, we’ll keep things safe by scoping this function to a specific working directory, so our agent doesn’t wander off into places it shouldn’t.

Assignment

Create a new function in your functions directory. Here's the signature I used:

def get_file_content(working_directory, file_path):

If the file_path is outside the working_directory, return a string with an error:

f'Error: Cannot read "{file_path}" as it is outside the permitted working directory'

If the file_path is not a file, again, return an error string:

f'Error: File not found or is not a regular file: "{file_path}"'

Read the file and return its contents as a string
I’ll list some useful standard library functions in the tips section below.
If the file is longer than 10000 characters, truncate it to 10000 characters and append this message to the end [...File "{file_path}" truncated at 10000 characters].
Instead of hard-coding the 10000 character limit, I stored it in a config.py file.

We don’t want to accidentally read a gigantic file and send all that data to the LLM… that’s a good way to burn through our token limits.

If any errors are raised by the standard library functions, catch them and instead return a string describing the error. Always prefix errors with “Error:”.
Create a new “lorem.txt” file in the calculator directory. Fill it with at least 20,000 characters of lorem ipsum text. You can generate some here.
Update your tests.py file. Remove all the calls to get_files_info, and instead test get_file_content("calculator", "lorem.txt"). Ensure that it truncates properly.
Remove the lorem ipsum test, and instead test the following cases:
get_file_content("calculator", "main.py")
get_file_content("calculator", "pkg/calculator.py")
get_file_content("calculator", "/bin/cat") (this should return an error string)
get_file_content("calculator", "pkg/does_not_exist.py") (this should return an error string)

Tips

os.path.abspath: Get an absolute path from a relative path
os.path.join: Join two paths together safely (handles slashes)
.startswith: Check if a string starts with a specific substring
os.path.isfile: Check if a path is a file

Example of reading from a file:

MAX_CHARS = 10000
with open(file_path, "r") as f:
    file_content_string = f.read(MAX_CHARS)

4️⃣ Write File

Up to this point, our program has been strictly read-only — it could peek inside directories and files, but not make any changes.

Now things are about to get both dangerous and fun. 🚀
We’re giving our agent the ability to write new files and even overwrite existing ones.

This step marks a big shift: instead of just observing, our agent can now create and modify — a true step toward acting like a real assistant.

Assignment

Create a new function in your functions directory. Here's the signature I used:

def write_file(working_directory, file_path, content):

If the file_path is outside of the working_directory, return a string with an error:

f'Error: Cannot write to "{file_path}" as it is outside the permitted working directory'

If the file_path doesn't exist, create it. As always, if there are errors, return a string representing the error, prefixed with "Error:".
Overwrite the contents of the file with the content argument.
If successful, return a string with the message:

f'Successfully wrote to "{file_path}" ({len(content)} characters written)'

It’s important to return a success string so that our LLM knows that the action it took actually worked. Feedback loops, feedback loops, feedback loops!

Remove your old tests from tests.py and add three new ones, as always print the results of each:
write_file("calculator", "lorem.txt", "wait, this isn't lorem ipsum")
write_file("calculator", "pkg/morelorem.txt", "lorem ipsum dolor sit amet")
write_file("calculator", "/tmp/temp.txt", "this should not be allowed")

Tips

os.path.exists: Check if a path exists
os.makedirs: Create a directory and all its parents
os.path.dirname: Return the directory name

Example of writing to a file:

with open(file_path, "w") as f:
    f.write(content)

5️⃣ Run Python

If you thought allowing an LLM to write files was a bad idea…

You ain’t seen nothin’ yet! (praise the basilisk)

It’s time to build the functionality for our Agent to run arbitrary Python code.

Now, it’s worth pausing to point out the inherent security risks here. We have a few things going for us:

We’ll only allow the LLM to run code in a specific directory (the working_directory).
We’ll use a 30-second timeout to prevent it from running indefinitely.

But aside from that… yes, the LLM can run arbitrary code that we (or it) places in the working directory… so be careful. As long as you only use this AI Agent for the simple tasks we’re doing in this course you should be just fine.

Do not give this program to others for them to use! It does not have all the security and safety features that a production AI agent would have. It is for learning purposes only.

Assignment

Create a new function in your functions directory called run_python_file. Here’s the signature to use:

def run_python_file(working_directory, file_path, args=[]):

If the file_path is outside the working directory, return a string with an error:

f'Error: Cannot execute "{file_path}" as it is outside the permitted working directory'

If the file_path doesn’t exist, return an error string:

f'Error: File "{file_path}" not found.'

If the file doesn’t end with “.py”, return an error string:

f'Error: "{file_path}" is not a Python file.'

Use the subprocess.run function to execute the Python file and get back a "completed_process" object. Make sure to:
Set a timeout of 30 seconds to prevent infinite execution
Capture both stdout and stderr
Set the working directory properly
Pass along the additional args if provided
Return a string with the output formatted to include:
The stdout prefixed with STDOUT:, and stderr prefixed with STDERR:. The "completed_process" object has a stdout and stderr attribute.
If the process exits with a non-zero code, include “Process exited with code X”
If no output is produced, return “No output produced.”
If any exceptions occur during execution, catch them and return an error string:

f"Error: executing Python file: {e}"

Update your tests.py file with these test cases, printing each result:
run_python_file("calculator", "main.py") (should print the calculator's usage instructions)
run_python_file("calculator", "main.py", ["3 + 5"]) (should run the calculator... which gives a kinda nasty rendered result)
run_python_file("calculator", "tests.py")
run_python_file("calculator", "../main.py") (this should return an error)
run_python_file("calculator", "nonexistent.py") (this should return an error)

🎯 Wrapping Up Chapter 2

In this chapter, we leveled up our agent by giving it a real set of skills:

🧮 Work with our calculator project
📂 List directories and inspect file metadata
📄 Read file contents safely
🗂️ Fetch specific files
✍️ Write and overwrite files
🐍 Run Python code inside the project

With these abilities, our agent has moved from being a simple observer to an active participant — capable of exploring, editing, and executing code in its environment.

But we’re not stopping here. 🚀
In Chapter 3, we’ll bring these functions together and integrate them with the LLM, so our agent can decide when and how to use them. That’s when the magic really starts. ✨

👉 Stay tuned for Chapter 3: Turning Functions into an Agent.