Connecting LangChain

LangChain covers the LLM pipeline surface: prompts, LCEL chains, document loaders, vector stores, memory abstractions, plus the LangGraph StateGraph for multi-agent routing. Flux supplies the operational layer around that pipeline (durable execution, retries, scheduling, secrets) without replacing any of it.

The integration follows the same pattern Flux uses for any Python library: call LangChain or LangGraph code inside a @task function. Flux persists state around the call; LangChain handles what happens inside.

Prerequisites

Install the LangChain packages your workflow needs alongside Flux. The examples on this page use Ollama for local inference:

# Conversational chain
pip install langchain-core langchain-ollama

# RAG pipeline (adds Chroma and text-splitting utilities)
pip install langchain-core langchain-ollama langchain-chroma langchain-community langchain-text-splitters

# LangGraph multi-agent graph
pip install langchain-core langchain-ollama langgraph

Pull a local model and start Ollama before running any example:

ollama pull llama3
ollama pull nomic-embed-text  # only needed for RAG
ollama serve

Wrapping a chain as a task

The pattern is the same regardless of chain complexity. Build the chain inside a @task function, call await chain.ainvoke(...), and return the result. Flux retries the whole call if it raises an exception.

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_ollama import ChatOllama
from flux import ExecutionContext, task, workflow
from flux.tasks import pause
from typing import Any


@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=60)
async def call_langchain_chat(
    history: list[dict[str, str]],
    user_message: str,
    system_prompt: str,
    model: str,
    ollama_url: str,
) -> tuple[list[dict[str, str]], str]:
    """Process one conversation turn via LangChain ChatOllama."""
    chat_history = InMemoryChatMessageHistory()
    for msg in history:
        if msg["role"] == "user":
            chat_history.add_user_message(msg["content"])
        elif msg["role"] == "assistant":
            chat_history.add_ai_message(msg["content"])

    chat_history.add_user_message(user_message)

    langchain_messages: list[Any] = [
        SystemMessage(content=system_prompt),
        *chat_history.messages,
    ]

    llm = ChatOllama(model=model, base_url=ollama_url)
    response = await llm.ainvoke(langchain_messages)
    chat_history.add_ai_message(response.content)

    serialized = [
        {"role": "user" if isinstance(m, HumanMessage) else "assistant", "content": m.content}
        for m in chat_history.messages
    ]
    return serialized, response.content

The retry_max_attempts=3 setting means Flux retries the full chain call — including the ChatOllama instantiation and the ainvoke — up to three times on failure. This covers transient connectivity problems with a local Ollama server.

Using the task in a workflow

A conversational workflow passes the message history between turns using pause(). Each resume call delivers the next user message without re-running earlier turns:

@workflow
async def conversational_agent_langchain(ctx: ExecutionContext[dict[str, Any]]):
    """
    Conversational agent backed by a LangChain ChatOllama chain.

    Input:
        message      (required) — first user message
        system_prompt (optional) — default: "You are a helpful AI assistant."
        model        (optional) — Ollama model, default "llama3"
        max_turns    (optional) — maximum conversation turns, default 10
        ollama_url   (optional) — Ollama server URL, default "http://localhost:11434"

    Resume input:
        message (required) — next user message
    """
    initial_input = ctx.input or {}
    system_prompt = initial_input.get("system_prompt", "You are a helpful AI assistant.")
    model = initial_input.get("model", "llama3")
    max_turns = initial_input.get("max_turns", 10)
    ollama_url = initial_input.get("ollama_url", "http://localhost:11434")

    first_message = initial_input.get("message", "")
    if not first_message:
        return {"error": "No message provided", "execution_id": ctx.execution_id}

    messages: list[dict[str, str]] = []
    messages, _ = await call_langchain_chat(
        messages, first_message, system_prompt, model, ollama_url
    )

    for turn in range(1, max_turns):
        resume_input = await pause(f"waiting_for_user_input_turn_{turn}")
        next_message = resume_input.get("message", "") if resume_input else ""
        if not next_message:
            return {
                "status": "ended",
                "turn_count": len(messages) // 2,
                "conversation_history": messages,
                "execution_id": ctx.execution_id,
            }
        messages, _ = await call_langchain_chat(
            messages, next_message, system_prompt, model, ollama_url
        )

    return {
        "status": "max_turns_reached",
        "conversation_history": messages,
        "turn_count": len(messages) // 2,
        "execution_id": ctx.execution_id,
    }

Start and continue a conversation from the CLI:

# Start
flux workflow run conversational_agent_langchain '{"message": "Why is the sky blue?"}'

# Resume with the execution ID from the first call
flux workflow resume conversational_agent_langchain <execution_id> \
    '{"message": "Why does it turn red at sunset?"}'

RAG pipeline

A retrieval-augmented generation pipeline needs two workflows: one to index documents, one to query them. Each stage is a separate Flux task. Splitting the pipeline this way means a failure in embedding — say, the Ollama server is slow — retries only the embedding step, not the document loading or text splitting that already finished.

from pathlib import Path
from langchain_chroma import Chroma
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from flux import ExecutionContext, task, workflow


@task
async def load_documents(docs_path: str) -> list[Document]:
    loader = DirectoryLoader(
        docs_path,
        glob="**/*.md",
        loader_cls=TextLoader,
        loader_kwargs={"encoding": "utf-8"},
        show_progress=False,
    )
    return loader.load()


@task
async def split_documents(
    documents: list[Document],
    chunk_size: int = 500,
    chunk_overlap: int = 50,
) -> list[Document]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
    return splitter.split_documents(documents)


@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=300)
async def build_vector_store(
    chunks: list[Document],
    collection_name: str,
    embedding_model: str = "nomic-embed-text",
    ollama_url: str = "http://localhost:11434",
) -> str:
    persist_dir = Path.home() / ".flux" / "rag_indexes" / collection_name
    persist_dir.mkdir(parents=True, exist_ok=True)

    embeddings = OllamaEmbeddings(model=embedding_model, base_url=ollama_url)
    Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        collection_name=collection_name,
        persist_directory=str(persist_dir),
    )
    return str(persist_dir)


@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=60)
async def generate_rag_response(
    query: str,
    retrieved_docs: list[Document],
    model: str = "llama3",
    ollama_url: str = "http://localhost:11434",
) -> dict[str, Any]:
    context = "\n\n".join(
        f"[Source: {doc.metadata.get('source', 'unknown')}]\n{doc.page_content.strip()}"
        for doc in retrieved_docs
    )

    prompt = ChatPromptTemplate.from_messages([
        ("system", "Answer questions based on the provided documentation context."),
        ("human", "Context:\n\n{context}\n\nQuestion: {question}"),
    ])

    chain = prompt | ChatOllama(model=model, base_url=ollama_url) | StrOutputParser()
    answer = await chain.ainvoke({"context": context, "question": query})

    return {
        "answer": answer,
        "sources": [{"source": doc.metadata.get("source", "unknown")} for doc in retrieved_docs],
    }

The two workflows are thin orchestrators that call these tasks in order:

@workflow
async def rag_index_langchain(ctx: ExecutionContext[dict[str, Any]]):
    """
    Index markdown documents into a Chroma vector store.

    Input:
        docs_path       (required) — path to directory of markdown files
        collection_name (required) — unique name for this collection
        chunk_size      (optional) — characters per chunk, default 500
        chunk_overlap   (optional) — chunk overlap, default 50
        embedding_model (optional) — Ollama embedding model, default "nomic-embed-text"
        ollama_url      (optional) — Ollama server URL
    """
    input_data = ctx.input or {}
    docs_path = input_data.get("docs_path")
    collection_name = input_data.get("collection_name")
    if not docs_path or not collection_name:
        return {"error": "docs_path and collection_name are required", "execution_id": ctx.execution_id}

    documents = await load_documents(docs_path)
    chunks = await split_documents(
        documents,
        input_data.get("chunk_size", 500),
        input_data.get("chunk_overlap", 50),
    )
    persist_path = await build_vector_store(
        chunks,
        collection_name,
        input_data.get("embedding_model", "nomic-embed-text"),
        input_data.get("ollama_url", "http://localhost:11434"),
    )

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


@workflow
async def rag_query_langchain(ctx: ExecutionContext[dict[str, Any]]):
    """
    Query a pre-indexed Chroma collection.

    Input:
        collection_name (required) — name used when indexing
        query           (required) — question to answer
        llm_model       (optional) — Ollama LLM, default "llama3"
        embedding_model (optional) — Ollama embedding model, default "nomic-embed-text"
        top_k           (optional) — chunks to retrieve, default 3
        ollama_url      (optional) — Ollama server URL
    """
    input_data = ctx.input or {}
    collection_name = input_data.get("collection_name")
    query = input_data.get("query")
    if not collection_name or not query:
        return {"error": "collection_name and query are required", "execution_id": ctx.execution_id}

    ollama_url = input_data.get("ollama_url", "http://localhost:11434")
    embedding_model = input_data.get("embedding_model", "nomic-embed-text")

    persist_dir = Path.home() / ".flux" / "rag_indexes" / collection_name
    embeddings = OllamaEmbeddings(model=embedding_model, base_url=ollama_url)
    vector_store = Chroma(
        collection_name=collection_name,
        embedding_function=embeddings,
        persist_directory=str(persist_dir),
    )
    retrieved_docs = await vector_store.asimilarity_search(
        query, k=input_data.get("top_k", 3)
    )

    result = await generate_rag_response(
        query, retrieved_docs,
        input_data.get("llm_model", "llama3"),
        ollama_url,
    )

    return {
        "query": query,
        "answer": result["answer"],
        "sources": result["sources"],
        "collection_name": collection_name,
        "execution_id": ctx.execution_id,
    }

Run both steps from the CLI:

# Index
flux workflow run rag_index_langchain '{
    "docs_path": "./docs",
    "collection_name": "my_docs"
}'

# Query (reuse the index as many times as needed)
flux workflow run rag_query_langchain '{
    "collection_name": "my_docs",
    "query": "How does task caching work?"
}'

LangGraph StateGraph

LangGraph’s StateGraph manages typed state, parallel node execution, and conditional edge routing inside a single graph. Wrapping the compiled graph in a Flux @task gives you durability and retries around the entire graph run.

import operator
from typing import Annotated, TypedDict
from langchain_ollama import ChatOllama
from langgraph.graph import END, START, StateGraph
from flux import ExecutionContext, task, workflow


class ReviewState(TypedDict):
    code: str
    model: str
    ollama_url: str
    skip_testing: bool
    reviews: Annotated[list[dict], operator.add]


async def make_reviewer(agent_name: str, state: ReviewState) -> dict:
    llm = ChatOllama(model=state["model"], base_url=state["ollama_url"])
    messages = [
        ("system", f"You are a {agent_name} code reviewer."),
        ("human", f"Review this code:\n```\n{state['code']}\n```\nReturn a JSON array of findings."),
    ]
    response = await llm.ainvoke(messages)
    return {"reviews": [{"agent": agent_name, "content": response.content}]}


def should_run_testing(state: ReviewState) -> str:
    return "skip" if state.get("skip_testing", False) else "run"


@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=300)
async def run_langgraph_review(
    code: str,
    model: str,
    ollama_url: str,
    skip_testing: bool = False,
) -> list[dict]:
    """Build and run the LangGraph review graph, returning all agent results."""
    builder = StateGraph(ReviewState)

    for name in ("security", "performance", "style"):
        builder.add_node(name, lambda s, n=name: make_reviewer(n, s))
        builder.add_edge(START, name)
        builder.add_edge(name, END)

    builder.add_node("testing", lambda s: make_reviewer("testing", s))
    builder.add_conditional_edges(START, should_run_testing, {"run": "testing", "skip": END})
    builder.add_edge("testing", END)

    graph = builder.compile()
    initial_state: ReviewState = {
        "code": code,
        "model": model,
        "ollama_url": ollama_url,
        "skip_testing": skip_testing,
        "reviews": [],
    }
    final_state = await graph.ainvoke(initial_state)
    return final_state["reviews"]


@workflow
async def multi_agent_code_review_langgraph(ctx: ExecutionContext[dict]):
    """
    Multi-agent code review via LangGraph StateGraph, wrapped in a Flux workflow.

    Input:
        code         (required) — source code to review
        model        (optional) — Ollama model, default "llama3.2"
        ollama_url   (optional) — Ollama server URL
        skip_testing (optional) — skip testing agent, default false
    """
    input_data = ctx.input or {}
    code = input_data.get("code")
    if not code:
        return {"error": "No code provided", "execution_id": ctx.execution_id}

    reviews = await run_langgraph_review(
        code,
        input_data.get("model", "llama3.2"),
        input_data.get("ollama_url", "http://localhost:11434"),
        input_data.get("skip_testing", False),
    )

    return {
        "reviews": reviews,
        "agents_run": [r["agent"] for r in reviews],
        "execution_id": ctx.execution_id,
    }

flux workflow run multi_agent_code_review_langgraph '{
    "code": "def login(user, pw):\n    q = f\"SELECT * FROM users WHERE name={user}\"\n    return db.execute(q)",
    "model": "llama3.2"
}'

# Skip the testing agent (demonstrates the conditional edge)
flux workflow run multi_agent_code_review_langgraph '{
    "code": "def add(a, b): return a + b",
    "skip_testing": true
}'

How the layers divide

LangChain and LangGraph own the LLM interaction. Flux owns everything outside that boundary.

Error handling

Wrap the chain or graph call in a try/except and re-raise as RuntimeError with a message that identifies the failure mode. Flux triggers its retry logic on any exception from the task body:

@task.with_options(retry_max_attempts=3, retry_delay=1, retry_backoff=2, timeout=60)
async def call_chain(query: str, model: str, ollama_url: str) -> str:
    try:
        llm = ChatOllama(model=model, base_url=ollama_url)
        response = await llm.ainvoke([("human", query)])
        return response.content
    except Exception as e:
        raise RuntimeError(
            f"LangChain call failed: {e}. "
            "Verify that Ollama is running (ollama serve) and the model is pulled."
        ) from e

After all retries are exhausted, Flux marks the task and its parent workflow as failed. Check the output with:

flux execution show <execution-id>

Running the full examples

Three complete examples live in examples/ai/langchain/ in the Flux source tree: conversational_agent.py, rag_pipeline.py, and multi_agent_code_review.py. Start Ollama first, then run them directly or via Flux:

# Pull required models
ollama pull llama3
ollama pull nomic-embed-text  # for the RAG example

# Start Flux
flux start server &
flux start worker worker-1 &

# Conversational agent
flux workflow run conversational_agent_langchain '{"message": "What is a Flux workflow?"}'

# RAG: index then query
flux workflow run rag_index_langchain '{"docs_path": "./examples/ai/docs", "collection_name": "flux_docs"}'
flux workflow run rag_query_langchain '{"collection_name": "flux_docs", "query": "How does caching work?"}'

# Multi-agent code review (LangGraph)
flux workflow run multi_agent_code_review_langgraph '{
    "code": "def process(data):\n    return [x*2 for x in data]",
    "model": "llama3.2"
}'