Course Content
LangChain Mastery
7 sections · 109 lessons
Write a function to implement a custom LangChain chain.
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.
1from langchain_core.runnables import chain23answer_chain = prompt | llm | StrOutputParser()45@chain6def guarded_answer(inputs: dict) -> str:7 question = inputs["question"].strip()8 if len(question) < 5:9 return "Could you tell me a bit more about your question?"10 if contains_card_number(question):11 return "Please don't share card numbers here. I can help without them."12 return answer_chain.invoke({"question": question})1314pipeline = load_user_context | guarded_answer15pipeline.invoke({"question": "How do I change my UPI PIN?", "user_id": "u_812"})@chainwraps the function in aRunnableLambdanamed after it, so it shows up asguarded_answerin traces.- Calling
answer_chain.invokeinside 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:
1import re2from langchain_core.runnables import RunnableGenerator34PHONE = re.compile(r"\b[6-9]\d{9}\b")56def redact_phones(chunks):7 buffer = ""8 for text in chunks:9 buffer += text10 cut = buffer.rfind(" ") # release only whole words11 if cut > 0:12 yield PHONE.sub("[phone]", buffer[:cut + 1])13 buffer = buffer[cut + 1:]14 yield PHONE.sub("[phone]", buffer)1516safe_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
from langchain_classic.chains.base import Chain # legacy; avoid for new codeA 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
- "
@chainorRunnableLambda?" — 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 toRunnableLambda(func, afunc=...). - "How do you give a custom step config, like callbacks?" — Add a second parameter named
configto the function; LangChain passes the run config in.