End-to-end RAG application

A multi-step RAG pipeline in Flux — document extraction, embedding, vector storage, retrieval, and answer generation, with every step durable.

We will build a retrieval-augmented generation pipeline in Flux. Given a folder of markdown documents, we want to embed them once, store the vectors, and then answer arbitrary questions by retrieving the most relevant chunks and feeding them to an LLM.

The interesting part is not the RAG recipe itself. The interesting part is splitting it across durable Flux tasks so that embedding (the expensive step) runs once and never has to run again, even when the question changes, the worker restarts, or we redeploy the query side of the system.

What you are building

Two workflows that share a vector index on disk.

The ingest workflow runs on a fresh corpus or when documents change. It reads markdown files, splits them into chunks, embeds each chunk, builds a FAISS index, and writes the index plus the chunk metadata to disk.

The query workflow runs on every user question. It loads the index, embeds the question, retrieves the top-k chunks by similarity, and asks an LLM to answer using those chunks as grounding context.

Separating ingest from query is the load-bearing decision. You can redeploy the query workflow ten times a day without re-embedding a thousand documents. You can also scale them independently: ingest is a batch job that runs occasionally on a beefy worker, while query is an online path that wants low-latency workers near your users. Pinning them to different worker pools by label is a one-line change once you have decided where the boundary goes.

Architecture

INGEST                                QUERY
======                                =====
markdown/                             user question
   |                                       |
   v                                       v
extract -> chunk -> embed -> store    retrieve <- (index on disk)
                                           |
                                           v
                                       generate -> answer + sources

Every arrow is a Flux task boundary. Every box is durable: if the worker crashes between “embed” and “store”, the embeddings sit in the event log and the workflow resumes from “store” on the next claim.

We will pick FAISS as the vector store because it is a single dependency, runs entirely on disk, and has no service to host. Production teams usually swap it for Pinecone, Weaviate, pgvector, or ChromaDB; the workflow shape does not change.

  1. Document extraction

    We start with a task that reads markdown files from a directory and returns them as (path, content) records. Nothing exciting yet — but wrapping it in @task means the file list becomes part of the durable event log. A later step that fails will not re-read the disk.

    from pathlib import Path
from flux.task import task


@task
async def load_markdown(docs_path: str) -> list[dict]:
    docs_dir = Path(docs_path)
    if not docs_dir.is_dir():
        raise ValueError(f"Not a directory: {docs_path}")

    files = list(docs_dir.glob("**/*.md"))
    if not files:
        raise ValueError(f"No markdown files under {docs_path}")

    return [
        {"path": str(f.relative_to(docs_dir)), "content": f.read_text(encoding="utf-8")}
        for f in files
    ]

    The task is plain async Python. Flux records its arguments and its return value as events; on replay, it skips straight to the recorded output.

  2. Chunking

    We split each document into overlapping character windows. The simple paragraph-aware splitter below is enough for prose; for code or tables you would want something smarter, but the pipeline shape is unchanged.

    @task
async def chunk_documents(
    documents: list[dict],
    chunk_size: int = 500,
    overlap: int = 50,
) -> list[dict]:
    chunks = []
    for doc in documents:
        content = doc["content"]
        start = 0
        idx = 0
        while start < len(content):
            end = start + chunk_size
            text = content[start:end].strip()
            if text:
                chunks.append({
                    "content": text,
                    "path": doc["path"],
                    "chunk_index": idx,
                })
                idx += 1
            start = end - overlap
    return chunks

    Chunking is deterministic and cheap, so we run it as a single task over the whole corpus. If chunking ever becomes the bottleneck (millions of documents), this is the natural place to fan out — wrap the per-document logic in its own task and call them through parallel(...).

  3. Embedding

    Embedding is where time and money live. A thousand chunks against a hosted embedding model is a few cents and a few seconds, both of which we want to spend exactly once.

    import numpy as np
from ollama import AsyncClient


@task.with_options(
    cache=True,
    retry_max_attempts=3,
    retry_delay=1,
    retry_backoff=2,
    timeout=120,
)
async def embed_chunk(text: str, model: str = "nomic-embed-text") -> list[float]:
    client = AsyncClient(host="http://localhost:11434")
    response = await client.embeddings(model=model, prompt=text)
    return response["embedding"]

    Three things to notice.

    cache=True is the key bit. Flux caches a task’s result against a fingerprint of its arguments. Cache hits are disk-backed and survive across executions, not just within one run. Re-embedding the same chunk text on a later workflow execution is a no-op; the task returns the cached vector without calling Ollama. See Ollama integration for the host configuration and model-pull details.

    retry_max_attempts=3 with exponential backoff covers transient Ollama and network failures. If the embed call returns 500 once, Flux retries before failing the task.

    timeout=120 caps the wait on any single embed. Long enough for a cold model load, short enough that a hung connection eventually surfaces as a TaskTimeoutError instead of stalling the workflow.

  4. Storage

    Once we have a vector per chunk, we build a FAISS index and write it to disk along with the chunk metadata.

    import faiss


@task
async def build_and_save_index(
    chunks: list[dict],
    embeddings: list[list[float]],
    index_name: str,
) -> dict:
    import json
    import pickle

    vectors = np.array(embeddings, dtype=np.float32)
    dim = vectors.shape[1]

    index = faiss.IndexFlatL2(dim)
    index.add(vectors)

    index_dir = Path.home() / ".flux" / "rag_indexes"
    index_dir.mkdir(parents=True, exist_ok=True)

    index_file = index_dir / f"{index_name}.faiss"
    chunks_file = index_dir / f"{index_name}.chunks.pkl"
    meta_file = index_dir / f"{index_name}.meta.json"

    faiss.write_index(index, str(index_file))
    with open(chunks_file, "wb") as f:
        pickle.dump(chunks, f)
    with open(meta_file, "w") as f:
        json.dump({"num_chunks": len(chunks), "dim": dim}, f)

    return {"index_name": index_name, "num_chunks": len(chunks)}

    The output is small — a dict with the index name and chunk count. The bulky bytes go straight to disk under ~/.flux/rag_indexes/. We deliberately do not return the FAISS bytes through the event log; vector indexes get big fast and the event log is not the right place for them.

    If you need event-log-tracked output storage instead of a side-effect write, set output_storage=LocalFileStorage() on the task. For an external index like Pinecone this step becomes a one-line API call.

  5. The ingest workflow

    Now we wire the four tasks together.

    from flux import ExecutionContext
from flux.tasks import parallel
from flux.workflow import workflow


@workflow
async def rag_ingest(ctx: ExecutionContext[dict]):
    params = ctx.input or {}
    docs_path = params["docs_path"]
    index_name = params["index_name"]

    documents = await load_markdown(docs_path)
    chunks = await chunk_documents(documents)

    embeddings = await parallel(
        *(embed_chunk(c["content"]) for c in chunks)
    )

    return await build_and_save_index(chunks, embeddings, index_name)

    The parallel(...) call fans the embedding step across every chunk. On a corpus of a thousand chunks talking to an Ollama instance with a hot model, that is roughly an order-of-magnitude speedup over awaiting them in a loop. Each chunk gets its own task event in the log, so a failure on chunk 437 retries only chunk 437 — not the whole batch.

    A subtle gotcha: parallel(...) runs all coroutines concurrently with asyncio.gather. If your embedding provider rate-limits aggressively, you may want to batch — call parallel(...) on groups of 50 at a time rather than the full list. The shape stays the same; only the iteration changes.

    A second gotcha: each chunk’s embed is a separate task event in the log. A thousand chunks is a thousand events, which is fine for the event log itself but matters for the UI rendering the run. If you want a tighter event count, group the chunks and embed each group in a single task that loops internally; you trade per-chunk retry granularity for log compactness. For most corpora this is not worth doing.

  6. Query and retrieve

    The query side loads the index from disk and finds the closest chunks for a question.

    @task.with_options(retry_max_attempts=3, timeout=60)
async def retrieve(
    query: str,
    index_name: str,
    top_k: int = 3,
) -> list[dict]:
    import pickle

    index_dir = Path.home() / ".flux" / "rag_indexes"
    index = faiss.read_index(str(index_dir / f"{index_name}.faiss"))
    with open(index_dir / f"{index_name}.chunks.pkl", "rb") as f:
        chunks = pickle.load(f)

    query_vec = np.array(
        [await embed_chunk(query)], dtype=np.float32
    )
    distances, indices = index.search(query_vec, top_k)

    return [
        {**chunks[i], "score": float(d)}
        for i, d in zip(indices[0], distances[0])
    ]

    We reuse embed_chunk for the query embedding. Because it is cached, asking the same question twice never re-embeds the question text. That is the “across executions” caching from step 3 paying off a second time on the read path.

  7. Answer generation

    For the final answer we use the Flux agent() primitive. It returns a task that wraps an LLM call with the right provider, system prompt, and tool-loop scaffolding.

    from flux.tasks.ai import agent


async def make_answer_agent():
    return await agent(
        system_prompt=(
            "Answer the user's question using only the provided context. "
            "If the context does not contain the answer, say so plainly. "
            "Cite sources by filename."
        ),
        model="anthropic/claude-sonnet-4-20250514",
        max_tokens=1024,
        stream=False,
    )

    agent() returns a @task. We do not need to register it separately; calling it inside a workflow generates events like any other task. Swap anthropic/claude-sonnet-4-20250514 for ollama/llama3 to run end-to-end locally with no API key. See Anthropic integration for the credentials and model-string details, or the LLM providers overview to compare options.

    Notice the prompt. “Answer using only the provided context” plus “say so plainly when the context does not contain the answer” are the two non-negotiable guardrails for a RAG system. Without the first, the model hallucinates from its training data and the retrieval becomes decorative. Without the second, the model produces confidently-wrong answers on out-of-corpus questions. Both lines belong in the system prompt, not glued onto the user message.

  8. The query workflow

    @workflow
async def rag_query(ctx: ExecutionContext[dict]):
    params = ctx.input or {}
    query = params["query"]
    index_name = params["index_name"]
    top_k = params.get("top_k", 3)

    relevant = await retrieve(query, index_name, top_k)

    context_text = "\n\n".join(
        f"[{c['path']}#{c['chunk_index']}]\n{c['content']}"
        for c in relevant
    )

    answer_agent = await make_answer_agent()
    answer = await answer_agent(
        instruction=query,
        context=context_text,
    )

    return {
        "query": query,
        "answer": answer,
        "sources": [c["path"] for c in relevant],
    }

    rag_query is the workflow you expose to users. Register it once, then call it with flux workflow run rag_query '{"query": "...", "index_name": "..."}' from the CLI, the REST API, or another workflow.

What to remember

Durability per step is what makes the rest possible. Retries, caching, replay, and crash recovery all hang off the fact that each task call is a journaled event with its arguments and its return value.

Caching saves real money on the embed step. Disk-backed and cross-execution means we are not just memoizing within one run; we are memoizing forever, against the fingerprint of the chunk text. Edit a chunk and only that chunk re-embeds.

parallel(...) matters once your corpus is larger than a few dozen chunks. The shape is identical to sequential, but you trade O(N) time for O(1) time bounded by your provider’s concurrency limit.

Retrieval and generation belong in separate workflows. They have different lifecycles — you swap models on the generate side often, you swap embedding strategy on the retrieve side rarely. Splitting them lets each side redeploy without invalidating the other.

Where this shows up next