1. 项目概述:为什么需要自定义LLM类?
在构建RAG(Retrieval-Augmented Generation)系统时,开发者经常遇到一个关键痛点:不同的大语言模型(LLM)API接口和参数差异巨大。上周我帮一家金融科技公司对接他们的内部知识库时,就发现直接调用原始API会导致代码里遍布if-else分支——当需要切换GPT-4、Claude或本地部署的Llama3时,整个调用逻辑都要重写。
这就是自定义LLM类的核心价值所在。通过构建统一的接口层,我们可以:
- 标准化输入输出格式(比如强制要求返回结构化JSON)
- 内置重试机制和异常处理
- 统一监控和日志记录
- 实现模型的热切换
在LangChain生态中,BaseLLM这个抽象基类已经帮我们定义了最基础的接口规范。但实际企业级应用中,我们往往需要扩展更多功能。比如最近在为某医疗系统开发问答机器人时,就要求所有LLM响应必须包含置信度分数和参考文献定位。
2. 核心架构设计
2.1 类继承关系
推荐采用多层继承结构:
python复制class EnterpriseLLM(BaseLLM):
"""企业级扩展基类"""
def _log_invocation(self, params: dict) -> str:
"""埋点追踪每次调用"""
...
class QwenLLM(EnterpriseLLM):
"""通义千问定制实现"""
def _call(self, prompt: str, **kwargs) -> str:
# 实际调用阿里云API的代码
...
class Llama3LLM(EnterpriseLLM):
"""本地部署的Llama3封装"""
def _call(self, prompt: str, **kwargs) -> str:
# 调用HF pipeline的代码
...
这种设计带来三个关键优势:
- 公共能力下沉到基类(如日志、监控)
- 各模型特有实现隔离在子类
- 新增模型只需实现
_call核心方法
2.2 必须实现的抽象方法
根据LangChain最新1.3.x版本的要求,自定义类必须实现:
_call():核心文本生成逻辑_identifying_params():返回模型标识参数_llm_type():返回模型类型字符串
特别提醒:从1.3.0开始,langchain-core和langchain-community拆分为独立包。如果你看到ImportError: cannot import name 'BaseLLM',请确保安装了正确版本:
bash复制pip install langchain-core==0.1.0 langchain-community==0.0.1
3. 实战:构建支持退避重试的LLM类
3.1 异常处理设计
在金融领域RAG系统中,我们发现LLM API有约3%的失败率。以下是经过验证的退避策略实现:
python复制from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type
)
class RetryableLLM(EnterpriseLLM):
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10),
retry=retry_if_exception_type((TimeoutError, APIError))
)
def _call(self, prompt: str, **kwargs) -> str:
# 实际API调用代码
response = call_api_with_timeout(prompt, timeout=30)
if response.status_code != 200:
raise APIError(f"API返回异常状态码: {response.status_code}")
return response.json()["text"]
关键参数说明:
wait_exponential:指数退避,从4秒开始,最大不超过10秒stop_after_attempt:最多重试3次- 特别捕获TimeoutError和自定义APIError
3.2 上下文管理增强
对于需要维护会话状态的场景(如医疗问诊),建议增加上下文缓存:
python复制from datetime import datetime, timedelta
class ContextAwareLLM(RetryableLLM):
def __init__(self):
self._session_cache = LRUCache(maxsize=1000)
def _call(self, prompt: str, session_id: str = None, **kwargs) -> str:
if session_id:
context = self._session_cache.get(session_id, [])
context.append({"role": "user", "content": prompt})
augmented_prompt = self._format_context(context)
else:
augmented_prompt = prompt
response = super()._call(augmented_prompt, **kwargs)
if session_id:
context.append({"role": "assistant", "content": response})
self._session_cache.set(session_id, context, ttl=timedelta(hours=2))
return response
这个实现可以:
- 通过LRU缓存维护会话历史
- 自动将对话上下文拼接为模型需要的格式
- 设置2小时过期时间避免内存泄漏
4. 性能优化技巧
4.1 流式输出处理
当处理长文档问答时,同步等待完整响应会导致用户体验卡顿。以下是流式处理的实现方案:
python复制from typing import AsyncIterator
class StreamingLLM(RetryableLLM):
async def stream(self, prompt: str, **kwargs) -> AsyncIterator[str]:
async with aiohttp.ClientSession() as session:
async with session.post(
self.api_url,
json={"prompt": prompt},
timeout=30
) as resp:
async for chunk in resp.content.iter_chunks():
yield chunk[0].decode("utf-8")
前端调用示例:
javascript复制const response = await fetch('/api/stream_llm', {
method: 'POST',
body: JSON.stringify({prompt: userQuestion})
});
const reader = response.body.getReader();
while(true) {
const {done, value} = await reader.read();
if(done) break;
document.getElementById('output').innerHTML += value;
}
4.2 批量请求处理
当需要处理大量相似查询时(如客服工单分类),批量调用可提升吞吐量:
python复制from concurrent.futures import ThreadPoolExecutor
class BatchLLM(RetryableLLM):
def batch_call(self, prompts: list[str], max_workers=4) -> list[str]:
with ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(self._call, p) for p in prompts]
return [f.result() for f in futures]
实测数据对比:
| 请求数量 | 串行处理(s) | 批量处理(s) |
|---|---|---|
| 10 | 12.3 | 3.2 |
| 100 | 124.7 | 28.5 |
5. 企业级功能扩展
5.1 敏感词过滤
在金融和医疗行业,我们需要防止模型输出违规内容:
python复制class ComplianceLLM(RetryableLLM):
def __init__(self, blocked_terms: set[str]):
self.blocked_terms = blocked_terms
def _call(self, prompt: str, **kwargs) -> str:
response = super()._call(prompt, **kwargs)
for term in self.blocked_terms:
if term.lower() in response.lower():
raise ComplianceError(f"响应包含违禁词: {term}")
return response
建议结合正则表达式实现更复杂的模式匹配:
python复制import re
medical_pattern = re.compile(r"\b(诊断|处方|治疗建议)\b", flags=re.IGNORECASE)
if medical_pattern.search(response):
raise ComplianceError("模型输出了医疗建议")
5.2 成本监控
对于按token计费的API,必须实时监控用量:
python复制class CostAwareLLM(RetryableLLM):
def __init__(self):
self._token_counter = Counter()
def _call(self, prompt: str, **kwargs) -> str:
response = super()._call(prompt, **kwargs)
input_tokens = estimate_tokens(prompt)
output_tokens = estimate_tokens(response)
self._token_counter.update({
"input": input_tokens,
"output": output_tokens
})
if self._token_counter.total() > 1000000: # 每月100万token限额
raise BudgetExceededError("本月API预算已用尽")
return response
6. 测试策略
6.1 单元测试要点
使用pytest编写测试时,重点验证:
python复制def test_llm_invocation(mock_llm):
# 测试正常调用
response = mock_llm("你好")
assert isinstance(response, str)
# 测试异常重试
with patch('llm.call_api', side_effect=TimeoutError):
with pytest.raises(LLMError):
mock_llm("超时测试")
def test_streaming(llm_server):
# 测试流式输出
chunks = []
async for chunk in llm_server.stream("流式测试"):
chunks.append(chunk)
assert len(chunks) > 1
6.2 压力测试方案
使用locust模拟高并发场景:
python复制from locust import HttpUser, task
class LLMUser(HttpUser):
@task
def invoke_llm(self):
self.client.post("/api/llm", json={
"prompt": "压力测试内容",
"max_tokens": 50
})
启动测试:
bash复制locust -f test_llm.py --headless -u 100 -r 10 -t 5m
关键指标监控:
- 平均响应时间 < 2秒
- 错误率 < 0.5%
- 99分位延迟 < 5秒
7. 部署最佳实践
7.1 容器化配置
推荐Dockerfile配置:
dockerfile复制FROM python:3.10-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
ENV LLM_API_KEY="your_key"
ENV MAX_CONCURRENT=4
CMD ["gunicorn", "-w 4", "-k uvicorn.workers.UvicornWorker", "app:server"]
关键优化点:
- 使用slim镜像减少体积
- 分层构建加速CI/CD
- 通过环境变量注入秘钥
- 根据CPU核心数设置worker数量
7.2 健康检查端点
必须实现的K8s健康检查:
python复制from fastapi import APIRouter
router = APIRouter()
@router.get("/health")
async def health_check():
try:
test_response = await llm("健康检查", max_tokens=1)
return {"status": "healthy"}
except Exception as e:
raise HTTPException(status_code=503, detail=str(e))
Prometheus监控指标示例:
python复制from prometheus_client import Counter
LLM_REQUESTS = Counter(
'llm_requests_total',
'Total LLM API requests',
['model', 'status']
)
class InstrumentedLLM(RetryableLLM):
def _call(self, prompt: str, **kwargs) -> str:
start_time = time.time()
try:
response = super()._call(prompt, **kwargs)
[LLM](https://taotoken.net?utm_source=ai)_REQUESTS.labels(model=self.model_name, status="success").inc()
return response
except Exception as e:
LLM_REQUESTS.labels(model=self.model_name, status="failed").inc()
raise
8. 常见问题排查
8.1 内存泄漏问题
症状:服务运行一段时间后内存持续增长
排查步骤:
- 使用memory_profiler定位泄漏点:
python复制@profile def process_batch(prompts): return [llm(p) for p in prompts] - 检查是否未正确关闭:
- 数据库连接
- 文件句柄
- 网络会话
8.2 响应时间波动
可能原因及解决方案:
| 现象 | 可能原因 | 解决方案 |
|---|---|---|
| 白天延迟高 | 共享API配额耗尽 | 升级企业版或设置速率限制 |
| 特定关键词触发慢 | 模型安全审查导致延迟 | 本地部署或使用轻量级模型 |
| 周期性延迟 spikes | 底层资源争用 | 分离计算节点或设置资源预留 |
8.3 模型输出不一致
调试方法:
- 设置确定性参数:
python复制llm = CustomLLM( temperature=0.3, # 降低随机性 top_p=0.9, seed=42 # 固定随机种子 ) - 记录完整推理参数:
python复制logging.info(f"LLM调用参数: {llm._identifying_params()}")
9. 进阶扩展方向
9.1 多模型路由
智能路由实现示例:
python复制class RouterLLM(BaseLLM):
def __init__(self, models: dict[str, BaseLLM]):
self.models = models
def route(self, prompt: str) -> str:
if "代码" in prompt:
return "claude" # Claude更擅长代码
elif len(prompt) > 1000:
return "gpt-4-32k" # 处理长文本
else:
return "gpt-3.5" # 默认选择
def _call(self, prompt: str, **kwargs) -> str:
model_name = self.route(prompt)
return self.models[model_name](prompt, **kwargs)
9.2 混合本地与云端
分级调用策略:
- 先尝试本地模型(如Llama3-70B)
- 如置信度低于阈值,转云模型
- 记录反馈数据持续优化本地模型
实现代码:
python复制class HybridLLM(BaseLLM):
def _call(self, prompt: str, threshold=0.7, **kwargs) -> str:
local_response = self.local_llm(prompt)
confidence = calculate_confidence(local_response)
if confidence < threshold:
cloud_response = self.cloud_llm(prompt)
self._log_weak_response(prompt, local_response, cloud_response)
return cloud_response
return local_response
10. 版本兼容性处理
随着LangChain生态快速迭代,需要注意:
-
接口变更追踪:
python复制try: from langchain_core.language_models import BaseLLM except ImportError: # 回退到旧版导入路径 from langchain.base import BaseLLM -
多版本适配层:
python复制class VersionAdapter: @staticmethod def get_chat_history(messages): if LANGCHAIN_VERSION >= (0, 1): return format_to_openai(messages) else: return messages -
依赖声明规范:
python复制# setup.py install_requires=[ 'langchain-core>=0.1.0,<0.2.0', 'langchain-community>=0.0.1,<0.1.0', ]
