realtimex-deeptutor 0.5.0.post1__py3-none-any.whl

src/services/rag/pipeline.py

@@ -0,0 +1,215 @@
+"""
+RAG Pipeline
+============
+
+Composable RAG pipeline with fluent API.
+"""
+
+import asyncio
+from pathlib import Path
+import shutil
+from typing import Any, Dict, List, Optional
+
+from src.logging import get_logger
+
+from .components.base import Component
+from .components.routing import FileTypeRouter
+from .types import Document
+
+# Default knowledge base directory
+DEFAULT_KB_BASE_DIR = str(
+    Path(__file__).resolve().parent.parent.parent.parent / "data" / "knowledge_bases"
+)
+
+
+class RAGPipeline:
+    """
+    Composable RAG pipeline.
+
+    Build custom RAG pipelines using a fluent API:
+
+        pipeline = (
+            RAGPipeline("custom", kb_base_dir="/path/to/kb")
+            .parser(PDFParser())
+            .chunker(SemanticChunker())
+            .embedder(OpenAIEmbedder())
+            .indexer(GraphIndexer())
+            .retriever(HybridRetriever())
+        )
+
+        await pipeline.initialize("kb_name", ["doc1.pdf"])
+        result = await pipeline.search("query", "kb_name")
+    """
+
+    def __init__(self, name: str = "default", kb_base_dir: Optional[str] = None):
+        """
+        Initialize RAG pipeline.
+
+        Args:
+            name: Pipeline name for logging
+            kb_base_dir: Base directory for knowledge bases
+        """
+        self.name = name
+        self.kb_base_dir = kb_base_dir or DEFAULT_KB_BASE_DIR
+        self.logger = get_logger(f"Pipeline:{name}")
+        self._parser: Optional[Component] = None
+        self._chunkers: List[Component] = []
+        self._embedder: Optional[Component] = None
+        self._indexers: List[Component] = []
+        self._retriever: Optional[Component] = None
+
+    # Fluent API methods
+    def parser(self, p: Component) -> "RAGPipeline":
+        """Set the document parser."""
+        self._parser = p
+        return self
+
+    def chunker(self, c: Component) -> "RAGPipeline":
+        """Add a chunker to the pipeline."""
+        self._chunkers.append(c)
+        return self
+
+    def embedder(self, e: Component) -> "RAGPipeline":
+        """Set the embedder."""
+        self._embedder = e
+        return self
+
+    def indexer(self, i: Component) -> "RAGPipeline":
+        """Add an indexer to the pipeline."""
+        self._indexers.append(i)
+        return self
+
+    def retriever(self, r: Component) -> "RAGPipeline":
+        """Set the retriever."""
+        self._retriever = r
+        return self
+
+    async def initialize(self, kb_name: str, file_paths: List[str], **kwargs) -> bool:
+        """
+        Run full initialization pipeline.
+
+        Uses FileTypeRouter to classify files and route them appropriately:
+        - PDF/complex files -> configured parser (e.g., PDFParser)
+        - Text files -> direct text reading (fast path)
+
+        Args:
+            kb_name: Knowledge base name
+            file_paths: List of file paths to process
+            **kwargs: Additional arguments passed to components
+
+        Returns:
+            True if successful
+        """
+        self.logger.info(f"Initializing KB '{kb_name}' with {len(file_paths)} files")
+
+        if not self._parser:
+            raise ValueError("No parser configured. Use .parser() to set one")
+
+        # Stage 1: Parse documents with file type routing
+        self.logger.info("Stage 1: Parsing documents...")
+
+        # Classify files by type
+        classification = FileTypeRouter.classify_files(file_paths)
+        self.logger.info(
+            f"File classification: {len(classification.needs_mineru)} complex, "
+            f"{len(classification.text_files)} text, "
+            f"{len(classification.unsupported)} unsupported"
+        )
+
+        documents = []
+
+        # Process complex files (PDF, etc.) with configured parser
+        for path in classification.needs_mineru:
+            self.logger.info(f"Parsing (parser): {Path(path).name}")
+            doc = await self._parser.process(path, **kwargs)
+            documents.append(doc)
+
+        # Process text files directly (fast path)
+        for path in classification.text_files:
+            self.logger.info(f"Parsing (direct text): {Path(path).name}")
+            content = await FileTypeRouter.read_text_file(path)
+            doc = Document(
+                content=content,
+                file_path=str(path),
+                metadata={
+                    "filename": Path(path).name,
+                    "parser": "direct_text",
+                },
+            )
+            documents.append(doc)
+
+        # Log unsupported files
+        for path in classification.unsupported:
+            self.logger.warning(f"Skipped unsupported file: {Path(path).name}")
+
+        # Stage 2: Chunk (sequential - later chunkers see earlier results)
+        if self._chunkers:
+            self.logger.info("Stage 2: Chunking...")
+            for chunker in self._chunkers:
+                for doc in documents:
+                    new_chunks = await chunker.process(doc, **kwargs)
+                    doc.chunks.extend(new_chunks)
+
+        # Stage 3: Embed
+        if self._embedder:
+            self.logger.info("Stage 3: Embedding...")
+            for doc in documents:
+                await self._embedder.process(doc, **kwargs)
+
+        # Stage 4: Index (can run in parallel)
+        if self._indexers:
+            self.logger.info("Stage 4: Indexing...")
+            await asyncio.gather(
+                *[indexer.process(kb_name, documents, **kwargs) for indexer in self._indexers]
+            )
+
+        self.logger.info(f"KB '{kb_name}' initialized successfully")
+        return True
+
+    async def search(self, query: str, kb_name: str, **kwargs) -> Dict[str, Any]:
+        """
+        Search the knowledge base.
+
+        Args:
+            query: Search query
+            kb_name: Knowledge base name
+            **kwargs: Additional arguments passed to retriever
+
+        Returns:
+            Search results dictionary
+        """
+        if not self._retriever:
+            raise ValueError("No retriever configured. Use .retriever() to set one")
+
+        return await self._retriever.process(query, kb_name=kb_name, **kwargs)
+
+    async def delete(self, kb_name: str) -> bool:
+        """
+        Delete a knowledge base.
+
+        Args:
+            kb_name: Knowledge base name
+
+        Returns:
+            True if successful
+        """
+        # Validate kb_name to prevent path traversal
+        if not kb_name or kb_name in (".", "..") or "/" in kb_name or "\\" in kb_name:
+            raise ValueError(f"Invalid knowledge base name: {kb_name}")
+
+        self.logger.info(f"Deleting KB '{kb_name}'")
+
+        kb_dir = Path(self.kb_base_dir) / kb_name
+        # Ensure the resolved path is within the base directory
+        kb_dir = kb_dir.resolve()
+        base_dir = Path(self.kb_base_dir).resolve()
+        if not kb_dir.is_relative_to(base_dir):
+            raise ValueError(f"Knowledge base path outside allowed directory: {kb_name}")
+
+        if kb_dir.exists():
+            shutil.rmtree(kb_dir)
+            self.logger.info(f"Deleted KB directory: {kb_dir}")
+            return True
+
+        self.logger.warning(f"KB directory not found: {kb_dir}")
+        return False

src/services/rag/pipelines/__init__.py

@@ -0,0 +1,32 @@
+"""
+Pre-configured Pipelines
+========================
+
+Ready-to-use RAG pipelines for common use cases.
+
+LightRAG and Academic pipelines are always available.
+LlamaIndex and RAGAnything require optional dependencies.
+"""
+
+# Always available pipelines
+from .academic import AcademicPipeline
+from .lightrag import LightRAGPipeline
+
+__all__ = [
+    "LightRAGPipeline",
+    "AcademicPipeline",
+]
+
+# Optional pipelines - import only if dependencies are available
+try:
+    from .llamaindex import LlamaIndexPipeline
+    __all__.append("LlamaIndexPipeline")
+except ImportError:
+    LlamaIndexPipeline = None  # type: ignore
+
+try:
+    from .raganything import RAGAnythingPipeline
+    __all__.append("RAGAnythingPipeline")
+except ImportError:
+    RAGAnythingPipeline = None  # type: ignore
+

src/services/rag/pipelines/academic.py

@@ -0,0 +1,44 @@
+"""
+Academic Pipeline
+=================
+
+Pipeline optimized for academic documents with numbered item extraction.
+"""
+
+from typing import Optional
+
+from ..components.chunkers import NumberedItemExtractor, SemanticChunker
+from ..components.embedders import OpenAIEmbedder
+from ..components.indexers import GraphIndexer
+from ..components.parsers import TextParser
+from ..components.retrievers import HybridRetriever
+from ..pipeline import RAGPipeline
+
+
+def AcademicPipeline(kb_base_dir: Optional[str] = None) -> RAGPipeline:
+    """
+    Create an academic document pipeline.
+
+    This pipeline uses:
+    - TextParser for document parsing (supports txt, md files)
+    - SemanticChunker for text chunking
+    - NumberedItemExtractor for extracting definitions, theorems, etc.
+    - OpenAIEmbedder for embedding generation
+    - GraphIndexer for knowledge graph indexing
+    - HybridRetriever for hybrid retrieval
+
+    Args:
+        kb_base_dir: Base directory for knowledge bases
+
+    Returns:
+        Configured RAGPipeline
+    """
+    return (
+        RAGPipeline("academic", kb_base_dir=kb_base_dir)
+        .parser(TextParser())
+        .chunker(SemanticChunker())
+        .chunker(NumberedItemExtractor())
+        .embedder(OpenAIEmbedder())
+        .indexer(GraphIndexer(kb_base_dir=kb_base_dir))
+        .retriever(HybridRetriever(kb_base_dir=kb_base_dir))
+    )

src/services/rag/pipelines/lightrag.py

@@ -0,0 +1,43 @@
+"""
+LightRAG Pipeline
+=================
+
+Pure LightRAG pipeline (text-only, no multimodal processing).
+Faster than RAGAnything for text-heavy documents.
+"""
+
+from typing import Optional
+
+from ..components.indexers import LightRAGIndexer
+from ..components.parsers import PDFParser
+from ..components.retrievers import LightRAGRetriever
+from ..pipeline import RAGPipeline
+
+
+def LightRAGPipeline(kb_base_dir: Optional[str] = None) -> RAGPipeline:
+    """
+    Create a pure LightRAG pipeline (text-only, no multimodal).
+
+    This pipeline uses:
+    - PDFParser for document parsing (extracts raw text from PDF/txt/md)
+    - LightRAGIndexer for knowledge graph indexing (text-only, fast)
+      * LightRAG handles chunking, entity extraction, and embedding internally
+      * No separate chunker/embedder needed - LightRAG does it all
+    - LightRAGRetriever for retrieval (uses LightRAG.aquery() directly)
+
+    Performance: Medium speed (~10-15s per document)
+    Use for: Business docs, text-heavy PDFs, when you need knowledge graph
+
+    Args:
+        kb_base_dir: Base directory for knowledge bases
+
+    Returns:
+        Configured RAGPipeline
+    """
+    return (
+        RAGPipeline("lightrag", kb_base_dir=kb_base_dir)
+        .parser(PDFParser())
+        # No chunker/embedder - LightRAG does everything internally
+        .indexer(LightRAGIndexer(kb_base_dir=kb_base_dir))
+        .retriever(LightRAGRetriever(kb_base_dir=kb_base_dir))
+    )

src/services/rag/pipelines/llamaindex.py

@@ -0,0 +1,313 @@
+"""
+LlamaIndex Pipeline
+===================
+
+True LlamaIndex integration using official llama-index library.
+"""
+
+import asyncio
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from llama_index.core import (
+    Document,
+    Settings,
+    StorageContext,
+    VectorStoreIndex,
+    load_index_from_storage,
+)
+from llama_index.core.base.embeddings.base import BaseEmbedding
+from llama_index.core.bridge.pydantic import PrivateAttr
+
+from src.logging import get_logger
+from src.services.embedding import get_embedding_client, get_embedding_config
+
+# Default knowledge base directory
+DEFAULT_KB_BASE_DIR = str(
+    Path(__file__).resolve().parent.parent.parent.parent.parent / "data" / "knowledge_bases"
+)
+
+
+class CustomEmbedding(BaseEmbedding):
+    """
+    Custom embedding adapter for OpenAI-compatible APIs.
+
+    Works with any OpenAI-compatible endpoint including:
+    - Google Gemini (text-embedding-004)
+    - OpenAI (text-embedding-ada-002, text-embedding-3-*)
+    - Azure OpenAI
+    - Local models with OpenAI-compatible API
+    """
+
+    _client: Any = PrivateAttr()
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._client = get_embedding_client()
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "custom_embedding"
+
+    async def _aget_query_embedding(self, query: str) -> List[float]:
+        """Get embedding for a query."""
+        embeddings = await self._client.embed([query])
+        return embeddings[0]
+
+    async def _aget_text_embedding(self, text: str) -> List[float]:
+        """Get embedding for a text."""
+        embeddings = await self._client.embed([text])
+        return embeddings[0]
+
+    def _get_query_embedding(self, query: str) -> List[float]:
+        """Sync version - called by LlamaIndex sync API."""
+        # Use nest_asyncio to allow nested event loops
+        import nest_asyncio
+
+        nest_asyncio.apply()
+        return asyncio.run(self._aget_query_embedding(query))
+
+    def _get_text_embedding(self, text: str) -> List[float]:
+        """Sync version - called by LlamaIndex sync API."""
+        # Use nest_asyncio to allow nested event loops
+        import nest_asyncio
+
+        nest_asyncio.apply()
+        return asyncio.run(self._aget_text_embedding(text))
+
+    async def _aget_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """Get embeddings for multiple texts."""
+        return await self._client.embed(texts)
+
+
+class LlamaIndexPipeline:
+    """
+    True LlamaIndex pipeline using official llama-index library.
+
+    Uses LlamaIndex's native components:
+    - VectorStoreIndex for indexing
+    - CustomEmbedding for OpenAI-compatible embeddings
+    - SentenceSplitter for chunking
+    - StorageContext for persistence
+    """
+
+    def __init__(self, kb_base_dir: Optional[str] = None):
+        """
+        Initialize LlamaIndex pipeline.
+
+        Args:
+            kb_base_dir: Base directory for knowledge bases
+        """
+        self.logger = get_logger("LlamaIndexPipeline")
+        self.kb_base_dir = kb_base_dir or DEFAULT_KB_BASE_DIR
+        self._configure_settings()
+
+    def _configure_settings(self):
+        """Configure LlamaIndex global settings."""
+        # Get embedding config
+        embedding_cfg = get_embedding_config()
+
+        # Configure custom embedding that works with any OpenAI-compatible API
+        Settings.embed_model = CustomEmbedding()
+
+        # Configure chunking
+        Settings.chunk_size = 512
+        Settings.chunk_overlap = 50
+
+        self.logger.info(
+            f"LlamaIndex configured: embedding={embedding_cfg.model} "
+            f"({embedding_cfg.dim}D, {embedding_cfg.binding}), chunk_size=512"
+        )
+
+    async def initialize(self, kb_name: str, file_paths: List[str], **kwargs) -> bool:
+        """
+        Initialize KB using real LlamaIndex components.
+
+        Args:
+            kb_name: Knowledge base name
+            file_paths: List of file paths to process
+            **kwargs: Additional arguments
+
+        Returns:
+            True if successful
+        """
+        self.logger.info(
+            f"Initializing KB '{kb_name}' with {len(file_paths)} files using LlamaIndex"
+        )
+
+        kb_dir = Path(self.kb_base_dir) / kb_name
+        storage_dir = kb_dir / "llamaindex_storage"
+        storage_dir.mkdir(parents=True, exist_ok=True)
+
+        try:
+            # Parse documents
+            documents = []
+            for file_path in file_paths:
+                file_path = Path(file_path)
+                self.logger.info(f"Parsing: {file_path.name}")
+
+                # Extract text based on file type
+                if file_path.suffix.lower() == ".pdf":
+                    text = self._extract_pdf_text(file_path)
+                else:
+                    with open(file_path, "r", encoding="utf-8") as f:
+                        text = f.read()
+
+                if text.strip():
+                    doc = Document(
+                        text=text,
+                        metadata={
+                            "file_name": file_path.name,
+                            "file_path": str(file_path),
+                        },
+                    )
+                    documents.append(doc)
+                    self.logger.info(f"Loaded: {file_path.name} ({len(text)} chars)")
+                else:
+                    self.logger.warning(f"Skipped empty document: {file_path.name}")
+
+            if not documents:
+                self.logger.error("No valid documents found")
+                return False
+
+            # Create index with LlamaIndex (run sync code in thread pool to avoid blocking)
+            self.logger.info(f"Creating VectorStoreIndex with {len(documents)} documents...")
+
+            # Run sync LlamaIndex code in thread pool to avoid blocking async event loop
+            loop = asyncio.get_event_loop()
+            index = await loop.run_in_executor(
+                None,  # Use default ThreadPoolExecutor
+                lambda: VectorStoreIndex.from_documents(documents, show_progress=True),
+            )
+
+            # Persist index
+            index.storage_context.persist(persist_dir=str(storage_dir))
+            self.logger.info(f"Index persisted to {storage_dir}")
+
+            self.logger.info(f"KB '{kb_name}' initialized successfully with LlamaIndex")
+            return True
+
+        except Exception as e:
+            self.logger.error(f"Failed to initialize KB: {e}")
+            import traceback
+
+            self.logger.error(traceback.format_exc())
+            return False
+
+    def _extract_pdf_text(self, file_path: Path) -> str:
+        """Extract text from PDF using PyMuPDF."""
+        try:
+            import fitz  # PyMuPDF
+
+            doc = fitz.open(file_path)
+            texts = []
+            for page in doc:
+                texts.append(page.get_text())
+            doc.close()
+            return "\n\n".join(texts)
+        except ImportError:
+            self.logger.warning("PyMuPDF not installed. Cannot extract PDF text.")
+            return ""
+        except Exception as e:
+            self.logger.error(f"Failed to extract PDF text: {e}")
+            return ""
+
+    async def search(
+        self,
+        query: str,
+        kb_name: str,
+        mode: str = "hybrid",
+        **kwargs,
+    ) -> Dict[str, Any]:
+        """
+        Search using LlamaIndex query engine.
+
+        Args:
+            query: Search query
+            kb_name: Knowledge base name
+            mode: Search mode (ignored, LlamaIndex uses similarity)
+            **kwargs: Additional arguments (top_k, etc.)
+
+        Returns:
+            Search results dictionary
+        """
+        self.logger.info(f"Searching KB '{kb_name}' with query: {query[:50]}...")
+
+        kb_dir = Path(self.kb_base_dir) / kb_name
+        storage_dir = kb_dir / "llamaindex_storage"
+
+        if not storage_dir.exists():
+            self.logger.warning(f"No LlamaIndex storage found at {storage_dir}")
+            return {
+                "query": query,
+                "answer": "No documents indexed. Please upload documents first.",
+                "content": "",
+                "mode": mode,
+                "provider": "llamaindex",
+            }
+
+        try:
+            # Load index from storage (run in thread pool)
+            loop = asyncio.get_event_loop()
+
+            def load_and_retrieve():
+                storage_context = StorageContext.from_defaults(persist_dir=str(storage_dir))
+                index = load_index_from_storage(storage_context)
+                top_k = kwargs.get("top_k", 5)
+
+                # Use retriever instead of query_engine to avoid LLM requirement
+                retriever = index.as_retriever(similarity_top_k=top_k)
+                nodes = retriever.retrieve(query)
+                return nodes
+
+            # Execute retrieval in thread pool to avoid blocking
+            nodes = await loop.run_in_executor(None, load_and_retrieve)
+
+            # Extract text from retrieved nodes
+            context_parts = []
+            for node in nodes:
+                context_parts.append(node.node.text)
+
+            content = "\n\n".join(context_parts) if context_parts else ""
+
+            return {
+                "query": query,
+                "answer": content,  # Return context for ChatAgent to use
+                "content": content,
+                "mode": mode,
+                "provider": "llamaindex",
+            }
+
+        except Exception as e:
+            self.logger.error(f"Search failed: {e}")
+            import traceback
+
+            self.logger.error(traceback.format_exc())
+            return {
+                "query": query,
+                "answer": f"Search failed: {str(e)}",
+                "content": "",
+                "mode": mode,
+                "provider": "llamaindex",
+            }
+
+    async def delete(self, kb_name: str) -> bool:
+        """
+        Delete knowledge base.
+
+        Args:
+            kb_name: Knowledge base name
+
+        Returns:
+            True if successful
+        """
+        import shutil
+
+        kb_dir = Path(self.kb_base_dir) / kb_name
+
+        if kb_dir.exists():
+            shutil.rmtree(kb_dir)
+            self.logger.info(f"Deleted KB '{kb_name}'")
+            return True
+
+        return False