"""BooksLM — Context collection and caching for directory-scoped AI chat. Collects markdown files from an Obsidian vault directory, applies secret redaction, builds a system prompt with file contents, and caches results for repeated queries. """ import hashlib import json import logging import os import time from pathlib import Path from typing import Any from backend.secret_redactor import redact_file_content logger = logging.getLogger("obsigate.bookslm") # ── Configuration limits ── BOOKSLM_MAX_FILES = int(os.getenv("BOOKSLM_MAX_FILES", "200")) BOOKSLM_MAX_TOTAL_CHARS = int(os.getenv("BOOKSLM_MAX_TOTAL_CHARS", "200000")) BOOKSLM_MAX_FILE_CHARS = int(os.getenv("BOOKSLM_MAX_FILE_CHARS", "30000")) # ── Cache ── _cache: dict[str, dict[str, Any]] = {} _CACHE_TTL = 300 # seconds def _cache_key(vault_path: Path, directory: str, file_mtimes: list[tuple[str, float]]) -> str: """Build a SHA-256 cache key from vault+directory+file modification times.""" raw = json.dumps({ "vault": str(vault_path), "dir": directory, "mtimes": sorted(file_mtimes), }, sort_keys=True) return hashlib.sha256(raw.encode()).hexdigest() def _should_skip(name: str) -> bool: """Return True if this file/directory name should be skipped.""" skip_prefixes = (".",) skip_names = {"_attachments", "node_modules", ".git", ".obsidian", "__pycache__"} if name in skip_names: return True if any(name.startswith(p) for p in skip_prefixes): return True return False def _file_priority(path: Path) -> tuple[int, float]: """Sort key: README/index first, then by modification time descending. Returns (priority_group, -mtime) so that: - Group 0: README* and index* files (come first) - Group 1: all other files (come after) Within each group, newer files come first. """ name_lower = path.stem.lower() if name_lower.startswith("readme") or name_lower.startswith("index"): group = 0 else: group = 1 try: mtime = path.stat().st_mtime except OSError: mtime = 0.0 return (group, -mtime) def collect_directory_context(vault_path: Path, directory: str) -> dict[str, Any]: """Walk a directory recursively, collect .md files with content. Args: vault_path: Absolute path to the vault root. directory: Relative directory path within the vault (empty = root). Returns: Dict with keys: files, total_chars, file_count, directory_tree. """ target_dir = (vault_path / directory).resolve() if directory else vault_path.resolve() vault_resolved = vault_path.resolve() # Safety: ensure target is within vault try: target_dir.relative_to(vault_resolved) except ValueError: logger.warning(f"Directory outside vault: {target_dir}") return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""} if not target_dir.exists() or not target_dir.is_dir(): return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""} # Check cache file_mtimes: list[tuple[str, float]] = [] md_files: list[Path] = [] try: for p in target_dir.rglob("*"): # Skip hidden dirs/files and special dirs parts = p.relative_to(target_dir).parts if any(_should_skip(part) for part in parts): continue if p.is_file() and p.suffix.lower() == ".md": md_files.append(p) try: file_mtimes.append((str(p.relative_to(target_dir)), p.stat().st_mtime)) except OSError: file_mtimes.append((str(p.relative_to(target_dir)), 0.0)) except PermissionError: logger.warning(f"Permission denied scanning {target_dir}") return {"files": [], "total_chars": 0, "file_count": 0, "directory_tree": ""} key = _cache_key(vault_resolved, directory, file_mtimes) if key in _cache: cached = _cache[key] if time.time() - cached.get("_ts", 0) < _CACHE_TTL: logger.debug(f"Cache hit for {directory}") return {k: v for k, v in cached.items() if k != "_ts"} # Sort by priority: README/index first, then by mtime descending md_files.sort(key=_file_priority) # Apply limits collected: list[dict[str, Any]] = [] total_chars = 0 for p in md_files: if len(collected) >= BOOKSLM_MAX_FILES: break if total_chars >= BOOKSLM_MAX_TOTAL_CHARS: break rel_path = str(p.relative_to(vault_resolved)).replace("\\", "/") try: content = p.read_text(encoding="utf-8", errors="replace") except Exception as e: logger.warning(f"Cannot read {rel_path}: {e}") continue # Redact secrets content = redact_file_content(content, rel_path) # Truncate if too long if len(content) > BOOKSLM_MAX_FILE_CHARS: content = content[:BOOKSLM_MAX_FILE_CHARS] + "\n\n[... tronqué]" remaining = BOOKSLM_MAX_TOTAL_CHARS - total_chars if len(content) > remaining: content = content[:remaining] + "\n\n[... tronqué]" title = p.stem.replace("-", " ").replace("_", " ").title() file_type = "markdown" collected.append({ "path": rel_path, "title": title, "content": content, "type": file_type, }) total_chars += len(content) # Build directory tree dir_tree = _build_directory_tree(target_dir, vault_resolved) result = { "files": collected, "total_chars": total_chars, "file_count": len(collected), "directory_tree": dir_tree, } # Store in cache _cache[key] = {**result, "_ts": time.time()} logger.info(f"Collected {len(collected)} files ({total_chars} chars) from {directory or '/'}") return result def _build_directory_tree(target_dir: Path, vault_root: Path) -> str: """Build a text representation of the directory tree (dirs + .md files).""" lines: list[str] = [] try: for p in sorted(target_dir.rglob("*")): parts = p.relative_to(target_dir).parts if any(_should_skip(part) for part in parts): continue if p.is_dir(): depth = len(p.relative_to(target_dir).parts) lines.append(f"{' ' * depth}{p.name}/") elif p.is_file() and p.suffix.lower() == ".md": depth = len(p.relative_to(target_dir).parts) lines.append(f"{' ' * depth}{p.name}") except PermissionError: pass return "\n".join(lines) def build_system_prompt(context: dict[str, Any]) -> str: """Build a system prompt for directory-scoped AI chat. Args: context: Output of collect_directory_context(). Returns: System prompt string with file contents. """ files = context.get("files", []) file_count = context.get("file_count", 0) total_chars = context.get("total_chars", 0) # Rough token estimate (1 token ≈ 4 chars) est_tokens = total_chars // 4 token_warning = "" if est_tokens > 100_000: token_warning = f"\n⚠️ Attention : le contexte est très volumineux (~{est_tokens:,} tokens estimés). Les réponses peuvent être moins précises.\n" prompt = ( "Tu es un assistant de recherche documentaire. " "Tu réponds UNIQUEMENT en te basant sur les documents fournis ci-dessous. " "Cite tes sources avec le nom du fichier quand tu utilises une information. " "Si l'information ne se trouve pas dans les documents, dis-le clairement." f"\n\n📚 Contexte : {file_count} fichier(s) ({total_chars:,} caractères)" f"{token_warning}\n" ) # Directory tree tree = context.get("directory_tree", "") if tree: prompt += f"\n📂 Arborescence du dossier :\n```\n{tree}\n```\n" # File contents prompt += "\n---\n" for f in files: prompt += f"\n## 📄 {f['title']} (`{f['path']}`)\n\n{f['content']}\n\n---\n" prompt += "\nFin du contexte. Réponds à la question de l'utilisateur en te basant uniquement sur ces documents." return prompt def invalidate_cache(vault_path: Path | None = None, directory: str | None = None) -> int: """Invalidate cache entries. Args: vault_path: If provided, only invalidate entries for this vault. directory: If provided, only invalidate entries for this directory. Returns: Number of cache entries removed. """ if vault_path is None and directory is None: count = len(_cache) _cache.clear() return count to_remove = [] for key, val in _cache.items(): # We can't easily reverse the hash, so clear everything if vault_path is given # For targeted invalidation, callers should use directory to_remove.append(key) for k in to_remove: del _cache[k] return len(to_remove)