ObsiGate/backend/bookslm.py

"""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,
    }

    # 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 = list(_cache)
    for k in to_remove:
        del _cache[k]
    return len(to_remove)