"""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 return bool(any(name.startswith(p) for p in skip_prefixes)) 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, "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS, "max_files": BOOKSLM_MAX_FILES, "scope": "directory", } # 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 _file_entry(target: Path, rel_path: str, remaining: int) -> dict[str, Any] | None: """Read, redact and truncate a single file into a context entry.""" suffix = target.suffix.lower() try: if suffix == ".pdf": from backend.pdf_reader import extract_pdf_text content = extract_pdf_text(target) file_type = "pdf" else: content = target.read_text(encoding="utf-8", errors="replace") file_type = "markdown" except Exception as e: logger.warning(f"Cannot read {rel_path}: {e}") return None content = redact_file_content(content, rel_path) if len(content) > BOOKSLM_MAX_FILE_CHARS: content = content[:BOOKSLM_MAX_FILE_CHARS] + "\n\n[... tronqué]" if len(content) > remaining: content = content[:remaining] + "\n\n[... tronqué]" return { "path": rel_path, "title": target.stem.replace("-", " ").replace("_", " ").title(), "content": content, "type": file_type, } def collect_files_context( vault_path: Path, rel_paths: list[str], scope: str = "documents", ) -> dict[str, Any]: """Collect an explicit list of files (open documents) as AI context. Unlike :func:`collect_directory_context`, this reads only the requested relative paths (markdown or PDF), never the whole directory. Paths outside the vault are silently ignored. Args: vault_path: Absolute path to the vault root. rel_paths: Relative file paths within the vault (in display order). scope: Context label exposed to the UI ("documents"). Returns: Same shape as :func:`collect_directory_context`. """ vault_resolved = vault_path.resolve() collected: list[dict[str, Any]] = [] total_chars = 0 seen: set[str] = set() for rel in rel_paths: if len(collected) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS: break if not rel or rel in seen: continue seen.add(rel) try: target = (vault_resolved / rel).resolve() target.relative_to(vault_resolved) except (ValueError, OSError): continue if not target.is_file(): continue entry = _file_entry(target, rel, BOOKSLM_MAX_TOTAL_CHARS - total_chars) if entry is None: continue collected.append(entry) total_chars += len(entry["content"]) return { "files": collected, "total_chars": total_chars, "file_count": len(collected), "directory_tree": "", "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS, "max_files": BOOKSLM_MAX_FILES, "scope": scope, } def empty_context(scope: str = "general") -> dict[str, Any]: """Return an empty context payload (used by the General assistant).""" return { "files": [], "total_chars": 0, "file_count": 0, "directory_tree": "", "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS, "max_files": BOOKSLM_MAX_FILES, "scope": scope, } 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], scope: str = "directory") -> str: """Build a system prompt for document-scoped AI chat. Args: context: Output of collect_directory_context()/collect_files_context(). scope: "directory" (whole folder) or "documents" (open files). 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" if scope == "documents": role = ( "Tu es un assistant de recherche documentaire intégré à ObsiGate. " "Tu réponds UNIQUEMENT en te basant sur les documents ouverts par l'utilisateur 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." ) context_label = "📄 Documents ouverts" else: role = ( "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." ) context_label = "📚 Contexte" prompt = ( f"{role}" f"\n\n{context_label} : {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 GENERAL_SYSTEM_PROMPT = """Tu es l'assistant intégré d'ObsiGate, une application web auto-hébergée pour consulter, rechercher et éditer des vaults Obsidian (Markdown). Tes deux rôles : 1. **Aider sur l'application** : expliquer la navigation, la recherche (full-text, filtres `tag:`, `created:`, `path:`), l'éditeur (CodeMirror, autosave, raccourcis), les onglets et le split view, les sauvegardes et la restauration, le partage public, l'export (HTML/Markdown/ePub/PDF), Mermaid, Excalidraw, les plugins, les thèmes, le mode hors-ligne, le MFA, etc. 2. **Proposer des actions concrètes** : créer un fichier ou un dossier dans un vault. Quand l'utilisateur demande explicitement de créer un fichier, inclus EXACTEMENT un bloc de ce type dans ta réponse (et rien d'autre à l'intérieur du bloc) : ```obsigate-action {"action": "create_file", "vault": "", "path": "", "content": ""} ``` Pour créer un dossier : ```obsigate-action {"action": "create_directory", "vault": "", "path": ""} ``` Règles : - Ne propose une action que si l'utilisateur la demande explicitement. - Explique en une phrase ce que fait l'action avant le bloc. - Utilise un chemin relatif se terminant par `.md` pour un fichier. - N'invente jamais un nom de vault : utilise l'un des vaults disponibles listés ci-dessous. - Réponds dans la langue de l'utilisateur, de façon concise et structurée (Markdown). """ def build_general_system_prompt(vaults: list[str] | None = None) -> str: """System prompt for the General assistant (app help + actions).""" prompt = GENERAL_SYSTEM_PROMPT if vaults: prompt += "\nVaults disponibles : " + ", ".join(sorted(vaults)) + "\n" else: prompt += "\nAucun vault n'est actuellement configuré.\n" 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 = list(_cache) for k in to_remove: del _cache[k] return len(to_remove)