pablan/backend/app/api/authoring/refine.py

"""Section refinement for the writing editor.

The user writes Markdown; after a pause the client asks the model to refine the
section the cursor is in. The whole document is context, but the model
regenerates ONLY that section (FIM-style), streamed back as SSE so the
suggestion appears progressively and can be aborted the moment the user resumes
typing.

The request body and the streamed response carry document text. That is fine on
this owner-scoped endpoint — the same trust boundary as
`GET /api/documents/{id}` — but nothing here logs content: metadata only.
"""

import asyncio
import logging
import time
import uuid
from collections.abc import AsyncIterator
from typing import Annotated

from fastapi import Depends
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field, ValidationError
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.authoring import grounding
from app.api.authoring.routing import authoring_router
from app.api.documents import readable_document, require_editor
from app.api.sse import sse
from app.auth.deps import get_current_user
from app.authoring.prompts import render_refine_prompt
from app.authoring.schema import AuthoringTemplate
from app.authoring.sections import ActiveSection, active_section, slice_lines
from app.db import get_db
from app.llm.client import NO_THINKING, chat_stream
from app.llm.errors import LLMError
from app.models import Template, User

router = authoring_router()
logger = logging.getLogger("pablan.authoring")


class RefineRequest(BaseModel):
    content_md: str = Field(max_length=100_000)
    cursor_line: int = Field(ge=1)


async def _template_for(
    db: AsyncSession, template_config_id: str | None
) -> AuthoringTemplate | None:
    """The blueprint a document was started from, if it still exists and still
    parses — it carries the persona, the temperature and the per-section hints
    that shape a refinement."""
    if not template_config_id:
        return None
    row = (
        await db.execute(
            select(Template).where(Template.config["id"].astext == template_config_id)
        )
    ).scalar_one_or_none()
    if row is None:
        return None
    try:
        return AuthoringTemplate.model_validate(row.config)
    except ValidationError:
        return None


@router.post("/{document_id}/refine")
async def refine_section(
    document_id: uuid.UUID,
    body: RefineRequest,
    user: Annotated[User, Depends(get_current_user)],
    db: Annotated[AsyncSession, Depends(get_db)],
) -> StreamingResponse:
    """Stream a matured version of the section at the cursor.

    SSE frames: one `section` frame with the exact line range the suggestion
    replaces, then `token` frames, then `done` (or `error`)."""
    document = await readable_document(db, document_id, user)
    require_editor(document, user)

    section = active_section(body.content_md, body.cursor_line)
    prefix, section_text, suffix = slice_lines(
        body.content_md, section.start_line, section.end_line
    )
    meta = document.meta or {}

    persona: str | None = None
    hint: str | None = None
    temperature = 0.4
    template = await _template_for(db, meta.get("template"))
    if template is not None:
        persona = template.persona
        temperature = template.model.temperature
        heading = grounding.leading_heading(section_text)
        if heading:
            hint = template.hint_for(heading)

    # Related, already-published knowledge the author may read. Rendered into
    # the prompt as grounding, AND surfaced to the editor's "?" inspector so the
    # author can see where a suggestion drew from.
    chunks = await grounding.for_section(db, section_text, user, document.id)
    messages = render_refine_prompt(
        section_text,
        prefix=prefix,
        suffix=suffix,
        persona=persona,
        hint=hint,
        # Background from the chat this capture came from, if any. Like the
        # document text, it travels only on this owner-scoped call and is
        # never logged.
        context=meta.get("context"),
        knowledge=[
            grounding.reference(chunk.title, chunk.heading_path, chunk.content)
            for chunk in chunks
        ],
    )
    references = [
        {"title": chunk.title, "heading_path": chunk.heading_path} for chunk in chunks
    ]

    return StreamingResponse(
        _stream_refine(messages, section, temperature, str(document.id), references),
        media_type="text/event-stream",
        headers={"cache-control": "no-cache", "x-accel-buffering": "no"},
    )


async def _stream_refine(
    messages: list[dict[str, str]],
    section: ActiveSection,
    temperature: float,
    document_id: str,
    references: list[dict[str, str]],
) -> AsyncIterator[str]:
    started = time.monotonic()
    outcome = "ok"
    token_events = 0
    # First: which lines "Accept" will overwrite, so the client can bind the
    # suggestion to an exact range even as the model streams.
    yield sse(
        "section", {"start_line": section.start_line, "end_line": section.end_line}
    )
    # What the suggestion is grounding on (the author's own readable material) —
    # titles + heading paths only, for the "?" inspector. Content-safe.
    if references:
        yield sse("grounding", {"references": references})
    try:
        async for token in chat_stream(
            messages, role="chat", temperature=temperature, extra_body=NO_THINKING
        ):
            token_events += 1
            yield sse("token", {"text": token})
        yield sse("done", {})
    except LLMError as exc:
        outcome = "error"
        yield sse("error", {"code": exc.code})
    except (asyncio.CancelledError, GeneratorExit):
        outcome = "aborted"
        raise
    finally:
        logger.info(
            "refine finished",
            extra={
                "event": "refine",
                "outcome": outcome,
                "duration_ms": round((time.monotonic() - started) * 1000),
                "token_events": token_events,
                "document_id": document_id,
            },
        )