hkuds--deeptutor/deeptutor/knowledge/initializer.py

#!/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
import shutil
from typing import Optional

from deeptutor.knowledge.naming import validate_knowledge_base_name
from deeptutor.knowledge.progress_tracker import ProgressStage, ProgressTracker
from deeptutor.services.config import resolve_llm_runtime_config
from deeptutor.services.rag.factory import normalize_provider_name
from deeptutor.services.rag.file_routing import FileTypeRouter
from deeptutor.services.rag.service import RAGService

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,
    ):
        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)

    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)
        )
        metadata["last_indexed_action"] = "create"

        with open(metadata_file, "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)

        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}")

        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,
        }

        with open(self.kb_dir / "metadata.json", "w", encoding="utf-8") as f:
            json.dump(metadata, indent=2, ensure_ascii=False, fp=f)

        self._register_to_config()

    def copy_documents(self, source_files: list[str]) -> list[str]:
        """Copy source documents into raw directory."""
        copied_files: list[str] = []
        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
            dest_path = self.raw_dir / source_path.name
            shutil.copy2(source_path, dest_path)
            copied_files.append(str(dest_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,
            f"Starting to process documents with {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,
                "No documents found to process",
                error="No documents found",
            )
            raise ValueError("No documents found to process")

        self.progress_tracker.update(
            ProgressStage.PROCESSING_DOCUMENTS,
            f"Found {len(doc_files)} documents, starting to process...",
            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,
                f"Embedding batches: {batch_num}/{total_batches} complete",
                current=batch_num,
                total=total_batches,
            )

        try:
            success = await rag_service.initialize(
                kb_name=self.kb_name,
                file_paths=file_paths,
                progress_callback=_on_progress,
            )
            if not success:
                self.progress_tracker.update(
                    ProgressStage.ERROR,
                    "Document processing failed",
                    error="RAG pipeline returned failure",
                )
                raise RuntimeError("RAG pipeline returned failure")

            self._update_metadata_with_provider(provider)
            self.progress_tracker.update(
                ProgressStage.PROCESSING_DOCUMENTS,
                "Documents processed successfully",
                current=len(doc_files),
                total=len(doc_files),
            )
        except Exception as e:
            error_msg = str(e)
            logger.error(f"Error processing documents: {error_msg}")
            self.progress_tracker.update(
                ProgressStage.ERROR,
                "Failed to process documents",
                error=error_msg,
            )
            raise

        await self.fix_structure()
        await self.display_statistics_generic()
        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())