Source code for connectors.haystack_llm_connector

"""
Haystack LLM 连接器模块 - 提供与Haystack框架集成的LLM接口
此模块保持向后兼容性,但建议使用新的LLMFactory
"""

import json
import logging
import os
from typing import Any, Callable, Dict, List, Optional

from haystack import component
from haystack.components.generators.chat import OpenAIChatGenerator
from haystack.dataclasses import ChatMessage
from haystack.utils import Secret

from .base_llm_connector import BaseLLMConnector
from .llm_factory import LLMFactory as _LLMFactory

# Preserve the documented legacy factory import without a second implementation.
LLMFactory = _LLMFactory

# 配置日志
logger = logging.getLogger(__name__)


[docs] class HaystackLLMConnector(BaseLLMConnector): """Haystack LLM 连接器"""
[docs] def __init__( self, api_key: str, model: str = "gpt-3.5-turbo", api_base_url: str = "https://api.openai.com/v1", timeout: int = 60, streaming_callback: Optional[Callable] = None, parameters: Optional[Dict[str, Any]] = None, ): """ 初始化Haystack LLM连接器 (向后兼容) Args: api_key (str): OpenAI API密钥 model (str): 模型名称,默认为"gpt-3.5-turbo" api_base_url (str): API基础URL,默认为``https://api.openai.com/v1`` timeout (int): 请求超时时间(秒),默认为60 streaming_callback (callable, optional): 用于流式响应的回调函数 parameters (dict, optional): LLM参数,如temperature、max_tokens等 """ super().__init__(api_key, model, timeout, parameters) self.api_base_url = api_base_url self.streaming_callback = streaming_callback # 初始化OpenAI生成器 self._init_generator()
def _get_provider_name(self) -> str: return "openai" def _init_generator(self) -> None: """初始化OpenAI Chat生成器""" try: self.generator = OpenAIChatGenerator( api_key=Secret.from_token(self.api_key), model=self.model, api_base_url=self.api_base_url, timeout=self.timeout, streaming_callback=self.streaming_callback, generation_kwargs=self.parameters, ) logger.info(f"成功初始化Haystack OpenAI生成器: {self.model}") except Exception as e: logger.error(f"初始化Haystack OpenAI生成器失败: {str(e)}") raise def _create_chat_message(self, message: Dict[str, str]) -> ChatMessage: """创建ChatMessage使用新的Haystack API""" role = message["role"] content = message["content"] # 角色映射到对应的ChatMessage工厂方法 role_mapping: Dict[str, Callable[[str], ChatMessage]] = { "user": ChatMessage.from_user, "assistant": ChatMessage.from_assistant, "system": ChatMessage.from_system, } # 使用角色特定方法或默认为user factory_method = role_mapping.get(role, ChatMessage.from_user) return factory_method(content)
[docs] def generate( self, messages: List[Dict[str, str]], generation_kwargs: Optional[Dict[str, Any]] = None ) -> Dict[str, Any]: """ 生成对话响应 Args: messages (list): 消息列表,格式为[{"role": "user", "content": "消息内容"}] generation_kwargs (dict, optional): 生成参数,覆盖默认参数 Returns: dict: 生成结果 """ try: # 转换消息格式为Haystack的ChatMessage (使用新API) normalized_messages = self.normalize_messages(messages) chat_messages = [self._create_chat_message(msg) for msg in normalized_messages] # 合并生成参数 params = self.parameters.copy() if generation_kwargs: params.update(generation_kwargs) # 生成响应 logger.debug( f"开始生成对话响应: {json.dumps([m.to_dict() for m in chat_messages], ensure_ascii=False)[:200]}..." ) response = self.generator.run(messages=chat_messages, generation_kwargs=params) # 提取结果 text = response["replies"][0].text if not isinstance(text, str) or not text: raise ValueError("LLM reply does not contain text") result = {"content": text, "role": "assistant", "model": self.model} logger.debug(f"生成对话响应成功: {result['content'][:100]}...") return result except Exception as e: logger.error(f"生成对话响应失败: {str(e)}") return {"content": f"生成对话响应失败: {str(e)}", "role": "assistant", "model": self.model, "error": str(e)}
[docs] def update_parameters(self, parameters: Dict[str, Any]) -> None: """ 更新LLM参数 Args: parameters (dict): 新的LLM参数 """ self.parameters.update(parameters) self.generator.generation_kwargs = self.parameters logger.info(f"已更新LLM参数: {json.dumps(parameters, ensure_ascii=False)}")
[docs] def set_model(self, model: str) -> None: """ 设置模型 Args: model (str): 模型名称 """ if model != self.model: self.model = model self._init_generator() logger.info(f"已更新模型: {model}")
[docs] class HaystackLLMFactory: """Haystack LLM工厂类,创建Haystack LLM连接器"""
[docs] @staticmethod def create_connector(config: Dict[str, Any], streaming_callback: Optional[Callable] = None) -> HaystackLLMConnector: """ 创建Haystack LLM连接器 Args: config (dict): LLM配置 streaming_callback (callable, optional): 用于流式响应的回调函数 Returns: HaystackLLMConnector: Haystack LLM连接器实例 """ api_key = config.get("api_key", os.environ.get("OPENAI_API_KEY", "")) model = config.get("model", "gpt-3.5-turbo") api_base_url = config.get("api_base_url", "https://api.openai.com/v1") timeout = config.get("timeout", 60) parameters = config.get("parameters", {}) return HaystackLLMConnector( api_key=api_key, model=model, api_base_url=api_base_url, timeout=timeout, streaming_callback=streaming_callback, parameters=parameters, )
[docs] @component class ConnectorChatGenerator: """Expose the existing provider-neutral generate contract to Haystack."""
[docs] def __init__(self, connector: BaseLLMConnector) -> None: self.connector = connector
[docs] @component.output_types(replies=List[ChatMessage]) def run( self, messages: List[ChatMessage], generation_kwargs: Optional[Dict[str, Any]] = None ) -> Dict[str, List[ChatMessage]]: normalized = [] for message in messages: text = message.text if not isinstance(text, str): raise ValueError("RAG generation requires text chat messages") normalized.append({"role": message.role.value, "content": text}) result = self.connector.generate(normalized, generation_kwargs=generation_kwargs) if result.get("error"): raise RuntimeError(str(result["error"])) content = result.get("content") if not isinstance(content, str) or not content: raise ValueError("Provider reply does not contain text") return {"replies": [ChatMessage.from_assistant(content)]}