deeptutor-knowledge 1.7.0__py3-none-any.whl

deeptutor/knowledge/initializer.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python
+"""Knowledge base initialization through the selected RAG provider."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+from datetime import datetime
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from deeptutor.multi_user.models import CurrentUser
+    from deeptutor.services.rag.pipelines.lightrag.indexing_policy import IndexingPolicySnapshot
+
+from deeptutor.knowledge.naming import validate_knowledge_base_name
+from deeptutor.knowledge.progress_tracker import ProgressStage, ProgressTracker
+from deeptutor.knowledge.staging import stage_source
+from deeptutor.services.config import resolve_llm_runtime_config
+from deeptutor.services.file_io import atomic_write_json
+from deeptutor.services.rag.factory import normalize_provider_name
+from deeptutor.services.rag.service import RAGService
+from deeptutor.services.setup.data_volume import (
+    DataVolumePermissionError,
+    ensure_data_volume_writable,
+    format_data_volume_permission_error,
+)
+from deeptutor.utils.file_types import FileTypeRouter
+
+logger = logging.getLogger(__name__)
+
+
+class KnowledgeBaseInitializer:
+    """Knowledge base initializer."""
+
+    def __init__(
+        self,
+        kb_name: str,
+        base_dir: str = "./data/knowledge_bases",
+        api_key: str | None = None,
+        base_url: str | None = None,
+        progress_tracker: ProgressTracker | None = None,
+        rag_provider: str | None = None,
+        indexing_snapshot: IndexingPolicySnapshot | None = None,
+        owner: CurrentUser | None = None,
+    ):
+        self.kb_name = validate_knowledge_base_name(kb_name)
+        self.base_dir = Path(base_dir)
+        self.kb_dir = self.base_dir / self.kb_name
+
+        self.raw_dir = self.kb_dir / "raw"
+        self.llamaindex_storage_dir = self.kb_dir / "llamaindex_storage"
+
+        self.api_key = api_key
+        self.base_url = base_url
+        self.progress_tracker = progress_tracker or ProgressTracker(self.kb_name, self.base_dir)
+        self.rag_provider = normalize_provider_name(rag_provider)
+        self.indexing_snapshot = indexing_snapshot
+        self.owner = owner
+        self.index_published = False
+
+    def _register_to_config(self) -> None:
+        """Register KB in kb_config.json with initializing state."""
+        try:
+            from deeptutor.knowledge.manager import KnowledgeBaseManager
+
+            manager = KnowledgeBaseManager(base_dir=str(self.base_dir))
+            manager.config = manager._load_config()
+            if self.kb_name in manager.config.get("knowledge_bases", {}):
+                return
+
+            manager.update_kb_status(
+                name=self.kb_name,
+                status="initializing",
+                progress={
+                    "stage": "initializing",
+                    "message": "Creating directory structure...",
+                    "percent": 0,
+                    "current": 0,
+                    "total": 0,
+                },
+            )
+            manager.config = manager._load_config()
+            manager.config.setdefault("knowledge_bases", {}).setdefault(self.kb_name, {})[
+                "rag_provider"
+            ] = self.rag_provider
+            manager._save_config()
+        except Exception as e:
+            logger.warning(f"Failed to register KB to config: {e}")
+
+    def _update_metadata_with_provider(self, provider: str) -> None:
+        metadata_file = self.kb_dir / "metadata.json"
+        metadata: dict = {}
+        if metadata_file.exists():
+            try:
+                with open(metadata_file, "r", encoding="utf-8") as f:
+                    metadata = json.load(f)
+            except Exception:
+                metadata = {}
+
+        metadata["rag_provider"] = provider
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        metadata["last_updated"] = timestamp
+        metadata["last_indexed_at"] = timestamp
+        metadata["last_indexed_count"] = len(
+            FileTypeRouter.collect_supported_files(self.raw_dir, recursive=True)
+        )
+        if provider == "llamaindex":
+            from deeptutor.knowledge.indexing_run import load_run
+
+            receipt = load_run(self.kb_dir)
+            if receipt:
+                metadata["last_indexed_count"] = sum(
+                    doc.get("status") == "completed" for doc in receipt["documents"].values()
+                )
+                metadata["indexing_state"] = receipt["state"]
+        metadata["last_indexed_action"] = "create"
+
+        atomic_write_json(metadata_file, metadata)
+
+        try:
+            from deeptutor.services.config import get_kb_config_service
+
+            service = get_kb_config_service()
+            service.set_rag_provider(self.kb_name, provider)
+            service.set_kb_config(self.kb_name, {"needs_reindex": False})
+        except Exception as config_err:
+            logger.warning(f"Failed to persist provider in centralized config: {config_err}")
+
+    def create_directory_structure(self) -> None:
+        """Create KB directory structure."""
+        logger.info(f"Creating directory structure for knowledge base: {self.kb_name}")
+        ensure_data_volume_writable(self.base_dir)
+
+        for dir_path in [
+            self.raw_dir,
+        ]:
+            dir_path.mkdir(parents=True, exist_ok=True)
+
+        metadata = {
+            "name": self.kb_name,
+            "created_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "description": f"Knowledge base: {self.kb_name}",
+            "version": "1.0",
+            "rag_provider": self.rag_provider,
+            "needs_reindex": False,
+        }
+
+        atomic_write_json(self.kb_dir / "metadata.json", metadata)
+
+        self._register_to_config()
+
+    def copy_documents(self, source_files: list[str]) -> list[str]:
+        """Stage source documents into the raw directory.
+
+        Sources arrive flattened to their basename, so a docs tree holding one
+        ``full.md`` per volume wants the same path over and over. Each document
+        is kept under a distinct name instead of overwriting the one before it
+        (#2041) — the same rule the incremental add path already followed.
+        """
+        copied_files: list[str] = []
+        staged: set[Path] = set()
+        for source in source_files:
+            source_path = Path(source)
+            if not source_path.exists() or not source_path.is_file():
+                logger.warning(f"Source file not found: {source}")
+                continue
+            staged_path, copied = stage_source(source_path, self.raw_dir / source_path.name)
+            if staged_path in staged:
+                continue
+            staged.add(staged_path)
+            copied_files.append(str(staged_path))
+            if copied and staged_path.name != source_path.name:
+                logger.warning(f"Name taken by another document, staged as {staged_path}")
+        return copied_files
+
+    async def process_documents(
+        self,
+    ) -> bool:
+        """Process documents with the KB's bound provider."""
+        provider = self.rag_provider
+
+        self.progress_tracker.update(
+            ProgressStage.PROCESSING_DOCUMENTS,
+            message_key="Starting to process documents with {{provider}} provider...",
+            message_params={"provider": provider},
+            current=0,
+            total=0,
+        )
+
+        # recursive=True so documents organized into folders are indexed too
+        # (folders are display-only and don't otherwise affect retrieval).
+        doc_files = FileTypeRouter.collect_supported_files(self.raw_dir, recursive=True)
+
+        if not doc_files:
+            self.progress_tracker.update(
+                ProgressStage.ERROR,
+                message_key="No documents found to process",
+                error="No documents found",
+            )
+            raise ValueError("No documents found to process")
+
+        self.progress_tracker.update(
+            ProgressStage.PROCESSING_DOCUMENTS,
+            message_key="Found {{count}} documents, starting to process...",
+            message_params={"count": len(doc_files)},
+            current=0,
+            total=len(doc_files),
+        )
+
+        rag_service = RAGService(
+            kb_base_dir=str(self.base_dir),
+            provider=provider,
+        )
+        file_paths = [str(doc_file) for doc_file in doc_files]
+
+        def _on_progress(batch_num, total_batches):
+            self.progress_tracker.update(
+                ProgressStage.PROCESSING_DOCUMENTS,
+                message_key="Embedding batches: {{current}}/{{total}} complete",
+                message_params={"current": batch_num, "total": total_batches},
+                current=batch_num,
+                total=total_batches,
+            )
+
+        def _on_image_progress(current: int, total: int):
+            self.progress_tracker.update(
+                ProgressStage.PROCESSING_DOCUMENTS,
+                message_key="Describing images: {{current}}/{{total}}",
+                message_params={"current": current, "total": total},
+                current=current,
+                total=total,
+            )
+
+        def _prepare_publication(candidate_root: Path) -> None:
+            self._update_metadata_with_provider(provider)
+            self.progress_tracker.update(
+                ProgressStage.COMPLETED,
+                message_key="Knowledge base initialization complete!",
+                current=1,
+                total=1,
+                indexed_count=len(doc_files),
+                index_changed=True,
+                index_action="create",
+                publication_version=candidate_root.name,
+            )
+            self.progress_tracker.verify_terminal(
+                current=1,
+                total=1,
+                indexed_count=len(doc_files),
+                index_action="create",
+                publication_version=candidate_root.name,
+            )
+
+        try:
+            success = await rag_service.initialize(
+                kb_name=self.kb_name,
+                file_paths=file_paths,
+                task_id=self.progress_tracker.task_id,
+                progress_callback=_on_progress,
+                image_progress_callback=_on_image_progress,
+                indexing_snapshot=self.indexing_snapshot,
+                before_publish=_prepare_publication if provider == "lightrag" else None,
+            )
+            if not success:
+                self.progress_tracker.update(
+                    ProgressStage.ERROR,
+                    message_key="Document processing failed",
+                    error="RAG pipeline returned failure",
+                )
+                raise RuntimeError("RAG pipeline returned failure")
+
+            self.index_published = provider == "lightrag"
+            try:
+                if provider != "lightrag":
+                    self._update_metadata_with_provider(provider)
+                    self.progress_tracker.update(
+                        ProgressStage.PROCESSING_DOCUMENTS,
+                        message_key="Documents processed successfully",
+                        current=len(doc_files),
+                        total=len(doc_files),
+                    )
+            except Exception:
+                if not self.index_published:
+                    raise
+                logger.warning(
+                    "LightRAG index for %s was published before metadata bookkeeping failed",
+                    self.kb_name,
+                    exc_info=True,
+                )
+        except PermissionError as e:
+            error_msg = format_data_volume_permission_error(self.kb_dir, cause=e)
+            logger.error("Error processing documents: %s", error_msg)
+            self.progress_tracker.update(
+                ProgressStage.ERROR,
+                message_key="Failed to process documents",
+                error=error_msg,
+            )
+            raise DataVolumePermissionError(error_msg) from e
+        except Exception as e:
+            error_msg = str(e)
+            logger.error(f"Error processing documents: {error_msg}")
+            self.progress_tracker.update(
+                ProgressStage.ERROR,
+                message_key="Failed to process documents",
+                error=error_msg,
+            )
+            raise
+
+        try:
+            await self.fix_structure()
+            await self.display_statistics_generic()
+        except Exception:
+            if not self.index_published:
+                raise
+            logger.warning(
+                "LightRAG index for %s was published before statistics bookkeeping failed",
+                self.kb_name,
+                exc_info=True,
+            )
+        return True
+
+    async def fix_structure(self) -> None:
+        """No-op retained for compatibility with previous pipelines."""
+        logger.info("Skipping legacy structure cleanup")
+
+    async def display_statistics_generic(self) -> None:
+        """Display basic statistics."""
+        raw_files = list(self.raw_dir.glob("*")) if self.raw_dir.exists() else []
+        from deeptutor.services.rag.index_probe import inspect_kb_versions
+
+        index_versions = inspect_kb_versions(self.kb_dir, self.rag_provider)
+
+        logger.info("=" * 50)
+        logger.info("Knowledge Base Statistics")
+        logger.info("=" * 50)
+        logger.info(f"Raw documents: {len(raw_files)}")
+        logger.info(f"Index versions: {len(index_versions)}")
+        logger.info(f"Provider used: {self.rag_provider}")
+        logger.info("=" * 50)
+
+
+async def initialize_knowledge_base(
+    kb_name: str,
+    source_files: list[str],
+    base_dir: str = "./data/knowledge_bases",
+    api_key: Optional[str] = None,
+    base_url: Optional[str] = None,
+    rag_provider: Optional[str] = None,
+) -> bool:
+    """Convenience initializer used by CLI wrappers."""
+    from deeptutor.knowledge.manager import KnowledgeBaseManager
+
+    manager = KnowledgeBaseManager(base_dir=base_dir)
+    initializer = KnowledgeBaseInitializer(
+        kb_name=kb_name,
+        base_dir=base_dir,
+        api_key=api_key,
+        base_url=base_url,
+        rag_provider=rag_provider,
+    )
+    try:
+        initializer.create_directory_structure()
+        copied_files = initializer.copy_documents(source_files)
+        await initializer.process_documents()
+        manager.update_kb_status(
+            name=kb_name,
+            status="ready",
+            progress={
+                "stage": "completed",
+                "message": "Knowledge base initialization complete!",
+                "percent": 100,
+                "current": 1,
+                "total": 1,
+                "file_name": "",
+                "error": None,
+                "timestamp": datetime.now().isoformat(),
+                "indexed_count": len(copied_files),
+                "index_changed": True,
+                "index_action": "create",
+            },
+        )
+        return True
+    except Exception as exc:
+        manager.update_kb_status(
+            name=kb_name,
+            status="error",
+            progress={
+                "stage": "error",
+                "message": "Knowledge base initialization failed",
+                "percent": 0,
+                "current": 0,
+                "total": 1,
+                "file_name": "",
+                "error": str(exc),
+                "timestamp": datetime.now().isoformat(),
+            },
+        )
+        raise
+
+
+async def main() -> None:
+    try:
+        llm_config = resolve_llm_runtime_config()
+        default_api_key = llm_config.api_key
+        default_base_url = llm_config.effective_url
+    except Exception:
+        default_api_key = ""
+        default_base_url = ""
+
+    parser = argparse.ArgumentParser(description="Initialize a new knowledge base from documents")
+    parser.add_argument("name", help="Knowledge base name")
+    parser.add_argument("--docs", nargs="+", help="Document files to process")
+    parser.add_argument("--docs-dir", help="Directory containing documents to process")
+    parser.add_argument("--base-dir", default="./knowledge_bases")
+    parser.add_argument("--api-key", default=default_api_key)
+    parser.add_argument("--base-url", default=default_base_url)
+    parser.add_argument("--skip-processing", action="store_true")
+
+    args = parser.parse_args()
+
+    doc_files: list[str] = []
+    if args.docs:
+        doc_files.extend(args.docs)
+    if args.docs_dir:
+        docs_dir = Path(args.docs_dir)
+        if docs_dir.exists() and docs_dir.is_dir():
+            doc_files.extend(str(f) for f in FileTypeRouter.collect_supported_files(docs_dir))
+
+    initializer = KnowledgeBaseInitializer(
+        kb_name=args.name,
+        base_dir=args.base_dir,
+        api_key=args.api_key,
+        base_url=args.base_url,
+    )
+    initializer.create_directory_structure()
+
+    if doc_files:
+        initializer.copy_documents(doc_files)
+
+    if not args.skip_processing:
+        await initializer.process_documents()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())