Building Custom Chains

While LangChain provides several useful pre-built chains like LLMChain and SequentialChain, you'll often encounter situations where your application requires a more tailored sequence of operations or specific logic connecting the components. This is where building custom chains becomes necessary. You might need to:

• Integrate custom data transformation steps between LLM calls.
• Implement conditional logic based on intermediate results.
• Combine outputs from multiple components in a specific way.
• Create chains with non-linear data flows.

The most flexible and recommended way to construct custom chains in modern LangChain is by using the LangChain Expression Language (LCEL). LCEL provides a declarative way to compose different components together using a syntax that mirrors the flow of data.

Composing with LangChain Expression Language (LCEL)

LCEL is built around the Runnable protocol. Most LangChain components, like LLM wrappers, prompt templates, and output parsers, implement this protocol. The fundamental operation in LCEL is piping (|) components together. When you pipe component A into component B (A | B), the output of A is automatically passed as the input to B.

Let's illustrate with a simple example building upon the concepts from Chapter 4. Suppose we want a chain that takes a topic, generates a technical question about it, and then uses another LLM call to provide a concise answer to that generated question.

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# Assume OPENAI_API_KEY is set in the environment
# Initialize models
model = ChatOpenAI(model="gpt-3.5-turbo")
qa_model = ChatOpenAI(model="gpt-4") # Using a potentially stronger model for QA

# 1. Prompt to generate a question
prompt_generate_question = ChatPromptTemplate.from_template(
    "Generate one concise technical question about the topic: {topic}"
)

# 2. Prompt to answer the generated question
prompt_answer_question = ChatPromptTemplate.from_template(
    "Provide a brief, clear answer to the following question:\n"
    "Question: {generated_question}\n"
    "Answer:"
)

# 3. Define the custom chain using LCEL
#    We need the original topic and the generated question for the second prompt.
#    RunnablePassthrough passes the input {'topic': '...'} along.
#    RunnableParallel executes branches and collects results in a dictionary.

chain = (
    {"generated_question": prompt_generate_question | model | StrOutputParser(),
     "topic": RunnablePassthrough()} # Pass the original topic through
    | prompt_answer_question
    | qa_model
    | StrOutputParser()
)

# Input for the chain
topic_input = {"topic": "vector databases"}

# Invoke the chain
result = chain.invoke(topic_input)

print(f"Topic: {topic_input['topic']}")
# Example Output (will vary):
# Generated Question (intermediate, not directly output by chain): What is the primary mechanism vector databases use for similarity search?
print(f"Generated Answer:\n{result}")
# Example Output (will vary):
# Generated Answer:
# The primary mechanism vector databases use for similarity search is Approximate Nearest Neighbor (ANN) algorithms. These algorithms efficiently find vectors in high-dimensional space that are closest (most similar) to a given query vector, without needing to compare the query against every single vector in the database.

In this example:

1. We define two separate prompt templates and LLM instances.
2. We use RunnableParallel (represented by the dictionary structure {...}) to run two branches:
   • The first branch (prompt_generate_question | model | StrOutputParser()) generates the question based on the input topic.
   • The second branch (RunnablePassthrough()) simply passes the original input dictionary (containing the topic) along.
3. The results of these parallel branches are collected into a dictionary: {'generated_question': '...', 'topic': '...'}.
4. This dictionary is then piped (|) into prompt_answer_question, which uses both keys.
5. The formatted prompt goes to the qa_model, and its output is parsed by StrOutputParser.

Data flow for the custom chain combining question generation and answering using LCEL.

Integrating Custom Python Functions

LCEL chains are not limited to LangChain components. You can integrate your own Python functions. LangChain automatically wraps standard Python functions passed within an LCEL sequence into RunnableLambda.

Let's modify the previous example slightly. Suppose after generating the question, we want to add a custom prefix using a Python function before passing it to the answering prompt.

# ... (imports and model definitions remain the same) ...

# Custom Python function
def add_question_prefix(question_text):
  print(f"--- Applying custom prefix ---") # For demonstration
  return f"Technical Query: {question_text}"

# Chain definition using the custom function
chain_with_custom_func = (
    {"intermediate_question": prompt_generate_question | model | StrOutputParser(),
     "topic": RunnablePassthrough()}
    | RunnableParallel(
        generated_question=lambda x: add_question_prefix(x["intermediate_question"]), # Apply function
        topic=lambda x: x["topic"] # Pass topic along
      )
    | prompt_answer_question # Expects 'generated_question' and 'topic'
    | qa_model
    | StrOutputParser()
)

# Invoke the modified chain
result_custom = chain_with_custom_func.invoke(topic_input)

print(f"Topic: {topic_input['topic']}")
# Example Output (will vary):
# --- Applying custom prefix ---
print(f"Generated Answer (with custom func):\n{result_custom}")
# Example Output (will vary):
# Generated Answer (with custom func):
# The primary mechanism vector databases use for similarity search is Approximate Nearest Neighbor (ANN) algorithms. These algorithms efficiently find vectors in high-dimensional space that are closest (most similar) to a given query vector, often using techniques like hashing or graph-based indexing, without exhaustively comparing against every entry.

Here, we used a lambda function within RunnableParallel to call our add_question_prefix function. LangChain handles wrapping this into a RunnableLambda behind the scenes, making the integration smooth.

Older Method: Subclassing Chain

Before LCEL became the standard, the primary way to create custom chains was by subclassing langchain.chains.base.Chain. This involves:

1. Defining input_keys (a list of expected input variable names).
2. Defining output_keys (a list of variable names the chain will output).
3. Implementing the _call method, which contains the core logic of the chain, taking inputs and returning a dictionary of outputs.

While still functional, subclassing is generally more verbose and less composable than using LCEL for most common workflow patterns. It might still be considered for chains requiring complex internal state management or logic that doesn't fit the sequential or parallel flow easily modeled by LCEL. However, for building most custom workflows, LCEL offers a more elegant and maintainable solution.

Building custom chains, particularly with LCEL, gives you fine-grained control over your LLM workflows, enabling you to construct sophisticated applications tailored precisely to your requirements. This composability is a significant advantage when moving past simple prompt-response interactions.