isage-middleware 0.2.4.3__cp311-none-any.whl

sage/middleware/operators/rag/promptor.py

@@ -0,0 +1,400 @@
+import json
+import os
+import time
+
+from jinja2 import Template
+
+from sage.common.core.functions import MapFunction as MapOperator
+
+QA_prompt_template_str = """Instruction:
+You are an intelligent assistant with access to a knowledge base. Answer the question below with reference to the provided context.
+Only give me the answer and do not output any other words.
+{%- if external_corpus %}
+Relevant corpus for the current question:
+{{ external_corpus }}
+{%- endif %}
+"""
+
+QA_short_answer_template_str = """Instruction:
+You are an intelligent assistant with access to a knowledge base. Answer the question below with reference to the provided context.
+Please provide a concise answer and conclude with 'So the final answer is: [your answer]'.
+{%- if external_corpus %}
+Relevant corpus for the current question:
+{{ external_corpus }}
+{%- endif %}
+"""
+
+summarization_prompt_template_str = """Instruction:
+You are an intelligent assistant. Summarize the content provided below in a concise and clear manner.
+Only provide the summary and do not include any additional information.
+{%- if external_corpus %}
+Content to summarize:
+{{ external_corpus }}
+{%- endif %}
+"""
+QA_prompt_template = Template(QA_prompt_template_str)
+QA_short_answer_template = Template(QA_short_answer_template_str)
+summarization_prompt_template = Template(summarization_prompt_template_str)
+
+query_profiler_prompt_template_str = """
+For the given query = how Trump earn his first 1 million dollars?: Analyze the language and internal structure of the query and provide the following information:
+
+1. Does it need joint reasoning across multiple documents?
+2. Provide a complexity profile for the query:
+   - Complexity: High / Low
+   - Joint Reasoning needed: Yes / No
+3. Does this query need input chunks to be summarized? If yes, provide a range in words for the summarized chunks.
+4. How many distinct pieces of information are needed to answer the query?
+
+database_metadata = The dataset consists of multiple chunks of information from Fortune 500 companies on financial reports from every quarter of 2023.
+chunk_size = 1024
+
+Estimate the query profile along with the database_metadata and chunk_size.
+
+Your output must be:
+- **Only a valid JSON object**
+- **No explanations, no formatting, no comments**
+- **No markdown code blocks or prose**
+- **Strictly conform to this schema:**
+
+{
+  "need_joint_reasoning": <true|false>,
+  "complexity": "High" or "Low",
+  "need_summarization": <true|false>,
+  "summarization_length": integer (30-200),
+  "n_info_items": integer (1-6)
+}
+"""
+query_profiler_prompt_template = Template(query_profiler_prompt_template_str)
+
+
+class QAPromptor(MapOperator):
+    """
+    QAPromptor is a prompt rag that generates a QA-style prompt using
+    an external corpus and a user query. This class is designed to prepare
+    the necessary prompt structure for a question-answering model.
+
+    Attributes:
+        config: Configuration data for initializing the prompt rag (e.g., model details, etc.).
+        prompt_template: A template used for generating the system prompt, typically includes context or instructions.
+    """
+
+    prompt_template: Template
+
+    def __init__(self, config, enable_profile=False, **kwargs):
+        super().__init__(**kwargs)
+
+        """
+        Initializes the QAPromptor instance with configuration and prompt template.
+
+        :param config: Dictionary containing configuration for the prompt rag.
+        """
+        self.config = config  # Store the configuration for later use
+        self.enable_profile = enable_profile
+
+        # 使用配置文件中的模板，如果没有则使用默认模板
+        self.use_short_answer = config.get("use_short_answer", False)  # 是否使用短答案模式
+
+        if "template" in config:
+            from jinja2 import Template
+
+            self.prompt_template = Template(config["template"])
+        else:
+            # 根据配置选择模板
+            if self.use_short_answer:
+                self.prompt_template = QA_short_answer_template
+            else:
+                self.prompt_template = QA_prompt_template  # Load the QA prompt template
+
+        # 只有启用profile时才设置数据存储路径
+        if self.enable_profile:
+            from sage.common.config.output_paths import get_sage_paths
+
+            try:
+                sage_paths = get_sage_paths()
+                self.data_base_path = str(sage_paths.states_dir / "promptor_data")
+            except Exception:
+                # Fallback to current working directory
+                if (
+                    self.ctx is not None
+                    and hasattr(self.ctx, "env_base_dir")
+                    and self.ctx.env_base_dir
+                ):
+                    self.data_base_path = os.path.join(
+                        self.ctx.env_base_dir, ".sage_states", "promptor_data"
+                    )
+                else:
+                    # 使用默认路径
+                    self.data_base_path = os.path.join(os.getcwd(), ".sage_states", "promptor_data")
+
+            os.makedirs(self.data_base_path, exist_ok=True)
+            self.data_records = []
+
+    def _save_data_record(self, query, external_corpus, prompt):
+        """保存提示词数据记录"""
+        if not self.enable_profile:
+            return
+
+        record = {
+            "timestamp": time.time(),
+            "query": query,
+            "external_corpus": external_corpus,
+            "prompt": prompt,
+        }
+        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"promptor_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}")
+
+    # sage_lib/functions/rag/qapromptor.py
+    def execute(self, data) -> list:
+        """
+        生成 ChatGPT 风格的 prompt（system+user 两条消息）。
+
+        支持多种输入格式：
+        1. (query, external_corpus_list_or_str)  # 元组格式
+        2. query_str  # 纯字符串
+        3. {"query": ..., "results": [...]}  # 字典格式（来自检索器）
+        4. {"question": ..., "context": [...]}  # 字典格式（来自测试）
+        """
+        self.logger.info(f"QAPromptor received data: {data}")
+        try:
+            # -------- 解析输入 --------
+            raw = data
+            original_data = data  # 保存原始数据以便返回
+
+            if isinstance(raw, dict):
+                # 字典格式输入 - 支持多种字段名
+                query = raw.get("query", raw.get("question", ""))
+
+                # 处理不同的上下文字段名
+                external_corpus_list = []
+
+                # 处理 refining_results 字段（来自 refiner - 压缩后的文档）
+                if "refining_results" in raw:
+                    results = raw.get("refining_results", [])
+                    for result in results:
+                        if isinstance(result, str):
+                            external_corpus_list.append(result)
+                        else:
+                            external_corpus_list.append(str(result))
+
+                # 处理 retrieval_results 字段（来自 retriever - 原始检索结果）
+                elif "retrieval_results" in raw:
+                    results = raw.get("retrieval_results", [])
+                    for result in results:
+                        if isinstance(result, dict) and "text" in result:
+                            external_corpus_list.append(result["text"])
+                        elif isinstance(result, str):
+                            external_corpus_list.append(result)
+                        else:
+                            external_corpus_list.append(str(result))
+
+                # 处理 context 字段（来自测试）
+                elif "context" in raw:
+                    context = raw.get("context", [])
+                    if isinstance(context, list):
+                        external_corpus_list.extend([str(c) for c in context])
+                    else:
+                        external_corpus_list.append(str(context))
+
+                # 处理 external_corpus 字段
+                elif "external_corpus" in raw:
+                    external_corpus = raw.get("external_corpus", "")
+                    if isinstance(external_corpus, list):
+                        external_corpus_list.extend([str(c) for c in external_corpus])
+                    else:
+                        external_corpus_list.append(str(external_corpus))
+
+                external_corpus = "\n".join(external_corpus_list)
+
+            elif isinstance(raw, tuple) and len(raw) == 2:
+                # 元组格式输入
+                query, external_corpus = raw
+                if isinstance(external_corpus, list):
+                    external_corpus = "\n".join(external_corpus)
+                # 对于元组输入，保持原有行为，返回query而不是原始数据
+                original_data = query
+            else:
+                # 字符串格式输入
+                query = str(raw)
+                external_corpus = ""
+                # 对于字符串输入，保持原有行为，返回query而不是原始数据
+                original_data = query
+
+            external_corpus = external_corpus or ""
+
+            # -------- system prompt --------
+            if external_corpus:
+                system_prompt = {
+                    "role": "system",
+                    "content": self.prompt_template.render(external_corpus=external_corpus),
+                }
+            else:
+                system_prompt = {
+                    "role": "system",
+                    "content": (
+                        "You are a helpful AI assistant. Answer the user's questions accurately."
+                    ),
+                }
+
+            # -------- user prompt --------
+            user_prompt = {
+                "role": "user",
+                "content": f"Question: {query}",
+            }
+            self.logger.info(
+                f"QAPromptor generated prompt: {system_prompt['content']} | {user_prompt['content']}"
+            )
+            prompt = [system_prompt, user_prompt]
+
+            # 保存数据记录（只有enable_profile=True时才保存）
+            if self.enable_profile:
+                self._save_data_record(query, external_corpus, prompt)
+
+            return [original_data, prompt]
+
+        except Exception as e:
+            self.logger.error("QAPromptor error: %s | input=%s", e, getattr(data, "data", ""))
+            fallback = [
+                {"role": "system", "content": "System encountered an error."},
+                {
+                    "role": "user",
+                    "content": (
+                        "Question: Error occurred. Please try again."
+                        f" (Original: {getattr(data, 'data', '')})"
+                    ),
+                },
+            ]
+            return fallback
+
+    def __del__(self):
+        """确保在对象销毁时保存所有未保存的记录"""
+        if hasattr(self, "enable_profile") and self.enable_profile:
+            try:
+                self._persist_data_records()
+            except Exception:
+                pass
+
+
+class SummarizationPromptor(MapOperator):
+    """
+    QAPromptor is a prompt rag that generates a QA-style prompt using
+    an external corpus and a user query. This class is designed to prepare
+    the necessary prompt structure for a question-answering model.
+
+    Attributes:
+        config: Configuration data for initializing the prompt rag (e.g., model details, etc.).
+        prompt_template: A template used for generating the system prompt, typically includes context or instructions.
+    """
+
+    prompt_template: Template
+
+    def __init__(self, config):
+        """
+        Initializes the QAPromptor instance with configuration and prompt template.
+
+        :param config: Dictionary containing configuration for the prompt rag.
+        """
+        super().__init__()
+        self.config = config  # Store the configuration for later use
+        self.prompt_template = (
+            summarization_prompt_template  # Load the summarization prompt template
+        )
+
+    def execute(self, data) -> list:
+        """
+        Generates a QA-style prompt for the input question and external corpus.
+
+        This method takes the query and external corpus, processes the corpus
+        into a single string, and creates a system prompt and user prompt based
+        on a predefined template.
+
+        :param data: A Data object containing a tuple. The first element is the query (a string),
+                     and the second is a list of external corpus (contextual information for the model).
+
+        :return: A Data object containing a list with two prompts:
+                 1. system_prompt: A system prompt based on the template with external corpus data.
+                 2. user_prompt: A user prompt containing the question to be answered.
+        """
+        # Unpack the input data into query and external_corpus
+        query, external_corpus = data
+
+        # Combine the external corpus list into a single string (in case it's split into multiple parts)
+        external_corpus = "".join(external_corpus)
+
+        # Prepare the base data for the system prompt, which includes the external corpus
+        base_system_prompt_data = {"external_corpus": external_corpus}
+
+        # query = data
+        # Create the system prompt using the template and the external corpus data
+        system_prompt = {
+            "role": "system",
+            "content": self.prompt_template.render(**base_system_prompt_data),
+        }
+        # system_prompt = {
+        #     "role": "system",
+        #     "content": ""
+        # }
+        # Create the user prompt using the query
+        user_prompt = {"role": "user", "content": f"Question: {query}"}
+
+        # Combine the system and user prompts into one list
+        prompt = [system_prompt, user_prompt]
+
+        # Return the prompt list wrapped in a Data object
+        return prompt
+
+
+class QueryProfilerPromptor(MapOperator):
+    """
+    QueryProfilerPromptor provides a prompt for profiling queries.
+
+    """
+
+    prompt_template: Template
+
+    def __init__(self, config):
+        """
+        Initializes the QueryProfilerPromptor instance with configuration and prompt template.
+
+        :param config: Dictionary containing configuration for the prompt rag.
+        """
+        super().__init__()
+        self.config = config  # Store the configuration for later use
+        self.prompt_template = (
+            query_profiler_prompt_template  # Load the query profiler prompt template
+        )
+
+    def execute(self, data) -> list:
+        """
+        Generates a profiling prompt for the input query.
+
+        :param data: A string representing the query to be profiled.
+
+        :return: A list containing the profiling prompt.
+        """
+        query = data
+        prompt = {
+            "role": "user",
+            "content": self.prompt_template.render(
+                query=query,
+                metadata=self.config.get("metadata", {}),
+                chunk_size=self.config.get("chunk_size", 1024),
+            ),
+        }
+        return [prompt]

sage/middleware/operators/rag/refiner.py

@@ -0,0 +1,231 @@
+"""
+Refiner Operator - SAGE RAG Pipeline Operator
+==============================================
+
+Uses isage-refiner (sage_refiner) for context compression in RAG pipelines.
+
+Installation:
+    pip install isage-refiner
+
+Usage:
+    from sage.middleware.operators.rag import RefinerOperator
+
+    config = {
+        "algorithm": "long_refiner",  # or "reform", "provence", etc.
+        "budget": 2048,
+        # LongRefiner specific
+        "base_model_path": "Qwen/Qwen2.5-3B-Instruct",
+        ...
+    }
+
+    env.map(RefinerOperator, config)
+"""
+
+import json
+import os
+import time
+from typing import Any
+
+from sage.common.config.output_paths import get_states_file
+from sage.common.core.functions import MapFunction as MapOperator
+
+
+class RefinerOperator(MapOperator):
+    """
+    Refiner Operator for SAGE RAG pipelines.
+
+    Wraps isage-refiner compressors (LongRefiner, REFORM, Provence, etc.)
+    for use in SAGE dataflow pipelines.
+
+    Config:
+        algorithm: str - "long_refiner", "reform", "provence", "llmlingua2", etc.
+        budget: int - Token budget for compression
+        enable_profile: bool - Enable data recording for debugging
+
+        # Algorithm-specific config passed through to compressor
+        base_model_path: str - For LongRefiner
+        score_model_path: str - For LongRefiner
+        ...
+    """
+
+    def __init__(self, config: dict, ctx=None):
+        super().__init__(config=config, ctx=ctx)
+        self.cfg = config
+        self.enable_profile = config.get("enable_profile", False)
+        self.compressor = None
+
+        # Data recording (only when enable_profile=True)
+        if self.enable_profile:
+            self.data_base_path = str(get_states_file("dummy", "refiner_data").parent)
+            os.makedirs(self.data_base_path, exist_ok=True)
+            self.data_records: list[dict] = []
+
+        self._init_compressor()
+
+    def _init_compressor(self):
+        """Initialize the compressor from isage-refiner."""
+        algorithm = self.cfg.get("algorithm", "long_refiner").lower()
+
+        try:
+            if algorithm == "long_refiner":
+                from sage_refiner import LongRefinerCompressor
+
+                self.compressor = LongRefinerCompressor(
+                    base_model_path=self.cfg.get("base_model_path", "Qwen/Qwen2.5-3B-Instruct"),
+                    query_analysis_module_lora_path=self.cfg.get(
+                        "query_analysis_module_lora_path", ""
+                    ),
+                    doc_structuring_module_lora_path=self.cfg.get(
+                        "doc_structuring_module_lora_path", ""
+                    ),
+                    global_selection_module_lora_path=self.cfg.get(
+                        "global_selection_module_lora_path", ""
+                    ),
+                    score_model_path=self.cfg.get("score_model_path", "BAAI/bge-reranker-v2-m3"),
+                    max_model_len=self.cfg.get("max_model_len", 25000),
+                    gpu_memory_utilization=self.cfg.get("gpu_memory_utilization", 0.5),
+                )
+
+            elif algorithm == "reform":
+                from sage_refiner import REFORMCompressor
+
+                self.compressor = REFORMCompressor(**self.cfg.get("reform_config", {}))
+
+            elif algorithm == "provence":
+                from sage_refiner import ProvenceCompressor
+
+                self.compressor = ProvenceCompressor(**self.cfg.get("provence_config", {}))
+
+            elif algorithm in ("simple", "none"):
+                # Simple truncation - no compression
+                self.compressor = None
+                self.logger.info("Using simple/none mode - no compression")
+
+            else:
+                raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+            self.logger.info(f"RefinerOperator initialized with algorithm: {algorithm}")
+
+        except ImportError as e:
+            raise ImportError(
+                f"Failed to import {algorithm} compressor. "
+                f"Install with: pip install isage-refiner\n"
+                f"Error: {e}"
+            ) from e
+
+    def execute(self, data: dict):
+        """Execute document compression.
+
+        Input format:
+            {
+                "query": str,
+                "retrieval_results": List[Dict],  # Retrieved documents
+            }
+
+        Output format:
+            {
+                "query": str,
+                "retrieval_results": List[Dict],  # Original (preserved)
+                "refining_results": List[str],    # Compressed document texts
+            }
+        """
+        if not isinstance(data, dict):
+            self.logger.error(f"Unexpected input format: {type(data)}")
+            return data
+
+        query = data.get("query", "")
+        docs = data.get("retrieval_results", [])
+
+        # Normalize documents to isage-refiner format
+        documents = self._normalize_documents(docs)
+
+        # Compress
+        try:
+            if self.compressor is None:
+                # Simple mode: just extract text
+                refined_texts = [
+                    doc.get("contents", doc.get("text", str(doc))) for doc in documents
+                ]
+            else:
+                budget = self.cfg.get("budget", 2048)
+                result = self.compressor.compress(
+                    question=query,
+                    document_list=documents,
+                    budget=budget,
+                )
+                # isage-refiner returns dict with various fields
+                refined_texts = result.get("compressed_context", "")
+                if isinstance(refined_texts, str):
+                    refined_texts = [refined_texts]
+
+        except Exception as e:
+            self.logger.error(f"Refiner execution failed: {e}")
+            refined_texts = [doc.get("contents", str(doc)) for doc in documents]
+
+        # Save data record if profiling
+        if self.enable_profile:
+            self._save_data_record(query, documents, refined_texts)
+
+        # Build output
+        result_data = data.copy()
+        result_data["refining_results"] = refined_texts
+
+        return result_data
+
+    def _normalize_documents(self, docs: list[str | dict]) -> list[dict[str, Any]]:
+        """Normalize documents to isage-refiner format (with 'contents' key)."""
+        normalized: list[dict[str, Any]] = []
+        for doc in docs:
+            if isinstance(doc, dict):
+                # isage-refiner expects 'contents' key
+                text = doc.get("contents") or doc.get("text") or str(doc)
+                normalized.append({"contents": text, **doc})
+            elif isinstance(doc, str):
+                normalized.append({"contents": doc})
+            else:
+                normalized.append({"contents": str(doc)})
+
+        return normalized
+
+    def _save_data_record(self, query: str, input_docs: list[dict], refined_docs: list[str]):
+        """Save data record (only when enable_profile=True)."""
+        if not self.enable_profile:
+            return
+
+        record = {
+            "timestamp": time.time(),
+            "query": query,
+            "input_docs": input_docs,
+            "refined_docs": refined_docs,
+            "budget": self.cfg.get("budget"),
+        }
+        self.data_records.append(record)
+
+        # Persist every 10 records
+        if len(self.data_records) >= 10:
+            self._persist_data_records()
+
+    def _persist_data_records(self):
+        """Persist data records to disk."""
+        if not self.enable_profile or not self.data_records:
+            return
+
+        timestamp = int(time.time())
+        filename = f"refiner_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.logger.info(f"Saved {len(self.data_records)} records to {path}")
+            self.data_records = []
+        except Exception as e:
+            self.logger.error(f"Failed to persist data records: {e}")
+
+    def __del__(self):
+        """Ensure data is saved on cleanup."""
+        if hasattr(self, "enable_profile") and self.enable_profile:
+            try:
+                self._persist_data_records()
+            except Exception:
+                pass