""" 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