ObsiGate/backend/bookslm_routes.py

"""BooksLM API routes — directory-scoped AI chat for Obsidian vaults."""

import json
import logging
from pathlib import Path

from fastapi import APIRouter, Depends, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from backend.auth.middleware import check_vault_access, require_auth
from backend.bookslm import (
    build_general_system_prompt,
    build_system_prompt,
    collect_directory_context,
    collect_files_context,
    empty_context,
)
from backend.indexer import get_vault_data, index
from backend.schemas import BooksLMContextResponse

logger = logging.getLogger("obsigate.bookslm_routes")
router = APIRouter(prefix="/api/ai/bookslm", tags=["BooksLM"])

VALID_MODES = {"directory", "documents", "general"}


# ── Request models ──


class BooksLMContextRequest(BaseModel):
    vault: str | None = Field(default=None, description="Vault name (required for directory/documents modes)")
    directory: str = Field(default="", description="Relative directory path within the vault")
    mode: str = Field(default="directory", description="Context mode: 'directory', 'documents' or 'general'")
    context_files: list[str] = Field(
        default_factory=list,
        description="Relative file paths to use as context (documents mode)",
    )


class BooksLMChatRequest(BaseModel):
    vault: str | None = Field(default=None, description="Vault name (required for directory/documents modes)")
    directory: str = Field(default="", description="Relative directory path within the vault")
    mode: str = Field(default="directory", description="Context mode: 'directory', 'documents' or 'general'")
    context_files: list[str] = Field(
        default_factory=list,
        description="Relative file paths to use as context (documents mode)",
    )
    message: str = Field(description="User message")
    conversation_history: list[dict[str, str]] = Field(
        default_factory=list,
        description="Previous conversation turns [{role, content}]",
    )
    provider: str | None = Field(
        default=None,
        description="AI provider override (e.g. 'deepseek', 'openrouter', 'gemini', 'nvidia', 'xiaomi', 'mistral', 'qwencloud'). "
                    "If not set, uses DEFAULT_PROVIDER.",
    )
    model: str | None = Field(
        default=None,
        description="Model name to use for this request. If not set, uses the provider's default model.",
    )


def _normalize_mode(mode: str | None) -> str:
    mode = (mode or "directory").lower()
    return mode if mode in VALID_MODES else "directory"


def _resolve_vault_path(vault: str | None, current_user):
    """Validate vault access and return (vault_name, vault_path)."""
    if not vault:
        raise HTTPException(status_code=400, detail="Champ 'vault' requis pour ce contexte")
    if not check_vault_access(vault, current_user):
        raise HTTPException(status_code=403, detail=f"Accès refusé à la vault '{vault}'")
    vault_data = get_vault_data(vault)
    if not vault_data:
        raise HTTPException(status_code=404, detail=f"Vault '{vault}' not found")
    return vault, Path(vault_data["path"])


def _build_context(mode: str, vault_path: Path | None, directory: str, context_files: list[str]):
    """Collect the context payload for the requested mode."""
    if mode == "general":
        return empty_context("general")
    if mode == "documents":
        ctx = collect_files_context(vault_path, context_files, scope="documents")  # type: ignore[arg-type]
        # No open document could be read → degrade gracefully to General.
        return ctx if ctx["file_count"] else empty_context("general")
    return collect_directory_context(vault_path, directory)  # type: ignore[arg-type]


# ── Endpoints ──


@router.post("/context", response_model=BooksLMContextResponse)
async def api_bookslm_context(
    req: BooksLMContextRequest,
    current_user=Depends(require_auth),
):
    """Collect context for the AI assistant.

    Three modes are supported:

    * ``directory`` — every supported file under a vault directory;
    * ``documents`` — only the explicitly listed open documents;
    * ``general`` — no document context (assistant for the app itself).
    """
    mode = _normalize_mode(req.mode)
    if mode == "general":
        return empty_context("general")

    _resolve_vault_path(req.vault, current_user)
    vault_path = Path(get_vault_data(req.vault)["path"])  # type: ignore[index]
    return _build_context(mode, vault_path, req.directory, req.context_files)


@router.post(
    "/chat",
    response_class=StreamingResponse,
    responses={200: {"content": {"text/event-stream": {}}, "description": "SSE token stream"}},
)
async def api_bookslm_chat(
    req: BooksLMChatRequest,
    current_user=Depends(require_auth),
):
    """Chat with the AI assistant.

    Builds a system prompt from the selected context (directory, open
    documents or general app knowledge), then streams the provider's answer
    as Server-Sent Events.
    """
    mode = _normalize_mode(req.mode)

    vault_path: Path | None = None
    if mode != "general":
        _resolve_vault_path(req.vault, current_user)
        vault_path = Path(get_vault_data(req.vault)["path"])  # type: ignore[index]

    # Collect context
    context = _build_context(mode, vault_path, req.directory, req.context_files)
    effective_mode = context.get("scope", mode)

    # Build system prompt
    if effective_mode == "general":
        system_prompt = build_general_system_prompt(list(index.keys()))
    elif effective_mode == "documents":
        system_prompt = build_system_prompt(context, scope="documents")
    else:
        if context["file_count"] == 0:
            raise HTTPException(status_code=404, detail="Aucun fichier markdown trouvé dans ce dossier")
        system_prompt = build_system_prompt(context, scope="directory")

    # Call AI provider
    from backend.ai import DEFAULT_PROVIDER, PROVIDERS, _call_deepseek_openrouter, _call_gemini

    # Build messages with conversation history
    messages_text = ""
    if req.conversation_history:
        for turn in req.conversation_history:
            role = turn.get("role", "user")
            content = turn.get("content", "")
            if role == "user":
                messages_text += f"\n\nUtilisateur : {content}"
            elif role == "assistant":
                messages_text += f"\n\nAssistant : {content}"

    # Current message
    user_prompt = req.message
    if messages_text:
        user_prompt = f"Historique de la conversation :{messages_text}\n\nQuestion actuelle : {req.message}"

    async def generate_sse():
        try:
            # Resolve provider: explicit override wins, else default.
            # Fall back to first available if the requested one isn't configured.
            cfg_name = (req.provider or DEFAULT_PROVIDER).lower()
            if cfg_name not in PROVIDERS or not PROVIDERS[cfg_name].get("api_key"):
                # Try next available provider
                for pname, pcfg in PROVIDERS.items():
                    if pcfg.get("api_key") and pname != "gemini":
                        cfg_name = pname
                        break
                else:
                    # No provider available at all
                    err = "Aucun fournisseur AI configuré (clés API manquantes)"
                    error_data = json.dumps({"error": err}, ensure_ascii=False)
                    yield f"event: error\ndata: {error_data}\n\n"
                    return

            # Optional per-request model override
            original_model = None
            if req.model and cfg_name in PROVIDERS:
                original_model = PROVIDERS[cfg_name].get("model")
                PROVIDERS[cfg_name]["model"] = req.model
            try:
                if cfg_name == "gemini":
                    response = await _call_gemini(user_prompt, system_prompt, temperature=0.3, max_tokens=4096)
                else:
                    response = await _call_deepseek_openrouter(
                        user_prompt, system_prompt,
                        provider=cfg_name,
                        temperature=0.3,
                        max_tokens=4096,
                    )
            finally:
                # Restore the original model so other calls aren't affected
                if original_model is not None and cfg_name in PROVIDERS:
                    PROVIDERS[cfg_name]["model"] = original_model

            # Send the full response as a single SSE event
            data = json.dumps({"token": response, "provider": cfg_name, "model": req.model or PROVIDERS.get(cfg_name, {}).get("model", "")}, ensure_ascii=False)
            yield f"event: message\ndata: {data}\n\n"
            yield "event: done\ndata: {}\n\n"
        except Exception as e:
            logger.error(f"BooksLM chat error: {e}")
            error_data = json.dumps({"error": str(e)}, ensure_ascii=False)
            yield f"event: error\ndata: {error_data}\n\n"

    return StreamingResponse(
        generate_sse(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        },
    )