自定义聊天模型
在本指南中，我们将学习如何使用 LangChain 抽象创建一个自定义聊天模型。
用标准的 BaseChatModel 接口包装你的 LLM，能让你以最小的代码改动在现有的 LangChain 程序中使用你的 LLM！
额外的好处是，你的 LLM 会自动成为一个 LangChain Runnable，并能开箱即用地受益于一些优化（例如通过线程池实现的批处理）、异步支持、astream_events API 等。
输入与输出
首先，我们需要谈谈 消息（messages），它们是聊天模型的输入和输出。
消息
聊天模型以消息作为输入，并返回一条消息作为输出。
LangChain 有一些内置的消息类型：
| 消息类型 | 描述 |
|---|---|
SystemMessage | 用于预设 AI 行为，通常作为输入消息序列的第一条传入。 |
HumanMessage | 表示来自与聊天模型交互的人的消息。 |
AIMessage | 表示来自聊天模型的消息。它可以是文本，也可以是调用工具的请求。 |
FunctionMessage / ToolMessage | 用于把工具调用结果传回模型的消息。 |
AIMessageChunk / HumanMessageChunk / ... | 每种消息类型的块（chunk）变体。 |
注意：
ToolMessage和FunctionMessage紧密对应 OpenAI 的function和tool角色。这是一个快速发展的领域，随着越来越多的模型加入函数调用能力，预计这个 schema 会有新增内容。
from langchain_core.messages import (
AIMessage,
BaseMessage,
FunctionMessage,
HumanMessage,
SystemMessage,
ToolMessage,
)流式变体
所有聊天消息都有一个名称中包含 Chunk 的流式变体。
from langchain_core.messages import (
AIMessageChunk,
FunctionMessageChunk,
HumanMessageChunk,
SystemMessageChunk,
ToolMessageChunk,
)当从聊天模型流式输出时会用到这些块，它们都定义了一个可相加的属性！
AIMessageChunk(content="Hello") + AIMessageChunk(content=" World!")AIMessageChunk(content='Hello World!')基础聊天模型
让我们实现一个聊天模型，它会回显提示中最后一条消息的前 n 个字符！
为此，我们将继承 BaseChatModel，并且需要实现以下内容：
| 方法/属性 | 描述 | 必需/可选 |
|---|---|---|
_generate | 用于从提示生成聊天结果 | 必需 |
_llm_type（属性） | 用于唯一标识模型类型。用于日志记录。 | 必需 |
_identifying_params（属性） | 表示用于追踪目的的模型参数化。 | 可选 |
_stream | 用于实现流式输出。 | 可选 |
_agenerate | 用于实现原生异步方法。 | 可选 |
_astream | 用于实现 _stream 的异步版本。 | 可选 |
提示：
_astream的实现会在_stream已实现时使用run_in_executor在单独线程中启动同步的_stream，否则回退使用_agenerate。如果你想复用_stream的实现，可以使用这个技巧，但如果你能实现原生异步的代码，那是更好的方案，因为那样的代码运行开销更小。
实现
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
from langchain_core.callbacks import (
AsyncCallbackManagerForLLMRun,
CallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel, SimpleChatModel
from langchain_core.messages import AIMessageChunk, BaseMessage, HumanMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_core.runnables import run_in_executor
class CustomChatModelAdvanced(BaseChatModel):
"""A custom chat model that echoes the first `n` characters of the input.
When contributing an implementation to LangChain, carefully document
the model including the initialization parameters, include
an example of how to initialize the model and include any relevant
links to the underlying models documentation or API.
Example:
.. code-block:: python
model = CustomChatModel(n=2)
result = model.invoke([HumanMessage(content="hello")])
result = model.batch([[HumanMessage(content="hello")],
[HumanMessage(content="world")]])
"""
model_name: str
"""The name of the model"""
n: int
"""The number of characters from the last message of the prompt to be echoed."""
def _generate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
"""Override the _generate method to implement the chat model logic.
This can be a call to an API, a call to a local model, or any other
implementation that generates a response to the input prompt.
Args:
messages: the prompt composed of a list of messages.
stop: a list of strings on which the model should stop generating.
If generation stops due to a stop token, the stop token itself
SHOULD BE INCLUDED as part of the output. This is not enforced
across models right now, but it's a good practice to follow since
it makes it much easier to parse the output of the model
downstream and understand why generation stopped.
run_manager: A run manager with callbacks for the LLM.
"""
# Replace this with actual logic to generate a response from a list
# of messages.
last_message = messages[-1]
tokens = last_message.content[: self.n]
message = AIMessage(
content=tokens,
additional_kwargs={}, # Used to add additional payload (e.g., function calling request)
response_metadata={ # Use for response metadata
"time_in_seconds": 3,
},
)
##
generation = ChatGeneration(message=message)
return ChatResult(generations=[generation])
def _stream(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
"""Stream the output of the model.
This method should be implemented if the model can generate output
in a streaming fashion. If the model does not support streaming,
do not implement it. In that case streaming requests will be automatically
handled by the _generate method.
Args:
messages: the prompt composed of a list of messages.
stop: a list of strings on which the model should stop generating.
If generation stops due to a stop token, the stop token itself
SHOULD BE INCLUDED as part of the output. This is not enforced
across models right now, but it's a good practice to follow since
it makes it much easier to parse the output of the model
downstream and understand why generation stopped.
run_manager: A run manager with callbacks for the LLM.
"""
last_message = messages[-1]
tokens = last_message.content[: self.n]
for token in tokens:
chunk = ChatGenerationChunk(message=AIMessageChunk(content=token))
if run_manager:
# This is optional in newer versions of LangChain
# The on_llm_new_token will be called automatically
run_manager.on_llm_new_token(token, chunk=chunk)
yield chunk
# Let's add some other information (e.g., response metadata)
chunk = ChatGenerationChunk(
message=AIMessageChunk(content="", response_metadata={"time_in_sec": 3})
)
if run_manager:
# This is optional in newer versions of LangChain
# The on_llm_new_token will be called automatically
run_manager.on_llm_new_token(token, chunk=chunk)
yield chunk
@property
def _llm_type(self) -> str:
"""Get the type of language model used by this chat model."""
return "echoing-chat-model-advanced"
@property
def _identifying_params(self) -> Dict[str, Any]:
"""Return a dictionary of identifying parameters.
This information is used by the LangChain callback system, which
is used for tracing purposes make it possible to monitor LLMs.
"""
return {
# The model name allows users to specify custom token counting
# rules in LLM monitoring applications (e.g., in LangSmith users
# can provide per token pricing for their model and monitor
# costs for the given LLM.)
"model_name": self.model_name,
}让我们测试一下 🧪
聊天模型将实现 LangChain 标准的 Runnable 接口，许多 LangChain 抽象都支持它！
model = CustomChatModelAdvanced(n=3, model_name="my_custom_model")model.invoke(
[
HumanMessage(content="hello!"),
AIMessage(content="Hi there human!"),
HumanMessage(content="Meow!"),
]
)AIMessage(content='Meo', response_metadata={'time_in_seconds': 3}, id='run-ddb42bd6-4fdd-4bd2-8be5-e11b67d3ac29-0')model.invoke("hello")AIMessage(content='hel', response_metadata={'time_in_seconds': 3}, id='run-4d3cc912-44aa-454b-977b-ca02be06c12e-0')model.batch(["hello", "goodbye"])[AIMessage(content='hel', response_metadata={'time_in_seconds': 3}, id='run-9620e228-1912-4582-8aa1-176813afec49-0'),
AIMessage(content='goo', response_metadata={'time_in_seconds': 3}, id='run-1ce8cdf8-6f75-448e-82f7-1bb4a121df93-0')]for chunk in model.stream("cat"):
print(chunk.content, end="|")c|a|t||请查看模型中 _astream 的实现！如果你不实现它，那么将不会有任何输出被流式传输！
async for chunk in model.astream("cat"):
print(chunk.content, end="|")c|a|t||让我们尝试使用 astream events API，它也有助于复核所有回调是否都已正确实现！
async for event in model.astream_events("cat", version="v1"):
print(event){'event': 'on_chat_model_start', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'name': 'CustomChatModelAdvanced', 'tags': [], 'metadata': {}, 'data': {'input': 'cat'}}
{'event': 'on_chat_model_stream', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='c', id='run-125a2a16-b9cd-40de-aa08-8aa9180b07d0')}}
{'event': 'on_chat_model_stream', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='a', id='run-125a2a16-b9cd-40de-aa08-8aa9180b07d0')}}
{'event': 'on_chat_model_stream', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='t', id='run-125a2a16-b9cd-40de-aa08-8aa9180b07d0')}}
{'event': 'on_chat_model_stream', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'tags': [], 'metadata': {}, 'name': 'CustomChatModelAdvanced', 'data': {'chunk': AIMessageChunk(content='', response_metadata={'time_in_sec': 3}, id='run-125a2a16-b9cd-40de-aa08-8aa9180b07d0')}}
{'event': 'on_chat_model_end', 'name': 'CustomChatModelAdvanced', 'run_id': '125a2a16-b9cd-40de-aa08-8aa9180b07d0', 'tags': [], 'metadata': {}, 'data': {'output': AIMessageChunk(content='cat', response_metadata={'time_in_sec': 3}, id='run-125a2a16-b9cd-40de-aa08-8aa9180b07d0')}}
/home/eugene/src/langchain/libs/core/langchain_core/_api/beta_decorator.py:87: LangChainBetaWarning: This API is in beta and may change in the future.
warn_beta(
贡献
我们感谢所有聊天模型集成的贡献。
下面是一份检查清单，帮助你确保你的贡献被加入 LangChain：
文档：
- 模型为所有初始化参数都包含 doc-string，因为它们会展示在 APIReference (opens in a new tab) 中。
- 如果模型由某项服务驱动，模型的类 doc-string 要包含指向模型 API 的链接。
测试：
- 为被重写的方法添加单元测试或集成测试。如果你重写了相应代码，请验证
invoke、ainvoke、batch、stream能正常工作。
流式输出（如果你实现它）：
- 实现 _stream 方法以使流式输出工作
停止令牌行为：
- 应尊重停止令牌
- 停止令牌应作为响应的一部分被包含
密钥 API Keys：
- 如果你的模型连接到某个 API，它很可能会在初始化时接受 API 密钥。对密钥使用 Pydantic 的
SecretStr类型，这样当人们打印模型时它们不会被意外打印出来。
标识参数：
- 在标识参数中包含
model_name
优化：
考虑提供原生异步支持，以减少模型带来的开销！
- 提供
_agenerate的原生异步版本（被ainvoke使用） - 提供
_astream的原生异步版本（被astream使用）