langchain-core LCEL chain with retries and fallbacks, tested offline: working versions (tested 2026-10-15)

langchain-core==1.6.10
pytest==9.1.1

from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import Runnable

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are terse."),
    ("human", "{question}"),
])

def fake_model(text: str) -> GenericFakeChatModel:
    # The messages iterator is consumed on every call, so build a fresh model per use.
    return GenericFakeChatModel(messages=iter([AIMessage(content=text)]))

class Flaky:
    """Raises `exc` for the first `fail_times` calls, then delegates to `inner`."""

    def __init__(self, inner: Runnable, fail_times: int, exc: type[Exception] = TimeoutError):
        self.inner, self.fail_times, self.exc = inner, fail_times, exc
        self.calls = 0

    def __call__(self, value):
        self.calls += 1
        if self.calls <= self.fail_times:
            raise self.exc("transient")
        return self.inner.invoke(value)

def build_chain(primary: Runnable, backup: Runnable, attempts: int = 3):
    # Retry happens first, on the primary only. The fallback runs once retries are exhausted.
    retried = primary.with_retry(
        retry_if_exception_type=(TimeoutError,),
        stop_after_attempt=attempts,
        wait_exponential_jitter=False,
    )
    guarded = retried.with_fallbacks([backup], exceptions_to_handle=(TimeoutError,))
    return prompt | guarded | StrOutputParser()

import pytest
from langchain_core.messages import AIMessage
from langchain_core.runnables import RunnableLambda
from chain import Flaky, build_chain, fake_model

INPUT = {"question": "hi"}

def test_happy_path_returns_text():
    chain = build_chain(fake_model("pong"), fake_model("backup"))
    assert chain.invoke(INPUT) == "pong"

def test_retry_recovers_before_fallback():
    primary = Flaky(fake_model("primary"), fail_times=2)
    chain = build_chain(RunnableLambda(primary), fake_model("backup"), attempts=3)
    assert chain.invoke(INPUT) == "primary"
    assert primary.calls == 3

def test_fallback_after_retries_exhausted():
    primary = Flaky(fake_model("primary"), fail_times=99)
    chain = build_chain(RunnableLambda(primary), fake_model("backup"), attempts=3)
    assert chain.invoke(INPUT) == "backup"
    assert primary.calls == 3

def test_fallback_receives_original_input_not_the_error():
    seen = []

    def capture(value):
        seen.append(value)
        return AIMessage(content="ok")

    primary = Flaky(fake_model("primary"), fail_times=99)
    chain = build_chain(RunnableLambda(primary), RunnableLambda(capture))
    assert chain.invoke(INPUT) == "ok"
    assert "hi" in seen[0].to_string()

def test_unlisted_exception_bypasses_retry_and_fallback():
    primary = Flaky(fake_model("primary"), fail_times=99, exc=ValueError)
    chain = build_chain(RunnableLambda(primary), fake_model("backup"))
    with pytest.raises(ValueError):
        chain.invoke(INPUT)
    assert primary.calls == 1

[pytest]
pythonpath = .