isage-middleware 0.2.4.3__cp311-none-any.whl

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

@@ -0,0 +1,363 @@
+"""Index Builder - Service for building RAG vector indices
+
+Layer: L4 (sage-middleware/operators/rag)
+"""
+
+import logging
+from collections.abc import Callable
+from contextlib import contextmanager
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from sage.middleware.operators.rag.index_builder.manifest import IndexManifest
+from sage.middleware.operators.rag.index_builder.storage import VectorStore
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def _optional_progress(show: bool, description: str, total: int | None = None):
+    """Context manager for optional Rich progress bar.
+
+    Args:
+        show: Whether to show progress bar (False = silent mode)
+        description: Task description
+        total: Total number of items (None for indeterminate)
+
+    Yields:
+        Progress task update function: update(advance=1)
+    """
+    if not show:
+        # Silent mode - yield a no-op update function
+        def noop(**kwargs):
+            pass
+
+        yield noop
+        return
+
+    try:
+        from rich.progress import (
+            BarColumn,
+            Progress,
+            SpinnerColumn,
+            TaskProgressColumn,
+            TextColumn,
+            TimeRemainingColumn,
+        )
+
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[cyan]{task.description}"),
+            BarColumn(bar_width=30),
+            TaskProgressColumn(),
+            TimeRemainingColumn(),
+            transient=True,  # Clear after completion
+        ) as progress:
+            task = progress.add_task(description, total=total)
+
+            def update(advance: int = 1, **kwargs):
+                progress.update(task, advance=advance, **kwargs)
+
+            yield update
+
+    except ImportError:
+        # Fallback if rich is not available
+        logger.info(f"[Progress] {description}")
+
+        def fallback_update(**kwargs):
+            pass
+
+        yield fallback_update
+
+
+class IndexBuilder:
+    """Service for building RAG vector indices with pluggable backends.
+
+    This class orchestrates the complete index building workflow, using
+    dependency injection to decouple from specific vector storage backends.
+
+    Architecture Pattern:
+        - L4 defines this builder (orchestration logic)
+        - L4 provides SageDB backend (sage.middleware.components.sage_db)
+        - L3 provides ChromaDB backend (sage.libs.integrations.chroma)
+        - L5 uses IndexBuilder with injected backend factory
+
+    Args:
+        backend_factory: Function creating VectorStore instances
+            Signature: (persist_path: Path, dim: int) -> VectorStore
+
+    Example:
+        >>> # In sage-cli (L5)
+        >>> from sage.middleware.operators.rag.index_builder import IndexBuilder
+        >>> from sage.middleware.components.sage_db import SageVDBBackend
+        >>>
+        >>> def factory(path: Path, dim: int):
+        ...     return SageVDBBackend(path, dim)
+        >>>
+        >>> builder = IndexBuilder(backend_factory=factory)
+        >>> manifest = builder.build_from_docs(
+        ...     source_dir=Path("docs"),
+        ...     persist_path=Path(".sage/db"),
+        ...     embedding_model=embedder,
+        ...     chunk_size=800,
+        ...     chunk_overlap=160,
+        ... )
+    """
+
+    def __init__(self, backend_factory: Callable[[Path, int], VectorStore]):
+        """Initialize builder with backend factory.
+
+        Args:
+            backend_factory: Factory function for creating VectorStore instances
+        """
+        self.backend_factory = backend_factory
+
+    def build_from_docs(
+        self,
+        source_dir: Path,
+        persist_path: Path,
+        embedding_model: Any,
+        index_name: str = "default",
+        chunk_size: int = 800,
+        chunk_overlap: int = 160,
+        document_processor: Callable[[Path], list[dict[str, Any]]] | None = None,
+        max_documents: int | None = None,
+        show_progress: bool = True,
+    ) -> IndexManifest:
+        """Build vector index from document directory.
+
+        This method orchestrates the complete index building process:
+        1. Create vector store backend
+        2. Process documents (via document_processor or default)
+        3. Chunk text content
+        4. Generate embeddings
+        5. Store vectors with metadata
+        6. Build/optimize index
+        7. Persist to disk
+        8. Return manifest
+
+        Args:
+            source_dir: Directory containing source documents
+            persist_path: Path to save the built index
+            embedding_model: Model with embed() and get_dim() methods
+            index_name: Unique identifier for this index
+            chunk_size: Size of text chunks in characters
+            chunk_overlap: Overlap between consecutive chunks
+            document_processor: Optional custom document processing function
+                If None, uses simple text extraction
+                Signature: (source_dir: Path) -> list[dict] where dict has:
+                    - "content": str (text content)
+                    - "metadata": dict (doc_path, title, heading, etc.)
+            max_documents: Optional limit on number of documents to process
+            show_progress: Show Rich progress bar (False = quiet mode)
+
+        Returns:
+            IndexManifest with build statistics and metadata
+
+        Raises:
+            FileNotFoundError: If source_dir doesn't exist
+            RuntimeError: If index building fails
+
+        Example:
+            >>> # Custom document processor for Markdown
+            >>> def process_markdown(source_dir: Path):
+            ...     chunks = []
+            ...     for file in source_dir.glob("**/*.md"):
+            ...         text = file.read_text()
+            ...         chunks.append({
+            ...             "content": text,
+            ...             "metadata": {"doc_path": str(file.relative_to(source_dir))}
+            ...         })
+            ...     return chunks
+            >>>
+            >>> manifest = builder.build_from_docs(
+            ...     source_dir=Path("docs"),
+            ...     persist_path=Path(".sage/db"),
+            ...     embedding_model=embedder,
+            ...     document_processor=process_markdown,
+            ... )
+        """
+        if not source_dir.exists():
+            raise FileNotFoundError(f"Source directory not found: {source_dir}")
+
+        logger.debug(f"Building index from {source_dir}")
+        logger.debug(f"Backend: {self.backend_factory}")
+        logger.debug(f"Chunk size: {chunk_size}, overlap: {chunk_overlap}")
+
+        # Create vector store backend
+        dim = embedding_model.get_dim()
+        store = self.backend_factory(persist_path, dim)
+        logger.debug(f"Created vector store with dimension {dim}")
+
+        # Process documents
+        if document_processor is None:
+            # Default: simple text file processing
+            logger.debug(
+                "No document_processor provided, using default text extraction. "
+                "For better results, provide a custom processor."
+            )
+            processed_docs = self._default_document_processor(source_dir, max_documents)
+        else:
+            processed_docs = document_processor(source_dir)
+            if max_documents:
+                processed_docs = processed_docs[:max_documents]
+
+        logger.debug(f"Processed {len(processed_docs)} document sections")
+
+        # Import chunking utility
+        try:
+            from sage.common.utils.document_processing import (
+                chunk_text,
+                sanitize_metadata_value,
+                truncate_text,
+            )
+        except ImportError:
+            logger.debug("Cannot import chunking utilities from sage.common, using simple split")
+
+            def chunk_text(text: str, size: int, overlap: int) -> list[str]:
+                # Fallback: simple fixed-size chunking
+                chunks = []
+                start = 0
+                while start < len(text):
+                    end = min(len(text), start + size)
+                    chunks.append(text[start:end])
+                    start += size - overlap
+                return chunks
+
+            def sanitize_metadata_value(val: str) -> str:
+                # Remove problematic chars for JSON/C++ parser
+                return (
+                    val.replace("\\", "")
+                    .replace("\n", " ")
+                    .replace('"', "'")
+                    .replace("{", "(")
+                    .replace("}", ")")
+                )
+
+            def truncate_text(text: str, limit: int = 480) -> str:
+                return text[:limit] if len(text) > limit else text
+
+        # Embed and store (with chunking)
+        # First pass: count total chunks for accurate progress
+        all_chunks_data = []  # List of (chunk_text, base_metadata)
+        unique_docs = set()
+
+        for doc in processed_docs:
+            content = doc["content"]
+            base_metadata = doc["metadata"]
+
+            # Track unique documents
+            if "doc_path" in base_metadata:
+                unique_docs.add(base_metadata["doc_path"])
+
+            # Chunk the content
+            content_chunks = chunk_text(content, chunk_size, chunk_overlap)
+
+            for chunk_idx, chunk in enumerate(content_chunks):
+                all_chunks_data.append((chunk, base_metadata, chunk_idx))
+
+        total_chunks = len(all_chunks_data)
+        logger.debug(f"Total chunks to embed: {total_chunks}")
+
+        # Second pass: embed with accurate progress
+        with _optional_progress(show_progress, "Embedding", total=total_chunks) as progress_update:
+            for idx, (chunk, base_metadata, chunk_idx) in enumerate(all_chunks_data, start=1):
+                # Generate embedding
+                vector = embedding_model.embed(chunk)
+
+                # Create metadata for this chunk
+                metadata = {
+                    **base_metadata,
+                    "chunk": str(chunk_idx),
+                    "text": sanitize_metadata_value(truncate_text(chunk, limit=1200)),
+                }
+
+                # Sanitize all string values
+                metadata = {
+                    k: sanitize_metadata_value(str(v)) if isinstance(v, str) else str(v)
+                    for k, v in metadata.items()
+                }
+
+                # Store vector with metadata
+                store.add(vector, metadata)
+                progress_update(advance=1)
+
+                if idx % 500 == 0:
+                    logger.debug(f"Embedded {idx}/{total_chunks} chunks")
+
+        logger.debug(f"Added {total_chunks} vectors from {len(unique_docs)} documents")
+
+        # Build index
+        logger.debug("Building vector index...")
+        store.build_index()
+
+        # Persist to disk
+        logger.debug(f"Saving index to {persist_path}")
+        store.save(str(persist_path))
+
+        # Create manifest
+        manifest = IndexManifest(
+            index_name=index_name,
+            backend_type=type(store).__name__,
+            persist_path=persist_path,
+            source_dir=str(source_dir),
+            embedding_config={
+                "model": type(embedding_model).__name__,
+                "dim": dim,
+            },
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            num_documents=len(unique_docs),
+            num_chunks=total_chunks,
+            created_at=datetime.utcnow().isoformat(),
+        )
+
+        logger.debug(f"Index built successfully: {manifest}")
+        return manifest
+
+    def _default_document_processor(
+        self,
+        source_dir: Path,
+        max_documents: int | None = None,
+    ) -> list[dict[str, Any]]:
+        """Default document processor for plain text files.
+
+        This is a fallback processor that simply reads text files.
+        For production use, provide a custom processor that:
+        - Handles specific formats (Markdown, PDF, etc.)
+        - Implements smart chunking
+        - Preserves document structure
+
+        Args:
+            source_dir: Directory to scan
+            max_documents: Optional limit
+
+        Returns:
+            List of processed chunks with metadata
+        """
+        chunks = []
+        text_files = list(source_dir.glob("**/*.txt")) + list(source_dir.glob("**/*.md"))
+
+        if max_documents:
+            text_files = text_files[:max_documents]
+
+        for file_path in text_files:
+            try:
+                content = file_path.read_text(encoding="utf-8", errors="ignore")
+                rel_path = file_path.relative_to(source_dir)
+
+                chunks.append(
+                    {
+                        "content": content,
+                        "metadata": {
+                            "doc_path": str(rel_path),
+                            "title": file_path.stem,
+                            "text": content[:1000],  # Preview
+                        },
+                    }
+                )
+            except Exception as e:
+                logger.warning(f"Failed to process {file_path}: {e}")
+
+        return chunks

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

@@ -0,0 +1,101 @@
+"""Index Manifest - Metadata describing a built RAG index
+
+Layer: L4 (sage-middleware/operators/rag)
+"""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+
+@dataclass
+class IndexManifest:
+    """Metadata describing a built knowledge index.
+
+    This dataclass stores comprehensive metadata about a vector index,
+    including source information, embedding configuration, and statistics.
+
+    Attributes:
+        index_name: Unique identifier for this index
+        backend_type: Storage backend ("sagedb", "chromadb", "milvus", etc.)
+        persist_path: Path where the index is stored
+        source_dir: Original document directory
+        embedding_config: Embedding model configuration
+        chunk_size: Size of text chunks in characters
+        chunk_overlap: Overlap between chunks in characters
+        num_documents: Total number of documents indexed
+        num_chunks: Total number of text chunks (vectors) stored
+        created_at: ISO timestamp of index creation
+        metadata: Additional custom metadata
+
+    Example:
+        >>> manifest = IndexManifest(
+        ...     index_name="docs-public",
+        ...     backend_type="chromadb",
+        ...     persist_path=Path(".sage/vector_db"),
+        ...     source_dir="docs-public/docs_src",
+        ...     embedding_config={"method": "hash", "dim": 384},
+        ...     chunk_size=800,
+        ...     chunk_overlap=160,
+        ...     num_documents=124,
+        ...     num_chunks=2720,
+        ...     created_at=datetime.utcnow().isoformat(),
+        ... )
+    """
+
+    index_name: str
+    backend_type: str
+    persist_path: Path
+    source_dir: str
+    embedding_config: dict[str, Any]
+    chunk_size: int
+    chunk_overlap: int
+    num_documents: int
+    num_chunks: int
+    created_at: str
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert manifest to dictionary for serialization."""
+        return {
+            "index_name": self.index_name,
+            "backend_type": self.backend_type,
+            "persist_path": str(self.persist_path),
+            "source_dir": self.source_dir,
+            "embedding_config": self.embedding_config,
+            "chunk_size": self.chunk_size,
+            "chunk_overlap": self.chunk_overlap,
+            "num_documents": self.num_documents,
+            "num_chunks": self.num_chunks,
+            "created_at": self.created_at,
+            "metadata": self.metadata,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "IndexManifest":
+        """Create manifest from dictionary."""
+        data_copy = data.copy()
+        data_copy["persist_path"] = Path(data_copy["persist_path"])
+        return cls(**data_copy)
+
+    @property
+    def age_seconds(self) -> float:
+        """Get age of index in seconds since creation."""
+        created = datetime.fromisoformat(self.created_at)
+        return (datetime.utcnow() - created).total_seconds()
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if index contains any data."""
+        return self.num_chunks == 0
+
+    def __repr__(self) -> str:
+        """Readable representation of manifest."""
+        return (
+            f"IndexManifest("
+            f"name={self.index_name!r}, "
+            f"backend={self.backend_type!r}, "
+            f"docs={self.num_documents}, "
+            f"chunks={self.num_chunks})"
+        )

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

@@ -0,0 +1,131 @@
+"""Vector Store Protocol - Abstract interface for vector storage backends
+
+Layer: L4 (sage-middleware/operators/rag)
+"""
+
+from typing import Any, Protocol, runtime_checkable
+
+
+@runtime_checkable
+class VectorStore(Protocol):
+    """Abstract interface for vector storage backends.
+
+    This Protocol defines the contract that all vector storage implementations
+    must satisfy. It enables dependency injection and backend swapping without
+    tight coupling to specific implementations (SageDB, ChromaDB, Milvus, etc.).
+
+    Architecture Pattern:
+        - L4 (sage-middleware): Defines this Protocol + SageDB implementation
+        - L3 (sage-libs/integrations): Provides ChromaDB implementation
+        - L5 (sage-cli): Uses via factory injection
+
+    Example Implementation:
+        >>> class SageVDBBackend:
+        ...     def __init__(self, persist_path: Path, dim: int):
+        ...         from sage.middleware.components.sage_db import SageDB
+        ...         self.db = SageDB(dim)
+        ...         self.path = persist_path
+        ...
+        ...     def add(self, vector: list[float], metadata: dict) -> None:
+        ...         self.db.add(vector, metadata)
+        ...
+        ...     # ... implement other methods
+
+    Usage with IndexBuilder:
+        >>> def backend_factory(path: Path, dim: int) -> VectorStore:
+        ...     return SageDBBackend(path, dim)
+        >>>
+        >>> builder = IndexBuilder(backend_factory=backend_factory)
+    """
+
+    def add(self, vector: list[float], metadata: dict[str, Any]) -> None:
+        """Add a single vector with metadata to the store.
+
+        Args:
+            vector: Dense vector embedding (must match dimension)
+            metadata: Associated metadata (doc_path, title, heading, chunk, text, etc.)
+
+        Raises:
+            ValueError: If vector dimension doesn't match
+        """
+        ...
+
+    def build_index(self) -> None:
+        """Build/optimize the vector index for efficient search.
+
+        This is typically called after all vectors are added via `add()`.
+        Implementations may use various indexing strategies:
+        - Flat index (brute force)
+        - HNSW (Hierarchical Navigable Small World)
+        - IVF (Inverted File Index)
+        - Product Quantization
+
+        Raises:
+            RuntimeError: If index building fails
+        """
+        ...
+
+    def save(self, path: str) -> None:
+        """Persist the index to disk.
+
+        Args:
+            path: Absolute path to save location
+
+        Raises:
+            IOError: If save fails
+        """
+        ...
+
+    def load(self, path: str) -> None:
+        """Load a previously saved index from disk.
+
+        Args:
+            path: Absolute path to load from
+
+        Raises:
+            FileNotFoundError: If index doesn't exist
+            IOError: If load fails
+        """
+        ...
+
+    def search(
+        self,
+        query_vector: list[float],
+        top_k: int = 5,
+        filter_metadata: dict[str, Any] | None = None,
+    ) -> list[dict[str, Any]]:
+        """Search for nearest neighbor vectors.
+
+        Args:
+            query_vector: Query embedding
+            top_k: Number of results to return
+            filter_metadata: Optional metadata filters (e.g., {"doc_path": "intro.md"})
+
+        Returns:
+            List of results, each containing:
+                - vector: The matched vector
+                - metadata: Associated metadata
+                - distance/score: Similarity score
+
+        Example:
+            >>> results = store.search([0.1, 0.2, ...], top_k=5)
+            >>> for result in results:
+            ...     print(result["metadata"]["title"], result["score"])
+        """
+        ...
+
+    def get_dim(self) -> int:
+        """Get the vector dimension of this store.
+
+        Returns:
+            Vector dimension (e.g., 384 for BGE-small, 768 for BERT)
+        """
+        ...
+
+    def count(self) -> int:
+        """Get total number of vectors in the store.
+
+        Returns:
+            Total vector count
+        """
+        ...

sage/middleware/operators/rag/pipeline.py

@@ -0,0 +1,46 @@
+"""
+RAG Pipeline - RAG 系统的核心管道组件
+
+Layer: L4 (Middleware - Orchestration)
+This module orchestrates multiple RAG components (retriever, reranker, refiner, generator)
+into a cohesive pipeline. Pipeline/orchestration belongs in middleware, not libs.
+"""
+
+from typing import Any
+
+
+class RAGPipeline:
+    """RAG 管道主类 - 编排多个RAG组件"""
+
+    def __init__(self, retriever=None, generator=None, reranker=None, refiner=None):
+        self.retriever = retriever
+        self.generator = generator
+        self.reranker = reranker
+        self.refiner = refiner
+
+    def run(self, query: str, **kwargs) -> dict[str, Any]:
+        """运行 RAG 管道"""
+        # 1. 检索相关文档
+        if self.retriever:
+            documents = self.retriever.retrieve(query, **kwargs)
+        else:
+            documents = []
+
+        # 2. 重排序（可选）
+        if self.reranker and documents:
+            documents = self.reranker.rerank(query, documents, **kwargs)
+
+        # 3. 精化查询或文档（可选）
+        if self.refiner:
+            query, documents = self.refiner.refine(query, documents, **kwargs)
+
+        # 4. 生成回答
+        if self.generator:
+            response = self.generator.generate(query, documents, **kwargs)
+        else:
+            response = "No generator configured"
+
+        return {"query": query, "documents": documents, "response": response}
+
+
+__all__ = ["RAGPipeline"]

sage/middleware/operators/rag/profiler.py

@@ -0,0 +1,59 @@
+import json
+from dataclasses import dataclass
+
+from sage.common.core import FilterFunction
+
+
+@dataclass
+class QueryProfilerResult:
+    need_joint_reasoning: bool
+    complexity: str  # "High" or "Low"
+    need_summarization: bool
+    summarization_length: int  # 30-200
+    n_info_items: int  # 1-6
+
+    def __post_init__(self):
+        # 严格验证，抛出异常
+        if self.complexity not in ["High", "Low"]:
+            raise ValueError(f"complexity必须是'High'或'Low'，得到: {self.complexity}")
+
+        if not (30 <= self.summarization_length <= 200):
+            raise ValueError(
+                f"summarization_length必须在30-200之间，得到: {self.summarization_length}"
+            )
+
+        if not (1 <= self.n_info_items <= 6):
+            raise ValueError(f"n_info_items必须在1-6之间，得到: {self.n_info_items}")
+
+
+class Query_Profiler(FilterFunction):
+    def __init__(self, config, **kwargs):
+        super().__init__(**kwargs)
+
+    def execute(self, data):
+        js = json.loads(data)
+        # 使用解包创建对象并直接获取属性
+        profiler_result = QueryProfilerResult(
+            need_joint_reasoning=js.get("need_joint_reasoning", False),
+            complexity=js.get("complexity", "Low"),
+            need_summarization=js.get("need_summarization", False),
+            summarization_length=js.get("summarization_length", 30),
+            n_info_items=js.get("n_info_items", 1),
+        )
+
+        # 直接解包到变量
+        need_joint_reasoning = profiler_result.need_joint_reasoning
+        complexity = profiler_result.complexity
+        summarization_length = profiler_result.summarization_length
+
+        if need_joint_reasoning is False:
+            synthesis_method = "map_rerank"
+        else:
+            if complexity == "Low":
+                synthesis_method = "stuff"
+            else:
+                synthesis_method = "map_reduce"
+        num_chunks = [profiler_result.n_info_items, 3 * profiler_result.n_info_items]
+        intermediate_length_range = summarization_length
+
+        return [synthesis_method, num_chunks, intermediate_length_range]