While LangChain provides a rich set of pre-built components for interacting with Large Language Models (LLMs), constructing prompts, and parsing outputs, production environments often present unique requirements that necessitate extending or replacing these standard implementations. The framework's modular design, built upon well-defined interfaces, allows developers to inject custom logic precisely where needed. This section focuses on customizing the core building blocks: LLM wrappers, prompt templates, and output parsers, enabling you to tailor LangChain's behavior to specific model interactions, complex input formatting, and non-standard output structures. Mastering these customizations is significant for building specialized and efficient LLM applications.
The standard LLM and ChatModel wrappers in LangChain handle communication with various model providers (OpenAI, Anthropic, Cohere, Hugging Face models, etc.). However, you might need a custom wrapper for several reasons:
To create a custom LLM wrapper, you typically inherit from langchain_core.language_models.llms.LLM
and implement the _call
method for synchronous execution and _acall
for asynchronous execution. For chat models, inherit from langchain_core.language_models.chat_models.BaseChatModel
and implement _generate
and _agenerate
.
A crucial method is _generate
, which takes a list of messages and optional stop
sequences, processes them, and returns a ChatResult
containing ChatGeneration
objects. You also need to implement the _llm_type
property to identify your custom model.
Let's illustrate with a conceptual example of a custom chat model wrapper that adds a simple prefix to every user message and logs the interaction duration.
import time
from typing import Any, List, Optional
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage, AIMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatResult
# Assume we have an existing chat model instance we want to wrap
# from langchain_openai import ChatOpenAI
# underlying_chat_model = ChatOpenAI()
class CustomLoggedChatWrapper(BaseChatModel):
"""
A conceptual wrapper that adds a prefix to user messages
and logs interaction time.
"""
underlying_model: BaseChatModel # The actual model to call
@property
def _llm_type(self) -> str:
return "custom_logged_chat_wrapper"
def _generate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
start_time = time.time()
# Modify messages (example: add prefix to HumanMessage)
processed_messages = []
for msg in messages:
if isinstance(msg, HumanMessage):
processed_messages.append(
HumanMessage(content=f"[User Inquiry] {msg.content}", additional_kwargs=msg.additional_kwargs)
)
else:
processed_messages.append(msg)
# Call the underlying model
# In a real scenario, you'd pass callbacks and kwargs appropriately
result = self.underlying_model._generate(processed_messages, stop=stop, run_manager=run_manager, **kwargs)
end_time = time.time()
duration = end_time - start_time
print(f"Custom Wrapper: Interaction took {duration:.2f} seconds.")
# Potentially modify the result before returning
# For example, logging token usage if available in result.llm_output
return result
# Implement _agenerate for async support similarly
async def _agenerate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
# Async implementation would use underlying_model._agenerate
# and asynchronous time tracking / logging
# (Implementation omitted for brevity)
pass
# Usage Example (conceptual, assuming underlying_chat_model is defined)
# custom_llm = CustomLoggedChatWrapper(underlying_model=underlying_chat_model)
# response = custom_llm.invoke([HumanMessage(content="Tell me about LangChain.")])
# print(response)
This example demonstrates wrapping an existing model, modifying inputs (processed_messages
), calling the underlying model's generation method, and adding custom logic (timing). A wrapper for an entirely unsupported model would involve making direct API calls within _generate
and constructing the ChatResult
manually. Remember to implement the asynchronous counterparts (_agenerate
, _acall
) for compatibility with asynchronous LangChain features like ainvoke
.
LangChain's PromptTemplate
and ChatPromptTemplate
offer flexibility for formatting inputs, but sometimes you need more intricate logic:
You can create custom prompt templates by inheriting from langchain_core.prompts.BasePromptTemplate
or langchain_core.prompts.BaseChatPromptTemplate
. The key method to implement is format_prompt
(which should return a PromptValue
) or simply format
(returning a formatted string) for simpler templates. For chat templates, you'll implement format_messages
.
Consider a scenario where we want a chat prompt template that adds a system message outlining the expected output format only if the user's query suggests a request for structured data (e.g., contains the word "list" or "table").
from typing import Any, Dict, List
from langchain_core.prompts import BaseChatPromptTemplate
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_core.pydantic_v1 import BaseModel, Field
class DynamicStructureChatPrompt(BaseChatPromptTemplate, BaseModel):
"""
A chat prompt template that conditionally adds a system message
about structured output.
"""
# Define input variables if needed, using Pydantic Fields
input_variables: List[str] = Field(default=["user_query"])
def format_messages(self, **kwargs: Any) -> List[BaseMessage]:
user_query = kwargs["user_query"]
messages: List[BaseMessage] = []
# Conditional logic for system message
if "list" in user_query.lower() or "table" in user_query.lower():
messages.append(
SystemMessage(content="Please provide the output as a structured list or table.")
)
else:
messages.append(
SystemMessage(content="You are a helpful assistant.")
)
messages.append(HumanMessage(content=user_query))
# Potentially add more messages based on other kwargs or logic
return messages
def _prompt_type(self) -> str:
return "dynamic-structure-chat-prompt"
# Usage Example:
dynamic_prompt = DynamicStructureChatPrompt()
query1 = "Summarize the benefits of LCEL."
messages1 = dynamic_prompt.format_messages(user_query=query1)
# messages1 would be [SystemMessage(...helpful assistant...), HumanMessage(...Summarize...)]
query2 = "Give me a list of vector stores supported by LangChain."
messages2 = dynamic_prompt.format_messages(user_query=query2)
# messages2 would be [SystemMessage(...structured list/table...), HumanMessage(...Give me a list...)]
# print(messages1)
# print(messages2)
This DynamicStructureChatPrompt
inspects the user_query
input variable to decide which system message to include. This allows for adaptive prompting based on the nature of the input, potentially guiding the LLM more effectively.
Output parsers are responsible for transforming the raw string or message output from an LLM into a more structured format (like a dictionary, a Pydantic object, or a custom domain object). While LangChain offers parsers for common formats (JSON, CSV, Pydantic models, structured function/tool calls), custom parsers are needed when:
To create a custom output parser, inherit from langchain_core.output_parsers.BaseOutputParser
. The primary method to implement is parse
, which takes the LLM's output string (or Generation
object via parse_result
) and returns the desired structured data. You might also implement get_format_instructions
to provide guidance to the LLM on how it should format its output.
Let's imagine an LLM sometimes returns a person's name and age, but the format varies (e.g., "Name: Alice, Age: 30", "Bob is 25 years old", "Age: 40 Name: Carol"). We can create a parser using regular expressions.
import re
from typing import Dict, Any
from langchain_core.output_parsers import BaseOutputParser
from langchain_core.exceptions import OutputParserException
class RegexNameAgeParser(BaseOutputParser[Dict[str, Any]]):
"""
Parses text to extract Name and Age using regular expressions,
handling varied formats.
"""
def parse(self, text: str) -> Dict[str, Any]:
"""Parses the output text to extract name and age."""
name_match = re.search(r"(?:Name:\s*|)(\b[A-Z][a-z]+)\b", text, re.IGNORECASE)
age_match = re.search(r"(?:Age:\s*|is\s*|)(\d+)\s*(?:years old|)", text, re.IGNORECASE)
if name_match and age_match:
name = name_match.group(1)
age = int(age_match.group(1))
return {"name": name, "age": age}
else:
# Handle cases where parsing fails
# Option 1: Raise an exception
raise OutputParserException(
f"Could not parse Name and Age from output: {text}"
)
# Option 2: Return a default/error structure
# return {"name": None, "age": None, "error": "Parsing failed"}
def get_format_instructions(self) -> str:
"""Instructions for the LLM on the desired format."""
# While this parser is robust, providing instructions can still help.
return "Please provide the name and age. For example: 'Name: John, Age: 35'"
@property
def _type(self) -> str:
return "regex_name_age_parser"
# Usage Example:
parser = RegexNameAgeParser()
output1 = "The user is Alice, Age: 30."
output2 = "Bob is 25 years old."
output3 = "Age: 40 Name: Carol"
output4 = "David's information is not available."
try:
parsed1 = parser.parse(output1) # {"name": "Alice", "age": 30}
parsed2 = parser.parse(output2) # {"name": "Bob", "age": 25}
parsed3 = parser.parse(output3) # {"name": "Carol", "age": 40}
# print(parsed1, parsed2, parsed3)
parsed4 = parser.parse(output4)
except OutputParserException as e:
print(e) # Prints the exception message for output4
# print(parser.get_format_instructions())
This parser uses re.search
to find patterns indicative of names and ages, demonstrating robustness to variations in the LLM's output. It also includes error handling via OutputParserException
and provides format instructions.
The beauty of LangChain's design, particularly with the LangChain Expression Language (LCEL), is that custom components conforming to the base interfaces integrate seamlessly with standard ones. You can pipe them together just like built-in components.
# Assume custom_prompt, custom_llm, and custom_parser are instantiated
# custom_prompt: An instance of DynamicStructureChatPrompt
# custom_llm: An instance of CustomLoggedChatWrapper (wrapping a real model)
# custom_parser: An instance of RegexNameAgeParser
# You can chain custom components together:
# chain = custom_prompt | custom_llm | custom_parser
# Or mix standard and custom components:
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
# standard_llm = ChatOpenAI(model="gpt-3.5-turbo")
# standard_parser = StrOutputParser()
# Example mix: Custom prompt, standard LLM, custom parser
# mixed_chain_1 = custom_prompt | standard_llm | custom_parser
# Example mix: Standard prompt, custom LLM, standard parser
# standard_prompt = ChatPromptTemplate.from_messages([
# ("system", "You are a helpful assistant."),
# ("human", "{question}")
# ])
# mixed_chain_2 = standard_prompt | custom_llm | standard_parser
# result = mixed_chain_1.invoke({"user_query": "Provide details for Eve, Age: 22."})
# print(result) # Might output: {'name': 'Eve', 'age': 22}
This composability allows you to precisely target which parts of your LLM interaction logic need customization without rewriting the entire chain. Start with standard components and introduce custom ones incrementally as specific requirements arise. Remember that thorough testing is essential, especially for custom output parsers dealing with potentially unpredictable LLM outputs. By leveraging custom components, you gain the fine-grained control necessary to build sophisticated, reliable, and production-ready LangChain applications.
© 2025 ApX Machine Learning