# How to Set Up Retrieval-Augmented Generation (RAG) in Ax: A Complete Guide

> Learn how to set up Retrieval-Augmented Generation (RAG) in Ax. This guide shows you how Ax provides a production-ready RAG implementation using a vector database and configuration.

- Repository: [Ax/ax](https://github.com/ax-llm/ax)
- Tags: how-to-guide
- Published: 2026-02-25

---

**Ax provides a production-ready RAG implementation through the `axRAG` factory function, which builds a multi-hop, self-healing pipeline on top of the AxFlow engine using just a vector database query function and a configuration object.**

The Ax framework (ax-llm/ax) offers a declarative approach to building complex AI pipelines. To set up Retrieval-Augmented Generation (RAG) in Ax, you use the `axRAG` factory exported from the core package, which orchestrates multi-hop retrieval, parallel sub-queries, and optional answer healing with minimal boilerplate.

## Core Architecture of axRAG

### The AxFlow Pipeline Engine

At the heart of Ax's RAG implementation lies **AxFlow**, a declarative pipeline engine that composes AI-powered nodes into a runtime graph. According to the source code in [`src/ax/prompts/rag.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/rag.ts), the `axRAG` factory returns an `AxFlow` instance that automatically handles async execution, retries, and parallel processing. The engine injects a mutable state object that travels through the pipeline, accumulating context and tracking execution metadata across hops.

### RAG Pipeline Nodes

The pipeline consists of ten specialized nodes with typed string signatures, defined at lines 36-75 of [`src/ax/prompts/rag.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/rag.ts). Each node performs a specific transformation:

- **`queryGenerator`** – Creates the initial search query from the user question.
- **`contextualizer`** – Merges retrieved documents with accumulated context.
- **`qualityAssessor`** – Scores answer completeness and identifies gaps.
- **`questionDecomposer`** – Breaks complex questions into parallel sub-queries.
- **`evidenceSynthesizer`** – Pools evidence from multiple retrieval streams.
- **`gapAnalyzer`** – Determines if additional information is required.
- **`answerGenerator`** – Produces the draft answer from current context.
- **`queryRefiner`** – Rewrites queries for subsequent hops.
- **`qualityValidator`** – Evaluates final answer quality against thresholds.
- **`answerHealer`** – Fetches healing context to improve low-quality answers.

### Four-Phase Execution Flow

The `axRAG` pipeline executes across four distinct phases implemented between lines 101-368:

1. **Multi-hop Retrieval** (lines 101-184): A `while` loop iteratively generates queries, executes the user-provided `queryFn` against your vector store, and refines context until reaching `maxHops` or satisfying the `qualityThreshold`.
2. **Parallel Sub-query Processing** (lines 186-259): Decomposes questions and runs `Promise.all` on multiple vector queries concurrently, maximizing throughput for complex questions.
3. **Answer Generation** (lines 261-285): Executes a single `answerGenerator` node to synthesize the draft response from aggregated context.
4. **Self-healing Loops** (lines 287-368): Conditionally validates answer quality and enters a healing `while` loop (lines 322-363) that fetches additional context via `answerHealer` until meeting the `qualityTarget` or exhausting three healing attempts.

## Configuration Options

When invoking `axRAG(queryFn, options)`, you tune behavior through the options object (default values defined at lines 24-28):

- **`maxHops`** – Maximum retrieval rounds for multi-hop exploration.
- **`qualityThreshold`** – Minimum completeness score to exit early from retrieval phases.
- **`maxIterations`** – Parallel sub-query cycles allowed in Phase 2.
- **`qualityTarget`** – Desired quality score for the healing loop to achieve.
- **`disableQualityHealing`** – Boolean flag to disable self-healing for faster execution.

The factory re-exports from [`src/ax/index.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/index.ts) at line 751, enabling direct import via `import { axRAG } from '@ax-llm/ax'`.

## Implementation Examples

### Minimal Single-Hop Setup

For straightforward use cases with minimal latency requirements, configure `axRAG` with healing disabled and limited hops:

```typescript
import { ai, axRAG } from '@ax-llm/ax';

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY,
});

// Vector database adapter
const queryVectorDB = async (query: string): Promise<string> => {
  const results = await yourVectorDB.query(query);
  return results.join('\n');
};

const fastRAG = axRAG(queryVectorDB, {
  maxHops: 2,
  qualityThreshold: 0.7,
  maxIterations: 1,
  disableQualityHealing: true,
});

const result = await fastRAG.forward(llm, {
  originalQuestion: 'What is renewable energy?',
});

console.log(result.finalAnswer);

```

### Production Multi-Hop Configuration

Enable full capabilities for complex research queries requiring deep exploration and quality validation:

```typescript
import { ai, axRAG } from '@ax-llm/ax';

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY,
});

const queryPinecone = async (query: string): Promise<string> => {
  const embedding = await llm.embed({ texts: [query] });
  const results = await pineconeDB.query({
    table: 'knowledge-base',
    values: embedding.embeddings[0],
    topK: 10,
  });
  return results.matches.map(m => m.metadata.content).join('\n');
};

const rag = axRAG(queryPinecone, {
  maxHops: 4,
  qualityThreshold: 0.75,
  maxIterations: 3,
  qualityTarget: 0.9,
  debug: true,
});

const out = await rag.forward(llm, {
  originalQuestion: 'What are the latest developments in quantum computing for cryptography?',
});

console.log('Answer:', out.finalAnswer);
console.log('Quality:', out.qualityAchieved);
console.log('Hops:', out.totalHops);
console.log('Healing:', out.healingAttempts);

```

Setting `debug: true` emits the full AxFlow execution trace, enabling observability into node-level execution and state transitions.

### Simple One-Shot Alternative

For single-retrieval scenarios without multi-hop logic, Ax provides **`axSimpleRAG`** (documented in [`docs/AXRAG.md`](https://github.com/ax-llm/ax/blob/main/docs/AXRAG.md)):

```typescript
import { ai, axSimpleRAG } from '@ax-llm/ax';

const llm = ai({ name: 'openai', apiKey: process.env.OPENAI_APIKEY });

const simple = axSimpleRAG(queryVectorDB);
const resp = await simple.forward(llm, { question: 'Define AI ethics.' });
console.log(resp.answer);

```

## How the Pipeline Works Under the Hood

**Declarative Flow Construction** – The pipeline uses fluent methods (`node`, `execute`, `map`, `while`, `branch`) to define execution graphs. AxFlow compiles these into an optimized runtime that handles the complex orchestration of the ten RAG nodes without manual async management.

**Static Type Validation** – Every node declares input/output schemas as string signatures (e.g., `'originalQuestion:string, previousContext?:string -> searchQuery:string, queryReasoning:string'`). This enables compile-time validation and automatic prompt generation based on type signatures.

**State Management** – The `map` step at lines 77-100 initializes runtime state including `accumulatedContext`, `currentHop`, and `healingAttempts`. This state object mutates as it passes through each phase, allowing nodes to access previous context without global variables.

**Parallel Execution** – Phase 2 leverages `Promise.all` (lines 222-229) to execute multiple vector queries concurrently, significantly reducing latency for decomposed questions.

**Self-Healing Mechanism** – The quality validation branch (lines 287-368) implements a feedback loop where `qualityValidator` outputs an issue list, and `answerHealer` retrieves targeted context to address those specific gaps, re-running generation until the answer meets the `qualityTarget` or hits the safety cap of three attempts.

## Summary

- **Import `axRAG`** from `@ax-llm/ax` to access the factory function defined in [`src/ax/prompts/rag.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/rag.ts).
- **Provide a `queryFn`** adapter that connects to your vector database (Pinecone, Weaviate, custom).
- **Configure thresholds** via `maxHops`, `qualityThreshold`, and `qualityTarget` to balance accuracy against latency.
- **Enable debugging** with `debug: true` to trace execution through the four-phase pipeline.
- **Use `axSimpleRAG`** for lightweight, single-hop retrieval when multi-hop complexity is unnecessary.

## Frequently Asked Questions

### What is the difference between `axRAG` and `axSimpleRAG`?

`axRAG` implements a full multi-hop pipeline with self-healing loops, parallel sub-query processing, and quality validation across four execution phases. `axSimpleRAG` provides a lightweight wrapper for single-retrieval scenarios without iterative refinement or healing capabilities, making it ideal for straightforward question-answering tasks where latency is critical.

### How does the self-healing mechanism work in Ax RAG?

The self-healing mechanism activates during Phase 4 when the `qualityValidator` node scores the draft answer below the `qualityTarget`. The pipeline enters a healing `while` loop (lines 322-363 in [`src/ax/prompts/rag.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/rag.ts)) where `answerHealer` retrieves additional context based on the validator's specific issue list, then re-runs `answerGenerator` with the enriched context. This continues for up to three attempts or until the quality threshold is satisfied.

### Can I use `axRAG` with any vector database?

Yes. The `axRAG` factory accepts any function matching the signature `(query: string) => Promise<string>` as its first argument. This abstraction allows you to connect Pinecone, Weaviate, Chroma, PostgreSQL with pgvector, or custom retrieval systems. The function simply needs to return relevant document content as a string, which the pipeline's `contextualizer` node then integrates into the accumulated state.

### Where can I find a complete working example of Ax RAG?

The repository includes a production-ready example in [`src/examples/advanced-rag.ts`](https://github.com/ax-llm/ax/blob/main/src/examples/advanced-rag.ts) demonstrating real-world configuration with environment variables, embedding generation, and detailed logging. Additionally, [`docs/AXRAG.md`](https://github.com/ax-llm/ax/blob/main/docs/AXRAG.md) contains usage snippets, performance tuning guides, and sample debug output traces for troubleshooting pipeline execution.