"""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 base64 import hashlib import json import logging import mimetypes import os import time from pathlib import Path from typing import Any from backend.media_types import is_media 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")) # Maximum size of an image sent to a vision model (bytes, before base64). BOOKSLM_MAX_IMAGE_BYTES = int(os.getenv("BOOKSLM_MAX_IMAGE_BYTES", "10000000")) IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".avif"} # ── 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() # #109-D3 — audio/video (and images) carry no extractable text; never feed # raw bytes to the model. Images are handled separately via vision data URLs. if is_media(suffix): return None 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 collect_adhoc_context( vault_path: Path, files: list[str] | None = None, directories: list[str] | None = None, scope: str = "general", ) -> dict[str, Any]: """Collect an ad-hoc mix of explicit files and directories. Backs the assistant ``@`` command: the user can attach extra files and directories to the current context without changing the base mode. Paths outside the vault are ignored by the underlying collectors. Args: vault_path: Absolute path to the vault root. files: Relative file paths to include. directories: Relative directory paths to include (recursive). scope: Context label exposed to the UI. Returns: Same shape as :func:`collect_directory_context`. """ vault_resolved = vault_path.resolve() collected: list[dict[str, Any]] = [] seen: set[str] = set() total_chars = 0 def _add(entries: list[dict[str, Any]]) -> None: nonlocal total_chars for entry in entries: if len(collected) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS: return path = entry.get("path") if not path or path in seen: continue seen.add(path) collected.append(entry) total_chars += len(entry.get("content", "")) if files: _add(collect_files_context(vault_resolved, files, scope=scope)["files"]) for directory in directories or []: _add(collect_directory_context(vault_resolved, directory)["files"]) 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 merge_contexts(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: """Merge two context payloads, de-duplicating files by path.""" files: list[dict[str, Any]] = [] seen: set[str] = set() total_chars = 0 for entry in list(base.get("files", [])) + list(extra.get("files", [])): path = entry.get("path") if not path or path in seen: continue if len(files) >= BOOKSLM_MAX_FILES or total_chars >= BOOKSLM_MAX_TOTAL_CHARS: break seen.add(path) files.append(entry) total_chars += len(entry.get("content", "")) tree = base.get("directory_tree", "") or extra.get("directory_tree", "") scope = extra.get("scope") or base.get("scope", "general") return { "files": files, "total_chars": total_chars, "file_count": len(files), "directory_tree": tree, "max_total_chars": BOOKSLM_MAX_TOTAL_CHARS, "max_files": BOOKSLM_MAX_FILES, "scope": scope, } def is_image_path(path: str) -> bool: """True when the path has a supported image extension.""" return Path(path or "").suffix.lower() in IMAGE_EXTENSIONS def load_vault_image_data_url(vault_path: Path, rel_path: str) -> str | None: """Read a vault image and return it as a ``data:`` URL for vision models. Returns ``None`` when the path is outside the vault, missing, too large, or not a supported image. """ if not rel_path or not is_image_path(rel_path): return None vault_resolved = vault_path.resolve() try: target = (vault_resolved / rel_path).resolve() target.relative_to(vault_resolved) except (ValueError, OSError): return None if not target.is_file(): return None try: if target.stat().st_size > BOOKSLM_MAX_IMAGE_BYTES: logger.warning("Image too large to send: %s", rel_path) return None raw = target.read_bytes() except OSError as exc: logger.warning("Cannot read image %s: %s", rel_path, exc) return None mime = mimetypes.guess_type(str(target))[0] or "image/png" encoded = base64.b64encode(raw).decode("ascii") return f"data:{mime};base64,{encoded}" 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", vault_name: str | None = None) -> 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). vault_name: Vault the context files belong to. When set, the prompt states it explicitly with write-tool guidance (BUG-046: without it the model invented vault names — e.g. "test" — and every confirmed ``append_to_file``/``edit_file`` call failed). 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." if vault_name: prompt += ( f"\n\nCes documents appartiennent au vault « {vault_name} ». Quand tu utilises un outil " "d'écriture (`append_to_file`, `edit_file`, `create_file`), passe TOUJOURS " f"exactement `\"vault\": \"{vault_name}\"` (jamais un nom inventé) et un `path` " "relatif au vault, identique à celui affiché ci-dessus. Pour créer un fichier " "dans un nouveau dossier, un seul `create_file` avec le chemin complet suffit " "(les dossiers parents sont créés automatiquement)." ) return prompt GENERAL_SYSTEM_HEADER = """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. """ # Text action protocol — used by the classic (non-agent) chat endpoint, where # the model has no native tool calling; the frontend turns each block into a # clickable “Apply” card. GENERAL_ACTION_TEXT_PROTOCOL = """ Quand l'utilisateur demande explicitement de créer un fichier, inclus un bloc de ce type dans ta réponse (un bloc par fichier, 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. - Pour créer un fichier dans un nouveau dossier, utilise **un seul** bloc `create_file` avec le chemin complet (ex. `"path": "Dossier/fichier.md"`) : les dossiers parents sont créés automatiquement, inutile d'émettre un `create_directory` séparé. - 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). """ # Agent mode: the model has native tools, so it must call them (function # calling) instead of emitting the text `obsigate-action` blocks — otherwise # the requested file is never created (BUG-053). GENERAL_ACTION_TOOL_PROTOCOL = """ Tu disposes d'outils natifs (function calling) pour lire, chercher et modifier les vaults : `create_file`, `create_directory`, `append_to_file`, `edit_file`, `read_file`, `search_fulltext`, etc. Quand l'utilisateur demande explicitement de créer un fichier, **appelle directement l'outil `create_file`** avec `{"vault": "", "path": "", "content": ""}`. Pour créer un dossier, appelle `create_directory`. Règles : - N'écris **jamais** de bloc ```obsigate-action``` : en mode agent, toutes les actions passent par les outils natifs. - Écris le contenu **complet** demandé dans l'argument `content` (ne le tronque pas, pas de « … » ni de ligne omise). - Pour créer un fichier dans un nouveau dossier, un seul appel `create_file` avec le chemin complet suffit (les dossiers parents sont créés automatiquement). - 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). """ # Backwards-compatible alias (classic chat prompt). GENERAL_SYSTEM_PROMPT = GENERAL_SYSTEM_HEADER + GENERAL_ACTION_TEXT_PROTOCOL def _format_app_context(app_context: dict[str, Any] | None, recent_files: list[dict[str, Any]] | None) -> str: """Render the live application state for the General assistant prompt. The General assistant has no document context; without this block it only knows the app exists. Passing what the user currently sees (open documents, current directory, active search, recently modified files) lets it answer "résume ce que je fais / où j'en suis" style questions. """ app_context = app_context or {} lines: list[str] = [] vault = app_context.get("vault") or app_context.get("current_vault") if vault: lines.append(f"- Vault sélectionné : {vault}") directory = app_context.get("directory") if directory: lines.append(f"- Répertoire courant : {directory}") current_path = app_context.get("current_path") if current_path: lines.append(f"- Document affiché dans le viewer : {current_path}") docs = app_context.get("open_documents") or [] rendered_docs = [] for doc in docs: if not isinstance(doc, dict) or not doc.get("path"): continue rendered_docs.append(f"{doc['path']} (vault {doc['vault']})" if doc.get("vault") else str(doc["path"])) if rendered_docs: lines.append("- Documents ouverts dans les onglets/panneaux : " + ", ".join(rendered_docs)) editing = app_context.get("editing") if isinstance(editing, dict) and editing.get("path"): surface = "Forge" if editing.get("surface") == "forge" else "l'éditeur" location = f"{editing['path']} (vault {editing['vault']})" if editing.get("vault") else str(editing["path"]) lines.append( f"- Document en cours d'édition dans {surface} : {location} — c'est le document affiché à la " "place de la vue lecture. Pour le mettre à jour, utilise les outils d'écriture " "(`edit_file`, `append_to_file`) : la modification est rechargée automatiquement dans " "l'éditeur et la vue lecture dès l'exécution de l'outil." ) query = app_context.get("search_query") if query: total = app_context.get("search_total") suffix = f" ({total} résultat(s))" if isinstance(total, int) else "" lines.append(f"- Recherche en cours : « {query} »{suffix}") results = app_context.get("search_results") or [] rendered_results = [str(r.get("path")) for r in results[:10] if isinstance(r, dict) and r.get("path")] if rendered_results: lines.append("- Premiers résultats affichés : " + ", ".join(rendered_results)) if recent_files: rendered_recent = [ f"{f.get('vault')}/{f.get('path')}" for f in recent_files[:10] if f.get("path") ] if rendered_recent: lines.append("- Fichiers récemment modifiés : " + ", ".join(rendered_recent)) if not lines: return "" return ( "## Contexte applicatif actuel\n" "Voici ce que l'utilisateur voit ou fait en ce moment dans ObsiGate. " "Sers-t'en pour comprendre sa demande ; ne le répète pas inutilement.\n" + "\n".join(lines) + "\n" ) def build_general_system_prompt( vaults: list[str] | None = None, app_context: dict[str, Any] | None = None, recent_files: list[dict[str, Any]] | None = None, agent: bool = False, ) -> str: """System prompt for the General assistant (app help + actions). ``app_context`` carries the live UI state (open documents, current directory, active search) and ``recent_files`` the last modified files, so the assistant knows what the user is doing rather than answering blind. ``agent`` selects the action protocol: the classic chat endpoint (no native tools) uses the text ``obsigate-action`` blocks, while the tool-calling agent endpoint must invoke the native tools instead (BUG-053). """ protocol = GENERAL_ACTION_TOOL_PROTOCOL if agent else GENERAL_ACTION_TEXT_PROTOCOL prompt = GENERAL_SYSTEM_HEADER + protocol if vaults: prompt += "\nVaults disponibles : " + ", ".join(sorted(vaults)) + "\n" else: prompt += "\nAucun vault n'est actuellement configuré.\n" block = _format_app_context(app_context, recent_files) if block: prompt += "\n" + block 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)