338 lines
12 KiB
Python
338 lines
12 KiB
Python
"""
|
||
LLM 客户端工具
|
||
|
||
提供统一的 LLM 接口,支持多种模型提供商。
|
||
使用 LangChain 的 ChatModel 抽象,简化不同厂商 API 的调用。
|
||
"""
|
||
from typing import Optional, List, Dict, Any, AsyncGenerator
|
||
from langchain_core.language_models.chat_models import BaseChatModel
|
||
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage, AIMessage
|
||
from langchain_core.callbacks import AsyncCallbackHandler
|
||
from core.config import settings
|
||
import time
|
||
|
||
|
||
def get_llm(
|
||
provider: str = "openai",
|
||
model: Optional[str] = None,
|
||
temperature: float = 0.7,
|
||
streaming: bool = False,
|
||
**kwargs
|
||
) -> BaseChatModel:
|
||
"""
|
||
获取 LLM 实例
|
||
|
||
Args:
|
||
provider: 模型提供商 ("openai", "anthropic", "ollama")
|
||
model: 模型名称 (如果不指定则使用配置中的默认值)
|
||
temperature: 温度参数 (0-2)
|
||
streaming: 是否启用流式输出
|
||
**kwargs: 其他参数传递给模型
|
||
|
||
Returns:
|
||
BaseChatModel: LangChain 的聊天模型实例
|
||
"""
|
||
|
||
if provider == "openai":
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
return ChatOpenAI(
|
||
model=model or settings.OPENAI_MODEL or "gpt-4",
|
||
temperature=temperature,
|
||
api_key=settings.OPENAI_API_KEY,
|
||
streaming=streaming,
|
||
**kwargs
|
||
)
|
||
|
||
elif provider == "anthropic":
|
||
from langchain_anthropic import ChatAnthropic
|
||
|
||
return ChatAnthropic(
|
||
model=model or settings.ANTHROPIC_MODEL or "claude-3-opus-20240229",
|
||
temperature=temperature,
|
||
api_key=settings.ANTHROPIC_API_KEY,
|
||
max_tokens=kwargs.pop("max_tokens", 4096),
|
||
streaming=streaming,
|
||
**kwargs
|
||
)
|
||
|
||
elif provider == "ollama":
|
||
try:
|
||
from langchain_ollama import ChatOllama
|
||
|
||
return ChatOllama(
|
||
model=model or settings.OLLAMA_MODEL or "llama3",
|
||
base_url=settings.OLLAMA_BASE_URL or "http://localhost:11434",
|
||
temperature=temperature,
|
||
**kwargs
|
||
)
|
||
except ImportError:
|
||
raise ImportError(
|
||
"langchain-ollama not installed. Run: pip install langchain-ollama"
|
||
)
|
||
|
||
else:
|
||
raise ValueError(f"Unsupported provider: {provider}. Use 'openai', 'anthropic', or 'ollama'")
|
||
|
||
|
||
# 便捷函数 - 常用配置
|
||
def get_fast_llm(provider: str = "openai") -> BaseChatModel:
|
||
"""获取快速响应的 LLM (低温度,适合事实性问题)"""
|
||
return get_llm(provider, temperature=0.3)
|
||
|
||
|
||
def get_creative_llm(provider: str = "openai") -> BaseChatModel:
|
||
"""获取创造性 LLM (高温度,适合创意写作)"""
|
||
return get_llm(provider, temperature=0.9)
|
||
|
||
|
||
def get_streaming_llm(provider: str = "openai") -> BaseChatModel:
|
||
"""获取支持流式输出的 LLM"""
|
||
return get_llm(provider, streaming=True)
|
||
|
||
|
||
class TokenUsageCallbackHandler(AsyncCallbackHandler):
|
||
"""
|
||
Token 使用回调处理器
|
||
|
||
用于捕获 LLM 调用的 token 使用情况
|
||
"""
|
||
def __init__(self):
|
||
self.prompt_tokens = 0
|
||
self.completion_tokens = 0
|
||
self.total_tokens = 0
|
||
self.response_content = ""
|
||
|
||
async def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str] = None, **kwargs):
|
||
"""LLM 开始时的回调"""
|
||
pass
|
||
|
||
async def on_llm_end(self, response, **kwargs):
|
||
"""LLM 结束时的回调，获取 token 统计"""
|
||
try:
|
||
# 从 response 中提取 token 信息
|
||
if hasattr(response, 'llm_output') and response.llm_output:
|
||
token_usage = response.llm_output.get('token_usage', {})
|
||
self.prompt_tokens = token_usage.get('prompt_tokens', 0)
|
||
self.completion_tokens = token_usage.get('completion_tokens', 0)
|
||
self.total_tokens = token_usage.get('total_tokens', 0)
|
||
except Exception as e:
|
||
print(f"[TokenUsageCallback] 提取 token 信息失败: {e}")
|
||
|
||
async def on_llm_new_token(self, token: str, **kwargs):
|
||
"""每个新 token 的回调（流式输出）"""
|
||
self.response_content += token
|
||
|
||
|
||
class LLMClient:
|
||
"""
|
||
LLM 客户端封装类
|
||
|
||
提供统一的异步接口，支持自定义API配置、流式输出和 token 统计
|
||
"""
|
||
|
||
async def chat_completion(
|
||
self,
|
||
messages: List[BaseMessage],
|
||
api_url: str,
|
||
api_key: str,
|
||
model: str = "gpt-3.5-turbo",
|
||
temperature: float = 1.0,
|
||
max_tokens: int = 500,
|
||
request_timeout: int = 60,
|
||
stream: bool = False,
|
||
**kwargs
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
调用 LLM API 生成回复
|
||
|
||
Args:
|
||
messages: LangChain 消息列表
|
||
api_url: API 地址
|
||
api_key: API 密钥
|
||
model: 模型名称
|
||
temperature: 温度参数
|
||
max_tokens: 最大 token 数
|
||
request_timeout: 请求超时时间（秒）
|
||
stream: 是否启用流式输出
|
||
**kwargs: 其他参数
|
||
|
||
Returns:
|
||
OpenAI 格式的响应字典，包含 token 使用信息
|
||
"""
|
||
try:
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
# 创建回调处理器
|
||
callback_handler = TokenUsageCallbackHandler()
|
||
|
||
# 创建自定义的 ChatOpenAI 实例
|
||
llm = ChatOpenAI(
|
||
model=model,
|
||
temperature=temperature,
|
||
api_key=api_key,
|
||
base_url=api_url if api_url else None,
|
||
max_tokens=max_tokens,
|
||
streaming=stream,
|
||
callbacks=[callback_handler],
|
||
request_timeout=request_timeout, # ✅ 设置超时时间
|
||
**kwargs
|
||
)
|
||
|
||
start_time = time.time()
|
||
|
||
if stream:
|
||
# 流式模式
|
||
full_content = ""
|
||
async for chunk in llm.astream(messages):
|
||
if hasattr(chunk, 'content'):
|
||
full_content += chunk.content
|
||
|
||
duration = time.time() - start_time
|
||
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": full_content
|
||
}
|
||
}
|
||
],
|
||
"usage": {
|
||
"prompt_tokens": callback_handler.prompt_tokens,
|
||
"completion_tokens": callback_handler.completion_tokens,
|
||
"total_tokens": callback_handler.total_tokens
|
||
},
|
||
"duration": duration
|
||
}
|
||
else:
|
||
# 非流式模式
|
||
response = await llm.ainvoke(messages)
|
||
duration = time.time() - start_time
|
||
|
||
return {
|
||
"choices": [
|
||
{
|
||
"message": {
|
||
"role": "assistant",
|
||
"content": response.content
|
||
}
|
||
}
|
||
],
|
||
"usage": {
|
||
"prompt_tokens": callback_handler.prompt_tokens,
|
||
"completion_tokens": callback_handler.completion_tokens,
|
||
"total_tokens": callback_handler.total_tokens
|
||
},
|
||
"duration": duration
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"[LLMClient] 调用失败: {e}")
|
||
raise
|
||
|
||
async def stream_chat(
|
||
self,
|
||
messages: List[BaseMessage],
|
||
api_url: str,
|
||
api_key: str,
|
||
model: str = "gpt-3.5-turbo",
|
||
temperature: float = 1.0,
|
||
max_tokens: int = 500,
|
||
request_timeout: int = 60,
|
||
**kwargs
|
||
) -> AsyncGenerator[Dict[str, Any], None]:
|
||
"""
|
||
流式调用 LLM API
|
||
|
||
Args:
|
||
messages: LangChain 消息列表
|
||
api_url: API 地址
|
||
api_key: API 密钥
|
||
model: 模型名称
|
||
temperature: 温度参数
|
||
max_tokens: 最大 token 数
|
||
request_timeout: 请求超时时间（秒）
|
||
**kwargs: 其他参数
|
||
|
||
Yields:
|
||
包含 token 片段的字典
|
||
"""
|
||
try:
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
print(f"\n[LLMClient] 🔧 创建 ChatOpenAI 实例")
|
||
print(f" - Model: {model}")
|
||
print(f" - API URL: {api_url[:50]}..." if len(api_url) > 50 else f" - API URL: {api_url}")
|
||
print(f" - Temperature: {temperature}")
|
||
print(f" - Max Tokens: {max_tokens}")
|
||
print(f" - Request Timeout: {request_timeout}s")
|
||
|
||
# 创建回调处理器
|
||
callback_handler = TokenUsageCallbackHandler()
|
||
|
||
# 创建自定义的 ChatOpenAI 实例
|
||
llm = ChatOpenAI(
|
||
model=model,
|
||
temperature=temperature,
|
||
api_key=api_key,
|
||
base_url=api_url if api_url else None,
|
||
max_tokens=max_tokens,
|
||
streaming=True,
|
||
callbacks=[callback_handler],
|
||
request_timeout=request_timeout, # ✅ 设置超时时间
|
||
**kwargs
|
||
)
|
||
|
||
start_time = time.time()
|
||
|
||
print(f"[LLMClient] 🚀 开始流式请求...")
|
||
print(f" - Messages 数量: {len(messages)}")
|
||
if messages:
|
||
first_msg_role = getattr(messages[0], 'role', 'unknown')
|
||
first_msg_preview = str(getattr(messages[0], 'content', ''))[:50]
|
||
print(f" - 第一条消息: [{first_msg_role}] {first_msg_preview}...")
|
||
|
||
chunk_count = 0
|
||
# 流式输出
|
||
async for chunk in llm.astream(messages):
|
||
if hasattr(chunk, 'content') and chunk.content:
|
||
chunk_count += 1
|
||
|
||
# 第一个 chunk 时记录
|
||
if chunk_count == 1:
|
||
first_chunk_time = time.time()
|
||
print(f"[LLMClient] ✨ 收到第一个 chunk (耗时: {first_chunk_time - start_time:.2f}s)")
|
||
|
||
yield {
|
||
"type": "chunk",
|
||
"content": chunk.content
|
||
}
|
||
|
||
duration = time.time() - start_time
|
||
|
||
print(f"[LLMClient] ✅ 流式请求完成")
|
||
print(f" - 总 Chunks: {chunk_count}")
|
||
print(f" - 耗时: {duration:.2f}秒")
|
||
print(f" - Prompt Tokens: {callback_handler.prompt_tokens}")
|
||
print(f" - Completion Tokens: {callback_handler.completion_tokens}")
|
||
print(f" - Total Tokens: {callback_handler.total_tokens}\n")
|
||
|
||
# 最后发送 token 使用信息
|
||
yield {
|
||
"type": "usage",
|
||
"usage": {
|
||
"prompt_tokens": callback_handler.prompt_tokens,
|
||
"completion_tokens": callback_handler.completion_tokens,
|
||
"total_tokens": callback_handler.total_tokens
|
||
},
|
||
"duration": duration
|
||
}
|
||
|
||
except Exception as e:
|
||
print(f"\n[LLMClient] ❌ 流式调用失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
raise
|