如何用 LangChain 自定义 Chat Model —— 测试工程师实践指南
在大模型(LLM)测试或集成时,标准API往往无法满足所有需求。比如你想做Mock、集成私有模型或特殊的业务逻辑,这时候就需要自定义 chat model。LangChain 为此提供了标准化的接口和抽象,简单易用、扩展性强。本篇将以通俗易懂的方式,带你一步步实现一个自定义 chat model,并讲解背后的设计原理和测试要点。
一、什么是 LangChain 的 Chat Model?
LangChain 把“聊天模型”抽象为一种消息输入、消息输出的接口。
常见消息类型有:
消息类型 | 说明 |
---|---|
SystemMessage | 系统消息,通常作为对话的第一条,用来设定AI行为 |
HumanMessage | 用户输入消息 |
AIMessage | AI 的回复,既可以是文本,也可以是函数调用的请求 |
FunctionMessage | 工具/函数调用结果 |
ToolMessage | 与 FunctionMessage 类似,随 OpenAI 新标准发展 |
…Chunk | 对应消息的“流式”输出切片,用于分步输出 |
流式消息Chunk:用于支持“边生成边输出”,比如大模型流式回复时每次只返回一部分token。
二、为什么要自定义 Chat Model?
- Mock/测试场景:开发阶段,用模型“鹦鹉学舌”或固定返回,方便回归和接口验证。
- 集成内部/私有模型:对接自研模型API,只需实现标准接口即可享受LangChain生态。
- 特殊业务逻辑:比如对输入做特殊过滤、日志追踪、A/B测试等。
三、如何自定义 Chat Model?(手把手教程)
1. 继承 BaseChatModel 并实现 _generate
方法
下面实现一个“鹦鹉模型”:只返回用户最后一句话的前n
个字符。
from typing import Any, Dict, List, Optional, Iterator
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from pydantic import Fieldclass ChatParrotLink(BaseChatModel):"""一个简单的自定义Chat Model:返回最后一句话的前n个字符。初始化参数:- parrot_buffer_length: int 要返回的字符数- model_name: str 模型名称"""model_name: str = Field(alias="model")parrot_buffer_length: intdef _generate(self,messages: List[BaseMessage],stop: Optional[List[str]] = None,run_manager=None,**kwargs: Any,) -> ChatResult:last_message = messages[-1]tokens = last_message.content[: self.parrot_buffer_length]message = AIMessage(content=tokens)return ChatResult(generations=[ChatGeneration(message=message)])@propertydef _llm_type(self) -> str:return "echoing-chat-model-advanced"
要点说明:
_generate
是核心——输入是历史消息列表(支持多轮),输出是 AI 的回复。- 必须有
_llm_type
属性,标记模型类型。
2. 如何实现流式输出(可选)
如果你想支持流式输出(比如边出字边推送),只需实现 _stream
方法:
def _stream(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager=None, **kwargs: Any,
) -> Iterator[ChatGenerationChunk]:last_message = messages[-1]tokens = last_message.content[: self.parrot_buffer_length]for t in tokens:yield ChatGenerationChunk(message=AIMessageChunk(content=t))yield ChatGenerationChunk(message=AIMessageChunk(content="")) # 结束
3. 使用示例
# 实例化模型
model = ChatParrotLink(parrot_buffer_length=3, model="my_custom_model")# 单轮对话
result = model.invoke([HumanMessage(content="hello!")])
print(result) # 输出: 'hel'# 多轮对话
result = model.invoke([HumanMessage(content="hello!"),AIMessage(content="Hi there!"),HumanMessage(content="Meow!"),
])
print(result) # 输出: 'Meo'# 流式输出
for chunk in model.stream("cat"):print(chunk.content, end="|") # 输出: c|a|t||
四、测试工程师关注的实践要点
- 输入输出类型清晰:输入为消息对象列表,输出为 AIMessage/AIMessageChunk。
- 可Mock、可回归:自定义模型可方便Mock,做接口回归测试。
- 支持批量、流式、异步:只要实现相应接口,LangChain自动支持批量/异步/流式测试。
- Stop Token支持:如需终止符,务必让模型在输出中包含终止符,便于后续解析和对齐。
- 参数可追踪:
_identifying_params
支持回溯测试用例用的参数,便于监控和测试。
五、进阶建议与注意事项
- 流式与异步建议结合使用:如需极致性能,考虑实现
_agenerate
(异步生成)和_astream
(异步流式)。 - 保密信息管理:如模型需API Key,建议用Pydantic的
SecretStr
避免日志泄露。 - 文档与测试覆盖:写好类和参数文档字符串,方便团队成员和自动化工具读取。
- 日志与监控:合理利用
run_manager
,可自动回调日志、统计token用量、链路追踪等。
六、下一步
- 你可以用同样方式包裹你自己的API、本地模型、甚至脚本Mock。
- 也可深入 LangChain 文档,学习如何让模型返回结构化数据、如何追踪token用量等。
总结:
自定义 Chat Model 提供了极大的灵活性,无论是接口测试、性能回归、Mock还是集成自研模型,LangChain 都给你准备好了标准化扩展点。测试工程师不需要深入AI模型细节,也能自信地实现并测试自己的Chat模型!
有问题欢迎留言交流,让测试更高效! 🚀