LangChain 自定义 LLM Wrapper封装自部署模型为标准接口的工程实践一、深度引言与场景痛点大家好我是赵咕咕。三月份的时候老板下了死命令所有涉及客户数据的场景不允许调用外部 API。这意味着我们之前基于 GPT-4 搭建的 Agent 系统需要全部切到自部署的模型上。公司内部跑着一套 vLLM 集群Host 着 Qwen2-72B 和 Llama-3-70B。问题来了LangChain 原生支持 OpenAI API 格式但我们的 vLLM 暴露的是兼容 OpenAI 风格的 HTTP 接口参数名略有不同比如top_kvstop_p而且错误处理机制完全不一样。更麻烦的是——我们的 Agent 代码里ChatOpenAI已经嵌入了 37 个地方。重写 37 处不可能。最好的方案是写一个符合 LangChain BaseLLM 接口的自定义 Wrapper让所有现有代码零改动切换。二、底层机制与原理深度剖析2.1 LangChain 的 LLM 抽象层LangChain 的 LLM 体系核心是BaseLLM纯文本补全和BaseChatModel聊天补全。所有 LLM 实现都继承这两个基类只要实现了它们要求的方法你就能无缝接入 LangChain 的 Chain、Agent、Tool 等所有上层抽象。2.2 BaseChatModel 的继承契约继承BaseChatModel需要实现三个核心方法_generate()接收 messages返回 LLM 生成结果。_stream()可选流式输出。_llm_type返回字符串标识用于日志和追踪。关键是_generate()的输入输出类型输入是list[BaseMessage]输出是ChatResult。只要你的自部署模型能接收并返回符合这个类型的数据LangChain 上层完全感知不到底层换了一个模型。2.3 自定义 Wrapper 的调用时序调用时序的核心设计思路格式转换在 Wrapper 内部完成外部传入的BaseMessage自动转为 vLLM 需要的 OpenAI Chat 格式响应也自动转回ChatResult。连接池复用使用httpx.AsyncClient的连接池避免每次请求重新建立 TCP 连接。指数退避重试2s → 4s → 8s最多 3 次重试。降级链路vLLM 集群完全不可用时自动切换到本机 Ollama 运行的小模型保住基础可用性。三、生产级代码实现import asyncio import logging import time from typing import Any, Iterator import httpx from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import ( AIMessage, BaseMessage, HumanMessage, SystemMessage, ) from langchain_core.outputs import ChatGeneration, ChatResult from pydantic import Field, PrivateAttr logger logging.getLogger(__name__) class CustomVLLMWrapper(BaseChatModel): 自部署 vLLM 模型的自定义 Wrapper。 继承 BaseChatModel实现 LangChain 标准接口 封装 vLLM OpenAI-compatible API 的调用细节。 model_name: str Field(defaultQwen2-72B-Instruct) api_base: str Field(defaulthttp://localhost:8000/v1) temperature: float Field(default0.0, ge0.0, le2.0) max_tokens: int Field(default4096, gt0) max_retries: int Field(default3, ge0) request_timeout: float Field(default120.0, gt0) # 降级配置 fallback_enabled: bool Field(defaultTrue) fallback_api_base: str Field(defaulthttp://localhost:11434) # 内部状态 _client: httpx.AsyncClient PrivateAttr() _fallback_client: httpx.AsyncClient | None PrivateAttr(defaultNone) def __init__(self, **data: Any): super().__init__(**data) self._init_clients() def _init_clients(self) - None: 初始化 HTTP 客户端含连接池。 limits httpx.Limits( max_keepalive_connections50, max_connections200, keepalive_expiry30.0, ) self._client httpx.AsyncClient( base_urlself.api_base, timeouthttpx.Timeout(self.request_timeout), limitslimits, ) if self.fallback_enabled: self._fallback_client httpx.AsyncClient( base_urlself.fallback_api_base, timeouthttpx.Timeout(60.0), limitslimits, ) property def _llm_type(self) - str: return fcustom_vllm::{self.model_name} property def _identifying_params(self) - dict[str, Any]: return { model_name: self.model_name, api_base: self.api_base, temperature: self.temperature, max_tokens: self.max_tokens, } def _generate( self, messages: list[BaseMessage], stop: list[str] | None None, run_manager: CallbackManagerForLLMRun | None None, **kwargs: Any, ) - ChatResult: 同步入口包装异步实现。 return asyncio.run(self._agenerate(messages, stop, run_manager, **kwargs)) async def _agenerate( self, messages: list[BaseMessage], stop: list[str] | None None, run_manager: CallbackManagerForLLMRun | None None, **kwargs: Any, ) - ChatResult: 异步生成——核心逻辑。 payload self._build_payload(messages, stop, **kwargs) # 尝试主模型带重试 for attempt in range(self.max_retries 1): try: response await self._client.post( /chat/completions, jsonpayload ) response.raise_for_status() data response.json() return self._parse_response(data) except httpx.HTTPStatusError as e: if e.response.status_code 429: # 限流 wait min(2 ** (attempt 1), 30) logger.warning( vLLM 限流 (429)等待 %ds 后重试 (%d/%d), wait, attempt 1, self.max_retries, ) await asyncio.sleep(wait) continue elif e.response.status_code 500: # 服务端错误 if attempt self.max_retries: wait min(2 ** (attempt 1), 16) logger.warning( vLLM 服务端错误 (%d)%ds 后重试 (%d/%d), e.response.status_code, wait, attempt 1, self.max_retries, ) await asyncio.sleep(wait) continue raise else: # 4xx 客户端错误不重试 logger.error(vLLM 请求错误: %s, e) raise except (httpx.TimeoutException, httpx.ConnectError) as e: if attempt self.max_retries: wait min(2 ** (attempt 1), 16) logger.warning( vLLM 连接超时%ds 后重试 (%d/%d), wait, attempt 1, self.max_retries, ) await asyncio.sleep(wait) continue raise RuntimeError(fvLLM 连接失败超过 {self.max_retries} 次重试) from e # 全部重试耗尽 → 降级 if self.fallback_enabled and self._fallback_client: logger.warning(vLLM 不可用降级到 Ollama 本地模型) return await self._fallback_generate(messages, stop) raise RuntimeError(vLLM 不可用且降级未启用) async def _fallback_generate( self, messages: list[BaseMessage], stop: list[str] | None None, ) - ChatResult: 降级链路使用本地 Ollama 模型。 try: prompt \n.join( f{self._role_name(m)}: {m.content} for m in messages ) response await self._fallback_client.post( # type: ignore[union-attr] /api/generate, json{ model: llama3:8b, prompt: prompt, stream: False, }, timeout30.0, ) response.raise_for_status() data response.json() return ChatResult( generations[ChatGeneration( messageAIMessage(contentdata.get(response, )) )], llm_output{ model: fallback::llama3:8b, is_fallback: True, }, ) except Exception as e: logger.error(降级模型也失败了: %s, e) raise RuntimeError(主模型和降级模型均不可用) from e def _build_payload( self, messages: list[BaseMessage], stop: list[str] | None, **kwargs: Any, ) - dict[str, Any]: 构建 OpenAI-compatible 请求体。 chat_messages [] for msg in messages: role self._role_name(msg) content str(msg.content) if msg.content else chat_messages.append({role: role, content: content}) payload: dict[str, Any] { model: self.model_name, messages: chat_messages, temperature: kwargs.get(temperature, self.temperature), max_tokens: kwargs.get(max_tokens, self.max_tokens), } if stop: payload[stop] stop return payload staticmethod def _role_name(msg: BaseMessage) - str: BaseMessage → OpenAI role 映射。 if isinstance(msg, SystemMessage): return system elif isinstance(msg, HumanMessage): return user elif isinstance(msg, AIMessage): return assistant return user staticmethod def _parse_response(data: dict[str, Any]) - ChatResult: 解析 vLLM 响应 → ChatResult。 choices data.get(choices, []) if not choices: logger.warning(vLLM 返回空 choices: %s, data) return ChatResult( generations[ChatGeneration(messageAIMessage(content))] ) choice choices[0] content choice.get(message, {}).get(content, ) finish_reason choice.get(finish_reason, stop) return ChatResult( generations[ChatGeneration( messageAIMessage(contentcontent), generation_info{finish_reason: finish_reason}, )], llm_output{ model: data.get(model, unknown), usage: data.get(usage, {}), is_fallback: False, }, ) async def aclose(self) - None: 关闭 HTTP 客户端连接。 await self._client.aclose() if self._fallback_client: await self._fallback_client.aclose() # ─── Agent 集成示例 ─── async def main(): from langchain.agents import AgentExecutor, create_openai_tools_agent from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder llm CustomVLLMWrapper( model_nameQwen2-72B-Instruct, api_basehttp://vllm-cluster.internal:8000/v1, temperature0.0, max_retries3, fallback_enabledTrue, ) prompt ChatPromptTemplate.from_messages([ (system, 你是一个智能助手。), (human, {input}), MessagesPlaceholder(variable_nameagent_scratchpad), ]) agent create_openai_tools_agent(llm, [], prompt) executor AgentExecutor(agentagent, tools[], verboseFalse) try: result await executor.ainvoke({input: 解释什么是 RAG}) print(result[output]) finally: await llm.aclose() if __name__ __main__: asyncio.run(main())代码中值得注意的几个点连接池复用httpx.AsyncClient的max_keepalive_connections配置了 50 个长连接池避免高并发下频繁 TCP 握手。分级重试429限流和 5xx服务端错误走指数退避重试4xx客户端错误如 400 Bad Request直接抛出不浪费重试次数。降级透明降级结果通过llm_output[is_fallback]标记上层可以感知到这次是降级回答质量可能下降。BaseMessage 自动映射_role_name()方法把 LangChain 的消息类型自动映射到 OpenAI 的 role 格式上游完全不需要关心。四、边界分析与架构权衡4.1 同步 vs 异步BaseChatModel._generate()是同步接口但底层的 HTTP 调用必然是异步的。这里用asyncio.run()包装同步入口实际逻辑全在_agenerate()中。如果你的 Agent 系统全部用ainvoke()链路_agenerate()会被直接调用不会触发asyncio.run()的开销。4.2 降级的代价降级链路的问题在于回答质量可能显著下降。从 Qwen2-72B 降到 Ollama 的 Llama-3-8B推理能力下降是必然的。建议在降级时通过generation_info标记让上游的监控系统能区分主模型回答和降级回答分别计算满意度。4.3 什么时候封装 Wrapper什么时候直接用 OpenAI 兼容模式场景推荐方案自部署模型完全兼容 OpenAI API直接ChatOpenAI(base_url...)自部署模型部分兼容但有差异自定义 Wrapper 封装差异需要多模型切换/降级/负载均衡必须自定义 Wrapper模型参数需要前置处理如截断 prompt自定义 Wrapper需要自定义认证头/鉴权逻辑自定义 Wrapper4.4 Token 计数问题vLLM 的 OpenAI 兼容端点不一定返回usage信息。如果你的 LangChain 上层依赖get_num_tokens()做上下文窗口管理需要在 Wrapper 里自己实现 token 计数——最简单的方式是用tiktoken做近似估算精度损失在 5% 以内。五、总结LangChain 的BaseChatModel继承体系为自部署模型提供了一个干净的接入点。通过实现_agenerate()方法你可以把任何自部署模型vLLM、Ollama、甚至是自研的推理服务无缝接入 LangChain 的 Agent 生态。而一个好的自定义 Wrapper 不只是把请求转发过去——它应该包含连接池管理减少 TCP 握手开销智能重试429 退避、5xx 重试、4xx 放过降级链路主模型挂了还有备胎格式映射BaseMessage 和 API 格式的自动转换这套方案我们线上跑了三个月日均 50 万次 LLM 调用vLLM 集群的可用性从裸调时的 99.5% 提升到了 99.95%重试降级贡献的 0.45 个百分点。自部署模型的封装不是技术难题是工程耐心题。把边角处理好了稳定性自然上来。下一篇预告Ruff、mypy、pytest 在 RAG 项目中的协作配置打造 CI 质量防线。