isage-middleware 0.2.4.3__cp311-none-any.whl

sage/middleware/operators/rag/generator.py

@@ -0,0 +1,340 @@
+import json
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from openai import OpenAI
+
+from sage.common.config.output_paths import get_states_file
+from sage.common.core.functions import MapFunction as MapOperator
+from sage.libs.integrations.huggingface import HFClient
+
+
+class OpenAIGenerator(MapOperator):
+    """
+    生成节点：调用 OpenAI-Compatible / SageLLM 等端点。
+
+    调用方式::
+        sub_conf = config["generator"]["sagellm"]   # <- 单端点子配置
+        gen = OpenAIGenerator(sub_conf)
+
+    其中 `sub_conf` 结构示例::
+
+        {
+          "method":     "openai",
+          "model_name": "gpt-4o-mini",
+          "base_url":   "http://localhost:8000/v1",
+          "api_key":    "xxx",  # pragma: allowlist secret
+          "seed":       42
+        }
+    """
+
+    def __init__(self, config: dict, enable_profile=False, **kwargs):
+        super().__init__(**kwargs)
+
+        # 直接持有子配置
+        self.config = config
+        self.enable_profile = enable_profile
+
+        # 实例化模型
+        # API key 优先级: 配置文件 > OPENAI_API_KEY
+        api_key = self.config.get("api_key") or os.getenv("OPENAI_API_KEY")
+
+        # 获取必需的配置参数（使用 .get() 提供默认值）
+        model_name = self.config.get("model_name") or self.config.get("model", "gpt-3.5-turbo")
+        # 展开环境变量（如果 model_name 包含环境变量）
+        model_name = os.path.expandvars(model_name)
+        base_url = self.config.get("base_url", "https://api.openai.com/v1")
+
+        # 直接使用 OpenAI 客户端（支持 sagellm 等 OpenAI 兼容 API）
+        self.model = OpenAI(
+            base_url=base_url,
+            api_key=api_key or "EMPTY",  # 本地服务可用任意 key
+        )
+        self.model_name = model_name
+        self.num = 1
+
+        # 只有启用profile时才设置数据存储路径
+        if self.enable_profile:
+            # Use unified output path system
+            self.data_base_path = str(get_states_file("dummy", "generator_data").parent)
+            os.makedirs(self.data_base_path, exist_ok=True)
+            self.data_records = []
+
+    def _save_data_record(self, query, prompt, response):
+        """保存生成数据记录"""
+        if not self.enable_profile:
+            return
+
+        record = {
+            "timestamp": time.time(),
+            "query": query,
+            "prompt": prompt,
+            "response": response,
+            "model_name": self.config.get("model_name") or self.config.get("model", "unknown"),
+        }
+        self.data_records.append(record)
+        self._persist_data_records()
+
+    def _persist_data_records(self):
+        """将数据记录持久化到文件"""
+        if not self.enable_profile or not self.data_records:
+            return
+
+        timestamp = int(time.time())
+        filename = f"generator_data_{timestamp}.json"
+        path = os.path.join(self.data_base_path, filename)
+
+        try:
+            with open(path, "w", encoding="utf-8") as f:
+                json.dump(self.data_records, f, ensure_ascii=False, indent=2)
+            self.data_records = []
+        except Exception as e:
+            self.logger.error(f"Failed to persist data records: {e}")
+
+    def execute(self, data: list[Any]) -> dict[str, Any]:
+        """
+        输入 : [original_data, prompt]  *或*  [prompt]
+        输出 : 完整的数据字典，包含 generated 字段
+
+        prompt 可以是:
+        - str: 普通字符串，将转换为 [{"role": "user", "content": prompt}]
+        - list[dict]: 已格式化的消息列表，直接传递给 OpenAI API
+        """
+        # 解析输入数据
+        if len(data) > 1:
+            # 来自QAPromptor: [original_data, prompt]
+            original_data = data[0]
+            prompt = data[1]
+        else:
+            # 直接prompt输入: [prompt]
+            original_data = {}
+            prompt = data[0]
+
+        # 提取user_query
+        if isinstance(original_data, dict):
+            user_query = original_data.get("query", original_data.get("question", ""))
+        else:
+            user_query = None
+
+        # 如果 prompt 是字符串，转换为标准消息格式
+        if isinstance(prompt, str):
+            messages = [{"role": "user", "content": prompt}]
+        elif isinstance(prompt, list) and all(isinstance(item, dict) for item in prompt):
+            # 如果已经是消息列表格式，直接使用
+            messages = prompt
+        else:
+            # 兜底处理：转换为字符串再构造消息
+            messages = [{"role": "user", "content": str(prompt)}]
+
+        # 准备生成参数（从配置中提取）
+        generate_kwargs = {}
+
+        # 支持的参数列表
+        supported_params = [
+            "max_tokens",
+            "temperature",
+            "top_p",
+            "enable_thinking",  # Qwen 特有参数：禁用思考过程输出
+            "stream",
+            "frequency_penalty",
+            "n",
+            "logprobs",
+        ]
+
+        # 从配置中提取参数并传递给 generate
+        for param in supported_params:
+            if param in self.config:
+                generate_kwargs[param] = self.config[param]
+
+        # 使用 OpenAI 客户端调用 chat completions API
+        completion = self.model.chat.completions.create(
+            model=self.model_name,
+            messages=messages,
+            **generate_kwargs,
+        )
+        response = completion.choices[0].message.content
+
+        self.num += 1
+
+        # 保存数据记录（只有enable_profile=True时才保存）
+        if self.enable_profile:
+            self._save_data_record(user_query, prompt, response)
+
+        self.logger.info(f"[{self.__class__.__name__}] Response: {response}")
+
+        # 构建完整的输出数据，保持上游数据
+        if isinstance(original_data, dict):
+            # 保持原始数据结构，添加generated字段
+            result = dict(original_data)
+            result["generated"] = response
+            # generate_time 由 MapOperator 自动添加
+            result["question"] = result.get(
+                "question",
+                {"query": user_query, "references": result.get("references", [])},
+            )
+            return result
+        else:
+            # 兼容原有tuple格式输出，但符合返回类型
+            return {
+                "query": user_query if user_query is not None else "",
+                "generated": response,
+                # generate_time 由 MapOperator 自动添加
+            }
+
+    def __del__(self):
+        """确保在对象销毁时保存所有未保存的记录"""
+        if hasattr(self, "enable_profile") and self.enable_profile:
+            try:
+                self._persist_data_records()
+            except Exception:
+                pass
+
+
+class HFGenerator(MapOperator):
+    """
+    HFGenerator is a generator rag that interfaces with a Hugging Face model
+    to generate responses based on input data.
+    """
+
+    def __init__(self, config, **kwargs):
+        """
+        Initializes the HFGenerator instance with configuration parameters.
+
+        :param config: Dictionary containing configuration for the generator, including
+                       the method and model name.
+        """
+        super().__init__(**kwargs)
+        self.config = config
+        # Apply the generator model with the provided configuration
+        self.model = HFClient(model_name=self.config["model_name"])
+
+    def execute(self, data: list, **kwargs) -> tuple[str, str]:
+        """
+        Executes the response generation using the configured Hugging Face model based on the input data.
+
+        :param data: Data object containing a list of input data.
+                     The expected format and the content of the data depend on the model's requirements.
+        :param kwargs: Additional parameters for the model generation (e.g., temperature, max_tokens, etc.).
+
+        :return: A Data object containing the generated response as a string.
+        """
+        # Generate the response from the Hugging Face model using the provided data and additional arguments
+        user_query = data[0] if len(data) > 1 else None
+
+        prompt = data[1] if len(data) > 1 else data[0]
+
+        response = self.model.generate(prompt, **kwargs)
+
+        print(f"\033[32m[ {self.__class__.__name__}]: Response: {response}\033[0m ")
+
+        # Return the generated response as a Data object
+        self.logger.info(f"\033[32m[ {self.__class__.__name__}]: Response: {response}\033[0m ")
+
+        return (
+            user_query if user_query is not None else "",
+            response if isinstance(response, str) else str(response),
+        )
+
+
+@dataclass
+class SageLLMRAGGenerator(MapOperator):
+    """
+    RAG 生成器 - 使用 SageLLM 引擎
+
+    通过 engine_type 参数选择底层 LLM 引擎：
+    - sagellm (默认): 使用 SageLLMGenerator，支持 auto/mock/cuda/ascend 后端
+
+    Example:
+        ```python
+        # 使用 sagellm 引擎（推荐）
+        generator = SageLLMRAGGenerator(
+            engine_type="sagellm",
+            backend_type="auto",
+            model_path="Qwen/Qwen2.5-7B-Instruct",
+            max_tokens=2048,
+        )
+        ```
+
+    Attributes:
+        engine_type: 引擎类型，支持 "sagellm"（默认）
+        backend_type: 后端类型，支持 "auto"/"mock"/"cuda"/"ascend"
+        model_path: 模型路径或 HuggingFace 模型 ID
+        max_tokens: 最大生成 token 数
+        temperature: 采样温度
+        top_p: nucleus 采样参数
+        timeout: 请求超时时间
+    """
+
+    # 引擎选择
+    engine_type: str = "sagellm"  # sagellm only
+    backend_type: str = "auto"  # auto/mock/cuda/ascend
+
+    # SageLLM 配置
+    model_path: str = ""
+    device_map: str = "auto"
+    dtype: str = "auto"
+
+    # 生成参数
+    max_tokens: int = 2048
+    temperature: float = 0.7
+    top_p: float = 0.95
+    top_k: int = 50
+
+    # 配置
+    timeout: float = 120.0
+    default_options: dict[str, Any] = field(default_factory=dict)
+
+    # 内部状态
+    _generator: Any = field(default=None, init=False, repr=False)
+
+    def __post_init__(self) -> None:
+        super().__init__()
+        self._init_generator()
+
+    def _init_generator(self) -> None:
+        """根据 engine_type 初始化底层生成器"""
+        if self.engine_type != "sagellm":
+            # 只支持 sagellm
+            raise ValueError(
+                f"Unsupported engine_type='{self.engine_type}'. "
+                f"Only 'sagellm' is supported. vLLM support has been removed in v0.3.0."
+            )
+
+        # 默认使用 sagellm
+        from sage.middleware.operators.llm import SageLLMGenerator
+
+        self._generator = SageLLMGenerator(
+            backend_type=self.backend_type,
+            model_path=self.model_path,
+            device_map=self.device_map,
+            dtype=self.dtype,
+            max_tokens=self.max_tokens,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            timeout=self.timeout,
+            default_options=self.default_options,
+        )
+
+    def execute(self, data: list[Any]) -> dict[str, Any]:
+        """
+        执行生成，委托给底层生成器
+
+        输入 : [original_data, prompt] 或 [prompt]
+        输出 : 包含 generated 字段的数据字典
+        """
+        result = self._generator.execute(data)
+
+        # 统一输出格式
+        if isinstance(result, dict):
+            return result
+        elif isinstance(result, tuple) and len(result) >= 2:
+            # Generator returns (original, text)
+            return {
+                "query": result[0] if result[0] else "",
+                "generated": result[1],
+            }
+        else:
+            return {"generated": str(result)}

sage/middleware/operators/rag/index_builder/__init__.py

@@ -0,0 +1,48 @@
+"""RAG Index Building Service (L4 Middleware Operator)
+
+This module provides index building functionality for RAG systems.
+It orchestrates document processing, embedding, and vector storage.
+
+Layer: L4 (sage-middleware/operators/rag)
+Dependencies:
+  - sage.libs.rag (L3) - chunk, document_loaders
+  - sage.middleware.components.sage_db (L4) - SageDB backend
+  - sage.common (L1) - embedding models
+
+Components:
+- VectorStore: Protocol defining vector storage interface
+- IndexManifest: Metadata describing a built index
+- IndexBuilder: Service for building vector indices
+
+Architecture Pattern:
+- L4 defines IndexBuilder (orchestration)
+- L4 provides SageDB backend implementation
+- L3 provides ChromaDB backend via integrations
+- L5 (sage-cli) uses IndexBuilder
+
+Example Usage:
+    >>> from sage.middleware.operators.rag.index_builder import IndexBuilder
+    >>> from sage.middleware.components.sage_db import SageVDBBackend
+    >>>
+    >>> # Create backend factory
+    >>> def factory(path, dim):
+    ...     return SageVDBBackend(path, dim)
+    >>>
+    >>> # Build index
+    >>> builder = IndexBuilder(backend_factory=factory)
+    >>> manifest = builder.build_from_docs(
+    ...     source_dir=Path("docs"),
+    ...     persist_path=Path(".sage/index"),
+    ...     embedding_model=embedder,
+    ... )
+"""
+
+from sage.middleware.operators.rag.index_builder.builder import IndexBuilder
+from sage.middleware.operators.rag.index_builder.manifest import IndexManifest
+from sage.middleware.operators.rag.index_builder.storage import VectorStore
+
+__all__ = [
+    "IndexBuilder",
+    "IndexManifest",
+    "VectorStore",
+]