Spaces:
Sleeping
Sleeping
| """Knowledge-base seeder — indexes reference documents under the KB tenant. | |
| Indexes PDFs and DOCX files from ``settings.knowledge_base_dirs`` under | |
| ``settings.knowledge_base_tenant_id`` (``__rics_kb__``). Used as the style | |
| and evidence fallback when a user has not yet uploaded their own documents. | |
| The seed is a no-op if the KB tenant already has chunks in the vector store, | |
| so it is safe to call on every startup. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import hashlib | |
| import logging | |
| from pathlib import Path | |
| from app.config import settings | |
| from app.ingest.pipeline import _get_thread_pool, _sync_embed, _sync_index, _sync_parse_and_chunk | |
| logger = logging.getLogger(__name__) | |
| _SUPPORTED: frozenset[str] = frozenset({".pdf", ".docx"}) | |
| def _resolve_kb_dirs() -> list[Path]: | |
| """Return all configured KB directories that exist on disk.""" | |
| dirs: list[Path] = [] | |
| repo_root = Path(__file__).parent.parent.parent | |
| for raw in settings.knowledge_base_dirs.split(","): | |
| raw = raw.strip() | |
| if not raw: | |
| continue | |
| candidate = Path(raw) if Path(raw).is_absolute() else repo_root / raw | |
| if candidate.is_dir(): | |
| dirs.append(candidate) | |
| else: | |
| logger.warning("KB directory not found, skipping: %s", candidate) | |
| return dirs | |
| def _kb_doc_id(path: Path) -> str: | |
| """Deterministic doc-id from file path — no DB record needed for system docs.""" | |
| return "kb_" + hashlib.sha1(str(path).encode()).hexdigest()[:16] | |
| async def seed_knowledge_base() -> int: | |
| """Index all supported files from knowledge_base_dirs under the KB tenant. | |
| Skips the entire seed if the KB tenant already has any chunks indexed, so | |
| re-running on every startup is cheap. | |
| Returns: | |
| Number of new files successfully indexed (0 when already seeded). | |
| """ | |
| if not settings.knowledge_base_enabled: | |
| logger.info("KB seeding disabled (knowledge_base_enabled=False)") | |
| return 0 | |
| from app.vectorstore.factory import get_vectorstore | |
| vs = get_vectorstore() | |
| tenant = settings.knowledge_base_tenant_id | |
| try: | |
| existing = vs.count(tenant_id=tenant) | |
| except Exception: | |
| existing = 0 | |
| if existing > 0: | |
| logger.info( | |
| "KB tenant '%s' already has %d chunks — skipping seed", tenant, existing | |
| ) | |
| return 0 | |
| dirs = _resolve_kb_dirs() | |
| if not dirs: | |
| logger.warning("No KB directories found — cannot seed knowledge base") | |
| return 0 | |
| files = [ | |
| f | |
| for d in dirs | |
| for f in sorted(d.iterdir()) | |
| if f.is_file() and f.suffix.lower() in _SUPPORTED | |
| ] | |
| if not files: | |
| logger.warning("No PDF/DOCX files found in KB directories: %s", dirs) | |
| return 0 | |
| logger.info( | |
| "Seeding KB tenant '%s' from %d file(s) across %d director(y/ies)", | |
| tenant, | |
| len(files), | |
| len(dirs), | |
| ) | |
| loop = asyncio.get_running_loop() | |
| pool = _get_thread_pool() | |
| indexed = 0 | |
| for fpath in files: | |
| doc_id = _kb_doc_id(fpath) | |
| try: | |
| chunks = await loop.run_in_executor( | |
| pool, _sync_parse_and_chunk, fpath, doc_id, tenant | |
| ) | |
| if not chunks: | |
| logger.debug("KB: no chunks produced for %s — skipping", fpath.name) | |
| continue | |
| texts = [c.text for c in chunks] | |
| vectors = await loop.run_in_executor(pool, _sync_embed, texts) | |
| await loop.run_in_executor(pool, _sync_index, chunks, vectors) | |
| logger.info("KB indexed: %s → %d chunks", fpath.name, len(chunks)) | |
| indexed += 1 | |
| except Exception as exc: | |
| logger.warning("KB seed failed for %s: %s", fpath.name, exc) | |
| logger.info("KB seeding complete: %d/%d files indexed", indexed, len(files)) | |
| return indexed | |