#附一 基于 LangChain 自定义 LLM
LangChain 为基于 LLM 开发自定义应用提供了高效的开发框架，便于开发者迅速地激发 LLM 的强大能力，搭建 LLM 应用。LangChain 也同样支持多种大模型，内置了 OpenAI、LLAMA 等大模型的调用接口。但是，LangChain 并没有内置所有大模型，它通过允许用户自定义 LLM 类型，来提供强大的可扩展性。
在本部分，我们以智谱为例，讲述如何基于 LangChain 自定义 LLM，让我们基于 LangChain 搭建的应用能够支持国内平台。
本部分涉及相对更多 LangChain、大模型调用的技术细节，有精力同学可以学习部署，如无精力可以直接使用后续代码来支持调用。
要实现自定义 LLM，需要定义一个自定义类继承自 LangChain 的 LLM 基类，然后定义两个函数：
① _generate 方法，接收一系列消息及其他参数，并返回输出；
② _stream 方法， 接收一系列消息及其他参数，并以流式返回结果。
首先我们导入所需的第三方库：
from typing import Any, Dict, Iterator, List, Optional
from zhipuai import ZhipuAI
from langchain_core.callbacks import (
CallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import (
AIMessage,
AIMessageChunk,
BaseMessage,
SystemMessage,
ChatMessage,
HumanMessage
)
from langchain_core.messages.ai import UsageMetadata
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
import time由于LangChain的消息类型是HumanMessage、AIMessage等格式，与一般模型接收的字典格式不一样，因此我们需要先定义一个将LangChain的格式转为字典的函数。
def _convert_message_to_dict(message: BaseMessage) -> dict:
"""把LangChain的消息格式转为智谱支持的格式
Args:
message: The LangChain message.
Returns:
The dictionary.
"""
message_dict: Dict[str, Any] = {"content": message.content}
if (name := message.name or message.additional_kwargs.get("name")) is not None:
message_dict["name"] = name
# populate role and additional message data
if isinstance(message, ChatMessage):
message_dict["role"] = message.role
elif isinstance(message, HumanMessage):
message_dict["role"] = "user"
elif isinstance(message, AIMessage):
message_dict["role"] = "assistant"
elif isinstance(message, SystemMessage):
message_dict["role"] = "system"
else:
raise TypeError(f"Got unknown type {message}")
return message_dict接着我们定义一个继承自 LLM 类的自定义 LLM 类：
# 继承自 LangChain 的 BaseChatModel 类
class ZhipuaiLLM(BaseChatModel):
"""自定义Zhipuai聊天模型。
"""
model_name: str = None
temperature: Optional[float] = None
max_tokens: Optional[int] = None
timeout: Optional[int] = None
stop: Optional[List[str]] = None
max_retries: int = 3
api_key: str | None = None上述初始化涵盖了我们平时常用的参数，也可以根据实际需求与智谱的 API 加入更多的参数。
接下来我们实现_generate方法：
def _generate(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> ChatResult:
"""通过调用智谱API从而响应输入。
Args:
messages: 由messages列表组成的prompt
stop: 在模型生成的回答中有该字符串列表中的元素则停止响应
run_manager: 一个为LLM提供回调的运行管理器
"""
# 列表推导式 将 messages 的元素逐个转为智谱的格式
messages = [_convert_message_to_dict(message) for message in messages]
# 定义推理的开始时间
start_time = time.time()
# 调用 ZhipuAI 对处理消息
response = ZhipuAI(api_key=self.api_key).chat.completions.create(
model=self.model_name,
temperature=self.temperature,
max_tokens=self.max_tokens,
timeout=self.timeout,
stop=stop,
messages=messages
)
# 计算运行时间 由现在时间 time.time() 减去 开始时间start_time得到
time_in_seconds = time.time() - start_time
# 将返回的消息封装并返回
message = AIMessage(
content=response.choices[0].message.content, # 响应的结果
additional_kwargs={}, # 额外信息
response_metadata={
"time_in_seconds": round(time_in_seconds, 3), # 响应源数据 这里是运行时间 也可以添加其他信息
},
# 本次推理消耗的token
usage_metadata={
"input_tokens": response.usage.prompt_tokens, # 输入token
"output_tokens": response.usage.completion_tokens, # 输出token
"total_tokens": response.usage.total_tokens, # 全部token
},
)
generation = ChatGeneration(message=message)
return ChatResult(generations=[generation])接下来我们实现另一核心的方法_stream，之前注释的代码本次不再注释：
def _stream(
self,
messages: List[BaseMessage],
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Iterator[ChatGenerationChunk]:
"""通过调用智谱API返回流式输出。
Args:
messages: 由messages列表组成的prompt
stop: 在模型生成的回答中有该字符串列表中的元素则停止响应
run_manager: 一个为LLM提供回调的运行管理器
"""
messages = [_convert_message_to_dict(message) for message in messages]
response = ZhipuAI().chat.completions.create(
model=self.model_name,
stream=True, # 将stream 设置为 True 返回的是迭代器，可以通过for循环取值
temperature=self.temperature,
max_tokens=self.max_tokens,
timeout=self.timeout,
stop=stop,
messages=messages
)
start_time = time.time()
# 使用for循环存取结果
for res in response:
if res.usage: # 如果 res.usage 存在则存储token使用情况
usage_metadata = UsageMetadata(
{
"input_tokens": res.usage.prompt_tokens,
"output_tokens": res.usage.completion_tokens,
"total_tokens": res.usage.total_tokens,
}
)
# 封装每次返回的chunk
chunk = ChatGenerationChunk(
message=AIMessageChunk(content=res.choices[0].delta.content)
)
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(res.choices[0].delta.content, chunk=chunk)
# 使用yield返回 结果是一个生成器 同样可以使用for循环调用
yield chunk
time_in_sec = time.time() - start_time
# Let's add some other information (e.g., response metadata)
# 最终返回运行时间
chunk = ChatGenerationChunk(
message=AIMessageChunk(content="", response_metadata={"time_in_sec": round(time_in_sec, 3)}, usage_metadata=usage_metadata)
)
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("", chunk=chunk)
yield chunk然后我们还需要定义一下模型的描述方法：
@property
def _llm_type(self) -> str:
"""获取此聊天模型使用的语言模型类型。"""
return self.model_name
@property
def _identifying_params(self) -> Dict[str, Any]:
"""返回一个标识参数的字典。
该信息由LangChain回调系统使用，用于跟踪目的，使监视llm成为可能。
"""
return {
"model_name": self.model_name,
}通过上述步骤，我们就可以基于 LangChain 定义智谱的调用方式了。我们将此代码封装在 zhipu_llm.py 文件中。