wehub-resource-sync
e4dcfc49aa
Tests / Import Check (Python 3.13) (push) Has been cancelled
Tests / Import Check (Python 3.14) (push) Has been cancelled
Tests / Python Tests (Python 3.11) (push) Has been cancelled
Tests / Python Tests (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.14) (push) Has been cancelled
Tests / Test Summary (push) Has been cancelled
Tests / Lint and Format (push) Has been cancelled
Tests / Web Node Tests (push) Has been cancelled
Tests / Import Check (Python 3.11) (push) Has been cancelled
Tests / Import Check (Python 3.12) (push) Has been cancelled
Tests / Python Tests (Python 3.13) (push) Has been cancelled
478 lines
17 KiB
Python
478 lines
17 KiB
Python
#!/usr/bin/env python
|
|
"""Incrementally add documents to an existing knowledge base."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
import shutil
|
|
from typing import List, Optional
|
|
|
|
from deeptutor.services.config import resolve_llm_runtime_config
|
|
from deeptutor.services.rag.factory import (
|
|
DEFAULT_PROVIDER,
|
|
has_ready_provider_index,
|
|
normalize_provider_name,
|
|
)
|
|
from deeptutor.services.rag.file_routing import FileTypeRouter
|
|
from deeptutor.services.rag.provider_binding import resolve_bound_provider
|
|
from deeptutor.services.rag.service import RAGService
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_BASE_DIR = "./data/knowledge_bases"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DocumentIndexFailure:
|
|
"""One file that could not be added to the provider index."""
|
|
|
|
file_path: Path
|
|
error: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DocumentIndexResult:
|
|
"""Structured incremental-index result.
|
|
|
|
A task can no longer infer success from "no exception": every staged file is
|
|
explicitly accounted for as either processed or failed.
|
|
"""
|
|
|
|
processed_files: list[Path]
|
|
failures: list[DocumentIndexFailure]
|
|
|
|
@property
|
|
def processed_count(self) -> int:
|
|
return len(self.processed_files)
|
|
|
|
@property
|
|
def failed_count(self) -> int:
|
|
return len(self.failures)
|
|
|
|
@property
|
|
def has_failures(self) -> bool:
|
|
return bool(self.failures)
|
|
|
|
def failure_summary(self, *, limit: int = 3) -> str:
|
|
if not self.failures:
|
|
return ""
|
|
shown = [f"{failure.file_path.name}: {failure.error}" for failure in self.failures[:limit]]
|
|
remaining = len(self.failures) - len(shown)
|
|
if remaining > 0:
|
|
shown.append(f"... and {remaining} more file(s)")
|
|
return "; ".join(shown)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RawDocumentRemoval:
|
|
"""Outcome of removing one raw document from a knowledge base."""
|
|
|
|
rel_path: str
|
|
was_indexed: bool
|
|
|
|
|
|
def _read_metadata(metadata_file: Path) -> dict:
|
|
"""Load a KB's metadata.json, returning {} when absent or unreadable."""
|
|
if not metadata_file.exists():
|
|
return {}
|
|
try:
|
|
with open(metadata_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
except Exception:
|
|
return {}
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def _write_metadata(metadata_file: Path, metadata: dict) -> None:
|
|
"""Persist a KB's metadata.json (pretty-printed, non-ASCII preserved)."""
|
|
with open(metadata_file, "w", encoding="utf-8") as f:
|
|
json.dump(metadata, f, indent=2, ensure_ascii=False)
|
|
|
|
|
|
def _raw_hash_key(file_path: Path, raw_dir: Path) -> str:
|
|
"""Stable ``file_hashes`` key for a staged file.
|
|
|
|
The key is the POSIX path relative to ``raw/`` (so folder uploads keep
|
|
distinct entries), falling back to the basename for anything that resolves
|
|
outside ``raw/``. Indexing and removal MUST share this rule, otherwise a
|
|
removed file's hash record would leak and silently skip a later re-add.
|
|
"""
|
|
try:
|
|
return file_path.resolve().relative_to(raw_dir.resolve()).as_posix()
|
|
except ValueError:
|
|
return file_path.name
|
|
|
|
|
|
def remove_raw_document(kb_dir: Path, file_path: Path) -> RawDocumentRemoval:
|
|
"""Delete one staged raw file and drop its indexed-hash record.
|
|
|
|
Deliberately decoupled from :class:`DocumentAdder` (whose constructor
|
|
requires a ready provider index): a KB stuck in an *error* state often has
|
|
no index at all, yet its raw/ files must still be removable so the user can
|
|
drop an unparseable document instead of deleting and rebuilding the whole
|
|
KB. Vectors are not touched here — the providers index whole KBs, so any
|
|
vectors from an already-indexed file are cleared by a subsequent re-index,
|
|
which the returned ``was_indexed`` flag lets the caller surface.
|
|
|
|
``file_path`` must already be sandbox-resolved under the KB's raw/ dir.
|
|
"""
|
|
raw_dir = kb_dir / "raw"
|
|
hash_key = _raw_hash_key(file_path, raw_dir)
|
|
|
|
target = file_path.resolve()
|
|
if target.exists():
|
|
target.unlink()
|
|
|
|
metadata_file = kb_dir / "metadata.json"
|
|
metadata = _read_metadata(metadata_file)
|
|
hashes = metadata.get("file_hashes")
|
|
was_indexed = isinstance(hashes, dict) and hash_key in hashes
|
|
if was_indexed:
|
|
del hashes[hash_key]
|
|
_write_metadata(metadata_file, metadata)
|
|
|
|
return RawDocumentRemoval(rel_path=hash_key, was_indexed=was_indexed)
|
|
|
|
|
|
class DocumentAdder:
|
|
"""Stage and index new files through a KB's bound RAG provider."""
|
|
|
|
def __init__(
|
|
self,
|
|
kb_name: str,
|
|
base_dir: str = DEFAULT_BASE_DIR,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
progress_tracker=None,
|
|
rag_provider: str | None = None,
|
|
):
|
|
self.kb_name = kb_name
|
|
self.base_dir = Path(base_dir)
|
|
self.kb_dir = self.base_dir / kb_name
|
|
|
|
if not self.kb_dir.exists():
|
|
raise ValueError(f"Knowledge base does not exist: {kb_name}")
|
|
|
|
self.raw_dir = self.kb_dir / "raw"
|
|
self.llamaindex_storage_dir = self.kb_dir / "llamaindex_storage"
|
|
self.legacy_rag_storage_dir = self.kb_dir / "rag_storage"
|
|
self.metadata_file = self.kb_dir / "metadata.json"
|
|
|
|
# Incremental adds must use the engine DeepTutor has bound to this KB. An
|
|
# explicit rag_provider (from the API, already matched against the KB)
|
|
# wins; otherwise use the shared binding resolver.
|
|
if rag_provider:
|
|
self.rag_provider = normalize_provider_name(rag_provider)
|
|
else:
|
|
self.rag_provider = self._provider_from_binding()
|
|
|
|
has_provider_index = has_ready_provider_index(self.kb_dir, self.rag_provider)
|
|
|
|
if (
|
|
self.rag_provider == DEFAULT_PROVIDER
|
|
and not has_provider_index
|
|
and self.legacy_rag_storage_dir.exists()
|
|
):
|
|
raise ValueError(
|
|
f"Knowledge base '{kb_name}' uses legacy index format and requires reindex before incremental add"
|
|
)
|
|
|
|
if not has_provider_index:
|
|
raise ValueError(f"Knowledge base not initialized ({self.rag_provider}): {kb_name}")
|
|
|
|
self.api_key = api_key
|
|
self.base_url = base_url
|
|
self.progress_tracker = progress_tracker
|
|
|
|
self.raw_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def _provider_from_binding(self) -> str:
|
|
return resolve_bound_provider(self.base_dir, self.kb_name)
|
|
|
|
def _get_file_hash(self, file_path: Path) -> str:
|
|
sha256_hash = hashlib.sha256()
|
|
with open(file_path, "rb") as f:
|
|
for byte_block in iter(lambda: f.read(65536), b""):
|
|
sha256_hash.update(byte_block)
|
|
return sha256_hash.hexdigest()
|
|
|
|
def get_ingested_hashes(self) -> dict[str, str]:
|
|
if self.metadata_file.exists():
|
|
try:
|
|
with open(self.metadata_file, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
return data.get("file_hashes", {})
|
|
except Exception:
|
|
return {}
|
|
return {}
|
|
|
|
def add_documents(self, source_files: List[str], allow_duplicates: bool = False) -> List[Path]:
|
|
"""Validate and stage files into raw/ before indexing."""
|
|
logger.info(f"Validating documents for '{self.kb_name}'...")
|
|
|
|
ingested_hashes = self.get_ingested_hashes()
|
|
files_to_process: list[Path] = []
|
|
|
|
for source in source_files:
|
|
source_path = Path(source)
|
|
if not source_path.exists() or not source_path.is_file():
|
|
logger.warning(f"Missing file: {source}")
|
|
continue
|
|
|
|
current_hash = self._get_file_hash(source_path)
|
|
if current_hash in ingested_hashes.values() and not allow_duplicates:
|
|
logger.info(f"Skipped (content already indexed): {source_path.name}")
|
|
continue
|
|
|
|
# Files already saved under raw/ (e.g. by the upload route, possibly
|
|
# inside a folder) are indexed in place — never flattened to the
|
|
# basename — so the uploaded folder structure is preserved verbatim.
|
|
if source_path.resolve().is_relative_to(self.raw_dir.resolve()):
|
|
files_to_process.append(source_path)
|
|
continue
|
|
|
|
dest_path = self.raw_dir / source_path.name
|
|
if dest_path.exists():
|
|
dest_hash = self._get_file_hash(dest_path)
|
|
if dest_hash == current_hash:
|
|
logger.info(f"Recovering staged file: {source_path.name}")
|
|
files_to_process.append(dest_path)
|
|
continue
|
|
if not allow_duplicates:
|
|
logger.info(f"Skipped (filename collision): {source_path.name}")
|
|
continue
|
|
|
|
shutil.copy2(source_path, dest_path)
|
|
logger.info(f"Staged to raw: {source_path.name}")
|
|
files_to_process.append(dest_path)
|
|
|
|
return files_to_process
|
|
|
|
async def process_new_documents(self, new_files: List[Path]) -> DocumentIndexResult:
|
|
"""Index staged files via the KB's bound provider."""
|
|
if not new_files:
|
|
return DocumentIndexResult(processed_files=[], failures=[])
|
|
|
|
rag_service = RAGService(kb_base_dir=str(self.base_dir), provider=self.rag_provider)
|
|
processed_files: list[Path] = []
|
|
failures: list[DocumentIndexFailure] = []
|
|
total_files = len(new_files)
|
|
|
|
for idx, doc_file in enumerate(new_files, 1):
|
|
try:
|
|
if self.progress_tracker is not None:
|
|
from deeptutor.knowledge.progress_tracker import ProgressStage
|
|
|
|
self.progress_tracker.update(
|
|
ProgressStage.PROCESSING_FILE,
|
|
f"Indexing {doc_file.name}",
|
|
current=idx,
|
|
total=total_files,
|
|
)
|
|
|
|
success = await rag_service.add_documents(self.kb_name, [str(doc_file)])
|
|
if success:
|
|
processed_files.append(doc_file)
|
|
self._record_successful_hash(doc_file)
|
|
logger.info(f"Processed: {doc_file.name}")
|
|
else:
|
|
error = "Provider returned failure without details."
|
|
failures.append(DocumentIndexFailure(doc_file, error))
|
|
logger.error(f"Failed to index: {doc_file.name}")
|
|
except Exception as e:
|
|
logger.exception(f"Failed {doc_file.name}: {e}")
|
|
failures.append(DocumentIndexFailure(doc_file, str(e)))
|
|
|
|
return DocumentIndexResult(processed_files=processed_files, failures=failures)
|
|
|
|
def _record_successful_hash(self, file_path: Path) -> None:
|
|
file_hash = self._get_file_hash(file_path)
|
|
metadata = _read_metadata(self.metadata_file)
|
|
hash_key = _raw_hash_key(file_path, self.raw_dir)
|
|
metadata.setdefault("file_hashes", {})[hash_key] = file_hash
|
|
_write_metadata(self.metadata_file, metadata)
|
|
|
|
def update_metadata(self, added_count: int) -> None:
|
|
"""Update metadata after incremental add."""
|
|
metadata: dict = {}
|
|
if self.metadata_file.exists():
|
|
try:
|
|
with open(self.metadata_file, "r", encoding="utf-8") as f:
|
|
metadata = json.load(f)
|
|
except Exception:
|
|
metadata = {}
|
|
|
|
metadata["rag_provider"] = self.rag_provider
|
|
metadata["needs_reindex"] = False
|
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
metadata["last_updated"] = timestamp
|
|
if added_count > 0:
|
|
metadata["last_indexed_at"] = timestamp
|
|
metadata["last_indexed_count"] = added_count
|
|
metadata["last_indexed_action"] = "upload"
|
|
|
|
history = metadata.get("update_history", [])
|
|
history.append(
|
|
{
|
|
"timestamp": metadata["last_updated"],
|
|
"action": "incremental_add",
|
|
"count": added_count,
|
|
"provider": self.rag_provider,
|
|
}
|
|
)
|
|
metadata["update_history"] = history
|
|
|
|
with open(self.metadata_file, "w", encoding="utf-8") as f:
|
|
json.dump(metadata, f, indent=2, ensure_ascii=False)
|
|
|
|
|
|
async def add_documents(
|
|
kb_name: str,
|
|
source_files: list[str],
|
|
base_dir: str = DEFAULT_BASE_DIR,
|
|
api_key: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
allow_duplicates: bool = False,
|
|
) -> int:
|
|
"""Convenience function used by CLI wrappers."""
|
|
from deeptutor.knowledge.manager import KnowledgeBaseManager
|
|
|
|
manager = KnowledgeBaseManager(base_dir=base_dir)
|
|
try:
|
|
manager.update_kb_status(
|
|
name=kb_name,
|
|
status="processing",
|
|
progress={
|
|
"stage": "processing_documents",
|
|
"message": "Processing uploaded documents...",
|
|
"percent": 0,
|
|
"current": 0,
|
|
"total": max(len(source_files), 1),
|
|
"file_name": "",
|
|
"error": None,
|
|
"timestamp": datetime.now().isoformat(),
|
|
},
|
|
)
|
|
|
|
adder = DocumentAdder(
|
|
kb_name=kb_name,
|
|
base_dir=base_dir,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
)
|
|
new_files = adder.add_documents(source_files, allow_duplicates=allow_duplicates)
|
|
if not new_files:
|
|
manager.update_kb_status(
|
|
name=kb_name,
|
|
status="ready",
|
|
progress={
|
|
"stage": "completed",
|
|
"message": "No new unique documents to process.",
|
|
"percent": 100,
|
|
"current": 1,
|
|
"total": 1,
|
|
"file_name": "",
|
|
"error": None,
|
|
"timestamp": datetime.now().isoformat(),
|
|
},
|
|
)
|
|
return 0
|
|
result = await adder.process_new_documents(new_files)
|
|
if result.has_failures:
|
|
raise RuntimeError(
|
|
f"Failed to index {result.failed_count}/{len(new_files)} file(s): "
|
|
f"{result.failure_summary()}"
|
|
)
|
|
adder.update_metadata(result.processed_count)
|
|
|
|
manager.update_kb_status(
|
|
name=kb_name,
|
|
status="ready",
|
|
progress={
|
|
"stage": "completed",
|
|
"message": f"Successfully processed {result.processed_count} files!",
|
|
"percent": 100,
|
|
"current": result.processed_count,
|
|
"total": max(len(new_files), 1),
|
|
"file_name": "",
|
|
"error": None,
|
|
"timestamp": datetime.now().isoformat(),
|
|
"indexed_count": result.processed_count,
|
|
"index_changed": result.processed_count > 0,
|
|
"index_action": "upload",
|
|
},
|
|
)
|
|
return result.processed_count
|
|
except Exception as exc:
|
|
manager.update_kb_status(
|
|
name=kb_name,
|
|
status="error",
|
|
progress={
|
|
"stage": "error",
|
|
"message": "Document upload failed",
|
|
"percent": 0,
|
|
"current": 0,
|
|
"total": max(len(source_files), 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="Incrementally add documents to a KB")
|
|
parser.add_argument("kb_name", help="KB Name")
|
|
parser.add_argument("--docs", nargs="+", help="Files")
|
|
parser.add_argument("--docs-dir", help="Directory")
|
|
parser.add_argument("--base-dir", default=DEFAULT_BASE_DIR)
|
|
parser.add_argument("--api-key", default=default_api_key)
|
|
parser.add_argument("--base-url", default=default_base_url)
|
|
parser.add_argument("--allow-duplicates", action="store_true")
|
|
|
|
args = parser.parse_args()
|
|
doc_files: list[str] = []
|
|
if args.docs:
|
|
doc_files.extend(args.docs)
|
|
if args.docs_dir:
|
|
p = Path(args.docs_dir)
|
|
doc_files.extend(str(f) for f in FileTypeRouter.collect_supported_files(p))
|
|
|
|
if not doc_files:
|
|
logger.error("No documents provided.")
|
|
return
|
|
|
|
processed_count = await add_documents(
|
|
kb_name=args.kb_name,
|
|
source_files=doc_files,
|
|
base_dir=args.base_dir,
|
|
api_key=args.api_key,
|
|
base_url=args.base_url,
|
|
allow_duplicates=args.allow_duplicates,
|
|
)
|
|
|
|
if processed_count:
|
|
logger.info(f"Done! Successfully added {processed_count} documents.")
|
|
else:
|
|
logger.info("No new unique documents to add.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|