RAG with durable retrieval

Build a RAG pipeline where each step — load, chunk, embed, index, retrieve, generate — is a Flux task, so embedding work checkpoints to disk and never re-runs on retries or agent restarts.

The expensive part of a RAG pipeline is embedding. Sending hundreds of document chunks through an embedding model can take minutes, and on a paid API it costs real money. Doing that work on every query, or after every worker restart, wastes time and budget that compounds with traffic.

Make each pipeline step a Flux task. Flux checkpoints every task’s output to its event log: embedding runs once, the result is stored, and any subsequent run that reaches the same task with the same arguments reads from the log instead of re-embedding. The pipeline splits into two workflows. One indexes (run once, or whenever the corpus changes). The other queries (run as many times as needed, reusing the saved index).

The pattern

rag_index_documents:
  load_markdown_documents → chunk_documents → generate_embeddings → build_faiss_index → [disk]

rag_query_documents:
  [disk] → retrieve_relevant_chunks → generate_rag_response

Splitting index-time work from query-time work is the core decision. It means embedding runs happen on a schedule or on corpus change, not per query. The index is written to ~/.flux/rag_indexes/ as serialized FAISS + pickled chunk metadata. Any query workflow can load it without re-doing the embedding work.

Within each workflow, every step is a @task. That gives you:

Durable shape

The skeleton uses six tasks and two workflows:

  1. load_markdown_documents — scans a directory recursively, reads every .md file, returns a list of {content, filename, path} dicts.
  2. chunk_documents — splits each document into fixed-size character windows with configurable overlap.
  3. generate_embeddings — calls the Ollama embedding API for each chunk; decorated with retry_max_attempts=3 and timeout=120 because network calls to a local model can stall.
  4. build_faiss_index — wraps the embeddings in a FAISS IndexFlatL2 and serializes it to bytes. Returning bytes rather than a live FAISS object is deliberate: Flux’s task output must be serializable for checkpointing.
  5. retrieve_relevant_chunks — deserializes the FAISS index, embeds the query, runs index.search(), and returns the top-k chunks with similarity scores.
  6. generate_rag_response — formats the retrieved chunks as numbered sources in a prompt and calls the Ollama chat API.

Complete solution

Install dependencies and pull the models you need:

pip install "flux-core[ai]" faiss-cpu numpy ollama
ollama pull llama3
ollama pull nomic-embed-text
from __future__ import annotations

import json
import pickle
from pathlib import Path
from typing import Any

import faiss
import numpy as np
from ollama import AsyncClient

from flux import ExecutionContext, task, workflow


# ---------------------------------------------------------------------------
# Step 1 — Load documents
# ---------------------------------------------------------------------------

@task
async def load_markdown_documents(docs_path: str) -> list[dict[str, str]]:
    docs_dir = Path(docs_path)
    if not docs_dir.exists():
        raise ValueError(f"Directory not found: {docs_path}")
    if not docs_dir.is_dir():
        raise ValueError(f"Path is not a directory: {docs_path}")

    documents = []
    md_files = list(docs_dir.glob("**/*.md"))
    if not md_files:
        raise ValueError(f"No markdown files found in: {docs_path}")

    for md_file in md_files:
        try:
            content = md_file.read_text(encoding="utf-8")
            documents.append({
                "content": content,
                "filename": md_file.name,
                "path": str(md_file.relative_to(docs_dir)),
            })
        except Exception as e:
            print(f"Warning: Failed to read {md_file}: {e}")

    return documents


# ---------------------------------------------------------------------------
# Step 2 — Chunk documents
# ---------------------------------------------------------------------------

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


# ---------------------------------------------------------------------------
# Step 3 — Generate embeddings  (retry + timeout: network call to local model)
# ---------------------------------------------------------------------------

@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=120)
async def generate_embeddings(
    texts: list[str],
    model: str = "nomic-embed-text",
    ollama_url: str = "http://localhost:11434",
) -> list[list[float]]:
    """
    Returns a list of float lists rather than a numpy array so the output
    is JSON-serializable and can be stored in Flux's event log unchanged.
    """
    client = AsyncClient(host=ollama_url)
    embeddings = []
    for text in texts:
        response = await client.embeddings(model=model, prompt=text)
        embeddings.append(response["embedding"])
    return embeddings


# ---------------------------------------------------------------------------
# Step 4 — Build FAISS index
# ---------------------------------------------------------------------------

@task
async def build_faiss_index(embeddings: list[list[float]]) -> bytes:
    """
    Accepts plain lists so this task is compatible with generate_embeddings'
    serializable output. Converts to float32 numpy array internally.
    Serializes the index to bytes — again, so Flux can checkpoint it.
    """
    arr = np.array(embeddings, dtype=np.float32)
    index = faiss.IndexFlatL2(arr.shape[1])
    index.add(arr)
    return faiss.serialize_index(index).tobytes()


# ---------------------------------------------------------------------------
# Indexing workflow  (run once per corpus)
# ---------------------------------------------------------------------------

@workflow
async def rag_index_documents(ctx: ExecutionContext[dict[str, Any]]):
    """
    Index documents and save to disk for later querying.

    Input:
        {
            "docs_path": "./path/to/docs",         # required
            "index_name": "my_docs",               # required
            "chunk_size": 500,                     # optional
            "overlap": 50,                         # optional
            "embedding_model": "nomic-embed-text", # optional
            "ollama_url": "http://localhost:11434"  # optional
        }
    """
    data = ctx.input or {}
    docs_path = data.get("docs_path")
    index_name = data.get("index_name")
    if not docs_path:
        return {"error": "Missing 'docs_path'", "execution_id": ctx.execution_id}
    if not index_name:
        return {"error": "Missing 'index_name'", "execution_id": ctx.execution_id}

    chunk_size     = data.get("chunk_size", 500)
    overlap        = data.get("overlap", 50)
    embedding_model = data.get("embedding_model", "nomic-embed-text")
    ollama_url     = data.get("ollama_url", "http://localhost:11434")

    documents   = await load_markdown_documents(docs_path)
    chunks      = await chunk_documents(documents, chunk_size, overlap)

    if not chunks:
        return {"error": "No chunks created", "execution_id": ctx.execution_id}

    chunk_texts  = [c["content"] for c in chunks]
    embeddings   = await generate_embeddings(chunk_texts, embedding_model, ollama_url)
    index_bytes  = await build_faiss_index(embeddings)

    # Persist index and metadata to disk
    index_dir = Path.home() / ".flux" / "rag_indexes"
    index_dir.mkdir(parents=True, exist_ok=True)

    (index_dir / f"{index_name}_index.faiss").write_bytes(index_bytes)

    with open(index_dir / f"{index_name}_chunks.pkl", "wb") as f:
        pickle.dump(chunks, f)

    metadata = {
        "index_name": index_name,
        "docs_path": str(docs_path),
        "num_documents": len(documents),
        "num_chunks": len(chunks),
        "chunk_size": chunk_size,
        "overlap": overlap,
        "embedding_model": embedding_model,
    }
    with open(index_dir / f"{index_name}_metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)

    return {
        "status": "indexed",
        "index_name": index_name,
        "num_documents": len(documents),
        "num_chunks": len(chunks),
        "execution_id": ctx.execution_id,
    }


# ---------------------------------------------------------------------------
# Step 5 — Retrieve relevant chunks  (retry: embeds query, then does ANN search)
# ---------------------------------------------------------------------------

@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=120)
async def retrieve_relevant_chunks(
    query: str,
    index_bytes: bytes,
    chunks: list[dict[str, Any]],
    embedding_model: str,
    ollama_url: str,
    top_k: int = 3,
) -> list[dict[str, Any]]:
    index = faiss.deserialize_index(np.frombuffer(index_bytes, dtype=np.uint8))

    # Re-use generate_embeddings logic directly (single query)
    client = AsyncClient(host=ollama_url)
    response = await client.embeddings(model=embedding_model, prompt=query)
    query_vec = np.array([response["embedding"]], dtype=np.float32)

    distances, indices = index.search(query_vec, top_k)

    results = []
    for idx, distance in zip(indices[0], distances[0]):
        chunk = chunks[idx]
        results.append({
            "content": chunk["content"],
            "filename": chunk["filename"],
            "path": chunk["path"],
            "chunk_index": chunk["chunk_index"],
            "similarity_score": float(distance),
        })
    return results


# ---------------------------------------------------------------------------
# Step 6 — Generate response
# ---------------------------------------------------------------------------

@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=60)
async def generate_rag_response(
    query: str,
    context_chunks: list[dict[str, Any]],
    model: str,
    ollama_url: str,
) -> dict[str, Any]:
    client = AsyncClient(host=ollama_url)

    context_parts = [
        f"[Source {i}: {chunk['filename']}]\n{chunk['content'].strip()}"
        for i, chunk in enumerate(context_chunks, 1)
    ]
    context = "\n\n".join(context_parts)

    prompt = f"""Context from documentation:

{context}

Question: {query}

Answer based on the context above. If the context doesn't contain relevant information, say so clearly."""

    response = await client.chat(
        model=model,
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant that answers questions based on provided documentation context. Always cite your sources.",
            },
            {"role": "user", "content": prompt},
        ],
    )

    return {
        "answer": response["message"]["content"],
        "sources": [
            {"filename": c["filename"], "path": c["path"]}
            for c in context_chunks
        ],
    }


# ---------------------------------------------------------------------------
# Query workflow  (run per question, reuses saved index)
# ---------------------------------------------------------------------------

@workflow
async def rag_query_documents(ctx: ExecutionContext[dict[str, Any]]):
    """
    Query a pre-indexed corpus using RAG.
    Run rag_index_documents first to create the index.

    Input:
        {
            "index_name": "my_docs",               # required
            "query": "Your question here",         # required
            "llm_model": "llama3",                 # optional
            "top_k": 3,                            # optional
            "ollama_url": "http://localhost:11434"  # optional
        }
    """
    data = ctx.input or {}
    index_name = data.get("index_name")
    query      = data.get("query")
    if not index_name:
        return {"error": "Missing 'index_name'", "execution_id": ctx.execution_id}
    if not query:
        return {"error": "Missing 'query'", "execution_id": ctx.execution_id}

    llm_model  = data.get("llm_model", "llama3")
    top_k      = data.get("top_k", 3)
    ollama_url = data.get("ollama_url", "http://localhost:11434")

    index_dir      = Path.home() / ".flux" / "rag_indexes"
    index_file     = index_dir / f"{index_name}_index.faiss"
    chunks_file    = index_dir / f"{index_name}_chunks.pkl"
    metadata_file  = index_dir / f"{index_name}_metadata.json"

    if not index_file.exists():
        return {
            "error": f"Index '{index_name}' not found. Run rag_index_documents first.",
            "execution_id": ctx.execution_id,
        }

    index_bytes = index_file.read_bytes()

    with open(chunks_file, "rb") as f:
        chunks = pickle.load(f)

    with open(metadata_file) as f:
        metadata = json.load(f)

    embedding_model = metadata.get("embedding_model", "nomic-embed-text")

    relevant_chunks = await retrieve_relevant_chunks(
        query, index_bytes, chunks, embedding_model, ollama_url, top_k,
    )

    result = await generate_rag_response(query, relevant_chunks, llm_model, ollama_url)

    return {
        "query": query,
        "answer": result["answer"],
        "sources": result["sources"],
        "num_sources": len(result["sources"]),
        "index_name": index_name,
        "num_chunks_indexed": metadata.get("num_chunks", 0),
        "execution_id": ctx.execution_id,
    }

Running it

Register both workflows with a running Flux server:

flux workflow register rag_agent_ollama.py

Index your corpus (run once):

flux workflow run rag_index_documents '{
    "docs_path": "./docs",
    "index_name": "my_project",
    "chunk_size": 500
}'

Query against the saved index:

flux workflow run rag_query_documents '{
    "index_name": "my_project",
    "query": "How does retry logic work?",
    "top_k": 3
}'

Check the result and see the full task trace:

flux workflow status rag_query_documents <execution_id>
flux execution show <execution_id> --detailed

The detailed view shows each task’s start time, duration, and whether it read from the event log or executed fresh. On a second query run with the same index_name, retrieve_relevant_chunks loads from disk (fast). On a retry of a failed rag_index_documents run, load_markdown_documents and chunk_documents replay from the checkpoint; only generate_embeddings re-runs from the point of failure.

Why serializable outputs matter

Two tasks in this pipeline make deliberate choices about their return types:

If you return types that Flux cannot serialize, the task output will not checkpoint. The pipeline still runs, but you lose replay protection on that step. Keep task outputs to JSON-safe primitives, bytes, or dataclasses with JSON-safe fields.

Variations

Different embedding models. Change embedding_model in the rag_index_documents input. The metadata file records which model was used, so rag_query_documents picks it up automatically. If you change the model, run rag_index_documents again to rebuild the index — query embeddings and corpus embeddings must come from the same model.

Swap LLM provider. generate_rag_response calls Ollama directly. To use OpenAI or Anthropic, replace the AsyncClient call with the provider’s SDK. The rest of the pipeline is unchanged because the LLM only touches the final task.

Chunking strategy. The current chunker splits on character count. For structured documents, replace chunk_documents with a sentence-aware or markdown-heading-aware splitter. The task signature stays the same; only the body changes.

Re-ranking. Add a rerank_chunks task between retrieve_relevant_chunks and generate_rag_response. A cross-encoder re-ranker (e.g. cross-encoder/ms-marco-MiniLM-L-6-v2) can significantly improve the top-k quality for a small compute cost. Putting it in its own task means it checkpoints independently and can be retried if the model download stalls.

Incremental indexing. To add documents without re-embedding the full corpus, load the existing index bytes from disk, deserialize them, call index.add() with the new embeddings, and serialize back. Wrap this in a separate rag_add_documents workflow that reads the existing metadata.json to verify the embedding model matches.

Multiple corpora. Pass different index_name values. Each corpus gets its own set of files under ~/.flux/rag_indexes/. A single rag_query_documents workflow can query any of them by name.

What’s next