Welcome back, future AI architect! In the previous chapter, we explored the foundational elements of AIPack and how .aip files package your AI agents. Now, we’re ready to tackle a core challenge in AI agent development: managing complexity.

Real-world problems rarely have simple, one-step solutions. Imagine an AI agent tasked with reviewing code, fixing bugs, and then writing documentation. Trying to cram all these responsibilities into a single, massive prompt often leads to chaotic outputs, missed steps, and frustrated users. This is where multi-stage markdown agents come in. They allow us to break down a grand challenge into a series of smaller, more manageable steps, just like a seasoned engineer breaks down a large software project.

In this chapter, you’ll learn how to decompose complex tasks into manageable, sequential stages using AIPack’s markdown agent format. We’ll explore how to define these stages, manage the flow of information between them, and introduce the power of Lua for dynamic control and decision-making within your agent’s workflow. By the end, you’ll be able to build sophisticated agents that can execute intricate, multi-step processes with clarity and precision, elevating your AI solutions beyond simple prompt-response systems.

Core Concepts: Orchestrating Agent Intelligence with Stages

At its heart, a multi-stage markdown agent is an AI agent whose execution is divided into distinct, ordered phases. Each phase (or “stage”) has a specific objective, takes certain inputs, and produces defined outputs, much like functions or microservices in a traditional programming application. This modular approach significantly improves reliability and debuggability.

What are Multi-Stage Markdown Agents?

Think of a complex recipe. You don’t just throw all ingredients into a pot at once. Instead, you might:

  1. Prep Vegetables: Chop and dice all produce.
  2. Sauté Aromatics: Cook onions and garlic until fragrant.
  3. Combine & Simmer: Add liquids and main ingredients, then let it cook slowly.
  4. Season & Serve: Adjust flavors and plate the dish.

Each of these is a distinct stage, building upon the previous one. Multi-stage markdown agents apply this same principle. They allow you to define a sequence of interactions with an underlying AI model, where the output of one stage becomes the input for the next. This sequential processing enables the agent to progressively refine its understanding and actions.

AIPack uses a simple yet powerful markdown-based syntax to define these stages within your .aip file. Each stage is separated by a horizontal rule (---), making the agent’s workflow clear, readable, and easy to follow.

The Role of Lua Logic for Dynamic Control

While markdown stages are excellent for defining sequential prompts, what if your agent needs to make decisions? What if you want it to conditionally proceed to the next stage based on the outcome of the current one, loop back, or even dynamically modify inputs?

This is where Lua logic becomes indispensable. AIPack allows you to embed Lua code directly within your markdown agent definitions. Lua is a lightweight, fast, and embeddable scripting language, making it perfect for adding dynamic control flow, data manipulation, and complex decision-making capabilities to your AI agents without significant overhead.

With Lua, you can:

  • Inspect and parse the structured output of a previous AI model call (e.g., JSON, YAML).
  • Perform conditional branching (if-else statements) to guide the agent’s path.
  • Iterate over lists or collections of data.
  • Transform, filter, or combine data before passing it to the next stage.
  • Control which stage executes next using _next_stage, or even end the agent’s execution prematurely using _final_output.

📌 Key Idea: Markdown defines what the agent asks and expects; Lua defines how the agent orchestrates its stages, manipulates data, and makes decisions about its workflow.

Agent Definition Structure: Markdown and Lua in Harmony

An AIPack multi-stage agent combines markdown for prompts and Lua for logic. Let’s look at the general structure:

# Agent Name

## Stage 1: Initial Task
This is the prompt for the first stage.
It describes what the AI should do and what format to output.
For example, "Identify key entities from the following text: {{.input}}"

---

## Stage 2: Conditional Logic
This section can contain Lua code to process the output of Stage 1
or make decisions about the next steps.

```lua
-- Lua code goes here
local stage1_output = _stage_output[1] -- Access output of previous stage
if string.find(stage1_output, "error") then
    _next_stage = "Error Handling Stage" -- Jump to a specific stage
else
    _next_stage = "Stage 3: Refinement" -- Or another specific stage
end

Stage 3: Refinement

This is the prompt for the third stage. It might take processed data from Lua or Stage 2 as input. “Refine the entities identified previously: {{.stage2_processed_data}}”


**Explanation of Key Variables and Syntax:**
*   **Markdown Headers (`#`, `##`):** Used for agent and stage titles, improving readability.
*   **Stage Separator (`---`):** A horizontal rule clearly delineates one stage from the next.
*   **Prompt Text:** The natural language instructions for the AI model.
*   **`{{.input}}`:** A placeholder for the initial input provided to the agent when it's run.
*   **Lua Blocks:** Any Lua code placed directly within a stage, enclosed in a fenced ````lua` block, will be executed *before* the prompt in that stage is sent to the AI model (if a prompt exists in that same stage).
*   **`_stage_output[N]`:** A special Lua variable that allows you to access the *raw text output* of stage `N` (1-based index). This is the exact text returned by the AI model in that stage.
*   **`_context`:** A Lua table acting as shared memory for the entire agent's execution. You can store and retrieve data here (`_context.my_variable = "value"`) and access it in subsequent markdown prompts using `{{.my_variable}}`.
*   **`_next_stage`:** A special Lua variable that, when set to a stage title (e.g., `_next_stage = "Error Handling Stage"`), dictates which stage the agent should execute next. If `_next_stage` is not set, the agent proceeds to the next sequential stage in the `.aip` file.
*   **`_final_output`:** A special Lua variable where you can set the final output of the agent, effectively ending its execution and returning this value.

### Context Flow and State Management

In a multi-stage agent, information needs to flow seamlessly between stages. The output of an earlier stage often becomes the input for a later one. AIPack manages this implicitly and explicitly:
*   The raw output of an AI model call for a given stage is available to subsequent Lua blocks and prompts via `_stage_output[N]`.
*   You can explicitly pass processed data between stages by assigning values to variables in the `_context` table within Lua. These `_context` variables are then referenced in subsequent markdown prompts using `{{.variable_name}}`.

🧠 **Important:** Carefully design the output format of each stage's AI model call. Structured outputs (like JSON or YAML) are much easier for Lua to parse and manipulate, leading to more robust and predictable workflows. Unstructured text is difficult to process programmatically.

Here's a visual representation of how stages and Lua logic might interact:

```mermaid
flowchart TD
    A[Start Agent] --> B[Stage 1 Initial Prompt];
    B --> C{Stage 2 Lua Logic};
    C -->|Condition A Met| D[Stage 3A Action A Prompt];
    C -->|Condition B Met| E[Stage 3B Action B Prompt];
    D --> F[End Agent];
    E --> F;

Step-by-Step Implementation: Building a Simple Code Review Agent

Let’s build a practical multi-stage agent: a code review assistant that first identifies potential issues in a Python code snippet and then, if issues are found, suggests specific fixes.

Step 1: Agent Initialization (.aip file basics)

First, create a new file named code_reviewer.aip.

name: Code Reviewer
description: An agent that reviews Python code for issues and suggests fixes.

This sets up the basic metadata for our agent.

Step 2: Defining the First Stage (Issue Identification)

Our first stage will focus on identifying potential issues in the provided Python code. We’ll instruct the AI to output these issues in a structured JSON format, making it easy for Lua to parse. This is crucial for reliable automation.

Add the following to your code_reviewer.aip file, below the description:

name: Code Reviewer
description: An agent that reviews Python code for issues and suggests fixes.

---

## Identify Code Issues
Review the following Python code for potential bugs, style violations (PEP 8), and areas for improvement.
Provide your findings as a JSON array of objects, where each object has "line", "type", and "description" fields. If no issues are found, return an empty JSON array `[]`.

```json
[
  {
    "line": <line_number>,
    "type": "<bug|style|improvement>",
    "description": "<description_of_issue>"
  }
]

Python Code to Review:

{{.input}}


**Explanation:**
*   We use `---` to start a new stage titled `## Identify Code Issues`.
*   The prompt clearly defines the task and, crucially, specifies the *expected JSON output format*. This is vital for the next Lua stage to reliably parse the AI's response.
*   `{{.input}}` will be replaced by the Python code we provide when running the agent.

### Step 3: Introducing Lua for Control Flow

Now, let's add a Lua stage that processes the output of `Identify Code Issues`. If issues are found, we'll proceed to a "Suggest Fixes" stage. If not, the agent will simply report "No issues found."

Add this new section *after* the `## Identify Code Issues` stage (i.e., after its `---` separator):

```aip
# ... (previous content) ...

---

## Identify Code Issues
# ... (previous prompt content) ...

---

## Process Issues and Decide Next Step
```lua
local issues_json_str = _stage_output[1] -- Get raw text output from the first stage
local issues_table = {}

-- AIPack's Lua environment provides built-in functions like `json.decode` to parse JSON.
-- We use pcall for safe execution, in case the AI returns malformed JSON.
local success, parsed_data = pcall(json.decode, issues_json_str)
if success and type(parsed_data) == "table" then
    issues_table = parsed_data
else
    -- Fallback for malformed JSON or empty output, assume no issues
    print("Warning: Could not parse issues JSON or it was empty. Assuming no issues.")
    issues_table = {}
end

_context.identified_issues = issues_table -- Store parsed issues in agent's context for later stages

if #issues_table > 0 then
    _next_stage = "Suggest Fixes" -- If issues exist, explicitly jump to the 'Suggest Fixes' stage
else
    _final_output = "No significant issues found in the code." -- No issues, set a final output and implicitly end
end


**Explanation:**
*   `_stage_output[1]` retrieves the raw text output from the first stage (`## Identify Code Issues`). Remember, stage indices are 1-based.
*   AIPack's Lua environment provides built-in functions like `json.decode` which attempts to parse the JSON string into a Lua table.
*   `pcall` is used for safe execution, catching potential errors during JSON parsing that might occur if the AI doesn't perfectly follow the output format.
*   `_context.identified_issues = issues_table` stores the parsed issues into the agent's shared `_context`. This makes the `issues_table` available to any subsequent stage using `{{.identified_issues}}`.
*   `if #issues_table > 0 then` checks if any issues were parsed. `#issues_table` gets the number of elements in the Lua table.
*   `_next_stage = "Suggest Fixes"` explicitly tells AIPack to jump to the stage titled `## Suggest Fixes` if issues are present.
*   `_final_output = "No significant issues found in the code."` sets the final message if no issues are found, effectively ending the agent's meaningful output here without proceeding to further stages.

### Step 4: Defining the Second Stage (Fix Suggestion)

Finally, we'll define the stage that suggests fixes, but *only* if the Lua logic determines there are issues. This stage will use the `_context.identified_issues` we stored.

Add this new section *after* the `## Process Issues and Decide Next Step` stage:

```aip
# ... (previous content) ...

---

## Process Issues and Decide Next Step
# ... (previous Lua content) ...

---

## Suggest Fixes
Based on the following identified issues, please provide specific, actionable code fixes for the Python code.
Present your suggestions clearly, referencing the line numbers where appropriate.

Identified Issues:
```json
{{.identified_issues | json.encode_pretty}}

Suggested Fixes:


**Explanation:**
*   This stage will only be reached if `_next_stage` was explicitly set to `"Suggest Fixes"` in the previous Lua block.
*   `{{.identified_issues | json.encode_pretty}}` retrieves the Lua table we stored in `_context.identified_issues`. AIPack's templating engine provides filters like `json.encode_pretty` (similar to `json.decode` but for encoding), which formats a Lua table into a pretty-printed JSON string for the AI model to read easily. This ensures the AI gets well-structured input for its task.

### Step 5: Running and Testing the Agent

Save your `code_reviewer.aip` file. Now, let's run it using the `aipack` CLI. Make sure you have `aipack` installed (as covered in Chapter 2) and an AI provider configured (e.g., Ollama, as discussed in Chapter 3).

Let's test with some problematic Python code. Create a file named `test_code.py` with the following content:

```python
# test_code.py
def calculate_area(radius):
    PI = 3.14
    # Missing return statement
    area = PI * radius * radius

def unused_function():
    print("I'm not used!")

class MyClass:
    def __init__(self, value):
        self.value = value
    # Method name is not snake_case
    def GetValue(self):
        return self.value

Now, run the agent, providing the content of test_code.py as input.

For Linux/macOS:

aipack run code_reviewer.aip -i "$(cat test_code.py)"

For Windows PowerShell:

aipack run code_reviewer.aip -i (Get-Content test_code.py -Raw)

You should observe the agent first identifying the issues (Stage 1), then the Lua logic processing them (Stage 2), and finally, if issues were found, providing suggestions (Stage 3). If you provide simple, perfect code, the agent should output “No significant issues found in the code.”

# Example for perfect code:
echo "def greet():\n    print('Hello')" > perfect_code.py
aipack run code_reviewer.aip -i "$(cat perfect_code.py)"
# Expected output: "No significant issues found in the code."

Mini-Challenge: Prioritizing Bug Fixes

Your turn! Let’s refine our code reviewer to exhibit more intelligent prioritization.

Challenge: Modify the code_reviewer.aip agent to prioritize “bug” issues over “style” or “improvement” issues. If any bugs are identified, the “Suggest Fixes” stage should only focus on those bugs. If no bugs are present, it should then suggest fixes for other types of issues (style, improvement).

Hint: In the ## Process Issues and Decide Next Step Lua block, after parsing issues_table, iterate through it to create a new Lua table containing only the “bug” issues. Store this in _context.bug_issues. Then, modify the if condition to check for _context.bug_issues first. If _context.bug_issues is empty, then store all identified_issues (excluding bugs) into a different context variable like _context.other_issues for the “Suggest Fixes” stage. Ensure the prompt for Suggest Fixes then correctly uses either {{.bug_issues | json.encode_pretty}} or {{.other_issues | json.encode_pretty}} based on your Lua logic.

What to observe/learn: This challenge will deepen your understanding of Lua’s table manipulation and conditional logic (if/else if/else) within AIPack’s workflow. You’ll learn how to filter data programmatically and dynamically adjust the agent’s focus, leading to even smarter and more targeted agent flows.

Common Pitfalls & Troubleshooting

Building multi-stage agents introduces new complexities. Here are some common issues and how to approach them:

  1. Unclear AI Model Outputs or Unexpected Input Formats:

    • Pitfall: The AI model in one stage outputs free-form text or an unexpected structure, and the subsequent Lua stage expects perfectly formatted JSON, leading to parsing errors.
    • Troubleshooting: Always be explicit in your markdown prompts about the exact format you expect the AI to return (e.g., “Return a JSON array like this: [...]”). Provide examples. Use print() statements generously in your Lua code to inspect the raw _stage_output values and the results of json.decode. This helps you see exactly what the AI returned and where parsing might have failed.

  2. Complex or Error-Prone Lua Logic:

    • Pitfall: Writing overly complex Lua blocks that are hard to read, debug, or maintain. Syntax errors in Lua will halt your agent’s execution.
    • Troubleshooting: Keep Lua blocks focused on a single responsibility (e.g., parsing, filtering, decision-making). Break down complex logic into smaller, more manageable steps. Use print() liberally to trace variable values and execution paths within your Lua code. Remember, Lua in AIPack is primarily for orchestration and data transformation, not for building full-fledged applications. Test small Lua snippets independently if possible.

  3. Token Limits in Multi-Stage Workflows:

    • Pitfall: While multi-stage agents help manage context by focusing each stage, passing excessively large amounts of data between stages can still hit token limits, especially when _stage_output is very large or _context accumulates too much information over many stages.
    • Troubleshooting: Be mindful of what data you store in _context. Only pass truly necessary information. Consider summarizing or extracting key details from large _stage_output before storing it or passing it to the next stage. Design prompts to be concise and to extract only the most relevant information from the AI, rather than asking for verbose explanations.

Summary

Congratulations! You’ve successfully delved into the powerful world of multi-stage markdown agents with AIPack. You now understand:

  • The necessity of breaking down complex tasks into smaller, manageable, and focused stages for improved reliability and clarity.
  • How markdown defines the declarative prompts and structure for each stage of your agent’s workflow.
  • The critical role of Lua logic for dynamic control flow, conditional execution, advanced data manipulation, and orchestrating interactions between stages.
  • How to manage context and data flow throughout an agent’s lifecycle using _stage_output, _context, and _next_stage variables.
  • Practical steps to build, run, and troubleshoot a multi-stage agent for a real-world task like intelligent code review.

By mastering multi-stage agents, you’re equipped to build more robust, reliable, and intelligent AI solutions that can handle complex, real-world problems. In our next chapter, we’ll explore how to integrate various AI model providers, from local models like Ollama to powerful cloud APIs, giving your agents the diverse intelligence and capabilities they need.

References

This page is AI-assisted and reviewed. It references official documentation and recognized resources where relevant.