如何用 LangChain 自定义 Chat Model —— 测试工程师实践指南-CSDN博客

在大模型（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 Field

class ChatParrotLink(BaseChatModel):
    """
    一个简单的自定义Chat Model：返回最后一句话的前n个字符。
    初始化参数：
      - parrot_buffer_length: int 要返回的字符数
      - model_name: str 模型名称
    """

    model_name: str = Field(alias="model")
    parrot_buffer_length: int

    def _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)])

    @property
    def _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||