LangChain Mastery

Course Content

LangChain Mastery

7 sections · 109 lessons

Write a function to implement a custom LangChain chain.


Why streaming redaction needs a buffercall9876543210nowforhelp012numberstartsnumber endsA per-chunk regex sees neither half as a phone number.
RunnableGenerator holds text until a word boundary, so a pattern split across two chunks is still caught.

What you need to know

"Custom chain" usually means a step with your own logic: a guard, a lookup, a calculation, or a choice about which sub-chain to call.

Python
from langchain_core.runnables import chainanswer_chain = prompt | llm | StrOutputParser()@chaindef guarded_answer(inputs: dict) -> str:    question = inputs["question"].strip()    if len(question) < 5:        return "Could you tell me a bit more about your question?"    if contains_card_number(question):        return "Please don't share card numbers here. I can help without them."    return answer_chain.invoke({"question": question})pipeline = load_user_context | guarded_answerpipeline.invoke({"question": "How do I change my UPI PIN?", "user_id": "u_812"})
  • @chain wraps the function in a RunnableLambda named after it, so it shows up as guarded_answer in traces.
  • Calling answer_chain.invoke inside it is fine: the inner call is traced as a child of this step.
  • The guards run before the model, so short or risky inputs cost nothing.

A step that works on a stream

A normal function gets the whole input at once, which breaks streaming. RunnableGenerator takes the stream of chunks and yields new chunks:

Python
import refrom langchain_core.runnables import RunnableGeneratorPHONE = re.compile(r"\b[6-9]\d{9}\b")def redact_phones(chunks):    buffer = ""    for text in chunks:        buffer += text        cut = buffer.rfind(" ")                  # release only whole words        if cut > 0:            yield PHONE.sub("[phone]", buffer[:cut + 1])            buffer = buffer[cut + 1:]    yield PHONE.sub("[phone]", buffer)safe_stream = prompt | llm | StrOutputParser() | RunnableGenerator(redact_phones)

The buffer matters: a phone number can be split across two chunks, and replacing chunk by chunk would miss it.

The legacy way

Python
from langchain_classic.chains.base import Chain   # legacy; avoid for new code

A Chain subclass had to declare input_keys, output_keys and _call. It did not stream and needed extra work for async. Mention it only to show you know why it was replaced.

When a function is not enough

If your custom logic needs loops, state that survives between steps, or a pause for human approval, it is a workflow, not a chain step. Build it as a LangGraph graph.

A real-life example

A fintech support bot over its help-centre docs kept getting messages that contained full card numbers. The team added a @chain step called guarded_answer that checks for 16-digit numbers with a Luhn check before any model call and returns a fixed safety message instead. In the first month it stopped about 1,100 messages with card numbers from reaching the model provider and the logs. Because it is a named Runnable, the LangSmith dashboard shows how often the guard fires, and the unit tests call guarded_answer.invoke directly with sample inputs and no model at all.

Follow-up questions to expect

  • "@chain or RunnableLambda?" — The same thing; the decorator is neater for a named function, RunnableLambda(fn) for wrapping an existing one.
  • "Can a custom step be async?" — Yes; decorate an async def, or pass both a sync and an async function to RunnableLambda(func, afunc=...).
  • "How do you give a custom step config, like callbacks?" — Add a second parameter named config to the function; LangChain passes the run config in.