Refactor AI Assist to background-task + buffer architecture
The assist flow previously tied the entire LLM generation to a single
POST request with no keepalives, causing NS_ERROR_NET_PARTIAL_TRANSFER
in Firefox when Hypercorn closed the connection during gaps between
chunks. This refactor decouples generation into a background task with
a buffer and a separate SSE stream — the same pattern used by chat.
- generation_buffer.py: Widen _buffers to support string keys, add
create/get/remove_assist_buffer() using "assist:{user_id}" keys,
fix cleanup log format for string keys
- generation_task.py: Add run_assist_generation() — lightweight
background task with no DB persistence or title generation
- notes.py: Replace single POST SSE route with POST /api/notes/assist
(returns 202) + GET /api/notes/assist/stream (SSE with 15s keepalives
and Last-Event-ID reconnection); 409 if already running
- useAssist.ts: Switch from apiStreamPost to apiPost + apiSSEStream
two-step pattern with named event mapping and stream handle cleanup
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1,4 +1,4 @@
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import date
|
||||
|
||||
@@ -7,7 +7,12 @@ from quart import Blueprint, Response, jsonify, request
|
||||
from fabledassistant.auth import login_required, get_current_user_id
|
||||
from fabledassistant.config import Config
|
||||
from fabledassistant.services.assist import build_assist_messages
|
||||
from fabledassistant.services.llm import stream_chat
|
||||
from fabledassistant.services.generation_buffer import (
|
||||
GenerationState,
|
||||
create_assist_buffer,
|
||||
get_assist_buffer,
|
||||
)
|
||||
from fabledassistant.services.generation_task import run_assist_generation
|
||||
from fabledassistant.services.notes import (
|
||||
convert_note_to_task,
|
||||
convert_task_to_note,
|
||||
@@ -191,7 +196,7 @@ async def get_backlinks_route(note_id: int):
|
||||
@notes_bp.route("/assist", methods=["POST"])
|
||||
@login_required
|
||||
async def assist_route():
|
||||
"""Stream an AI-assisted section edit via SSE."""
|
||||
"""Launch a background assist generation task."""
|
||||
uid = get_current_user_id()
|
||||
data = await request.get_json()
|
||||
|
||||
@@ -202,22 +207,50 @@ async def assist_route():
|
||||
if not target_section or not instruction:
|
||||
return jsonify({"error": "target_section and instruction are required"}), 400
|
||||
|
||||
# Reject if an assist generation is already running for this user
|
||||
existing = get_assist_buffer(uid)
|
||||
if existing and existing.state == GenerationState.RUNNING:
|
||||
return jsonify({"error": "Assist generation already running"}), 409
|
||||
|
||||
model = await get_setting(uid, "default_model", Config.OLLAMA_MODEL)
|
||||
messages = build_assist_messages(body, target_section, instruction)
|
||||
|
||||
async def generate():
|
||||
full_text = ""
|
||||
try:
|
||||
async for chunk in stream_chat(messages, model):
|
||||
full_text += chunk
|
||||
yield f"data: {json.dumps({'chunk': chunk})}\n\n"
|
||||
yield f"data: {json.dumps({'done': True, 'full_text': full_text})}\n\n"
|
||||
except Exception as e:
|
||||
logger.exception("Assist generation failed")
|
||||
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
||||
buf = create_assist_buffer(uid)
|
||||
asyncio.create_task(run_assist_generation(buf, messages, model))
|
||||
|
||||
return jsonify({"status": "started"}), 202
|
||||
|
||||
|
||||
@notes_bp.route("/assist/stream", methods=["GET"])
|
||||
@login_required
|
||||
async def assist_stream_route():
|
||||
"""SSE endpoint that tails the assist generation buffer."""
|
||||
uid = get_current_user_id()
|
||||
|
||||
buf = get_assist_buffer(uid)
|
||||
if buf is None:
|
||||
return jsonify({"error": "No active assist generation"}), 404
|
||||
|
||||
last_id_str = request.headers.get("Last-Event-ID") or request.args.get("last_event_id")
|
||||
last_id = int(last_id_str) if last_id_str is not None else -1
|
||||
|
||||
async def stream():
|
||||
cursor = last_id
|
||||
while True:
|
||||
pending = buf.events_after(cursor)
|
||||
for event in pending:
|
||||
yield buf.format_sse(event)
|
||||
cursor = event.index
|
||||
|
||||
if buf.state != GenerationState.RUNNING:
|
||||
break
|
||||
|
||||
got_new = await buf.wait_for_event(cursor, timeout=15.0)
|
||||
if not got_new:
|
||||
yield ": keepalive\n\n"
|
||||
|
||||
return Response(
|
||||
generate(),
|
||||
stream(),
|
||||
content_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
|
||||
Reference in New Issue
Block a user