From 08d27f0007db23e9093ad4a64e2c7d93a18c0a10 Mon Sep 17 00:00:00 2001 From: Victor Giers Date: Mon, 25 Aug 2025 23:56:26 +0200 Subject: [PATCH] feat: add message editing and regeneration endpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add PUT /sessions/{session_id}/messages/{index} to edit a user message and prune all subsequent messages. Add POST /sessions/{session_id}/regenerate to regenerate an assistant response from a given index (supports optional streaming). Introduce EditMessageRequest and RegenerateRequest schemas. Update the frontend to: • allow in‑place editing of user messages with an auto‑focusing textarea; • copy messages to clipboard; • regenerate assistant responses via a button. Add new CSS rules for message options bars and the editing textarea. --- backend/main.py | 109 ++++++++++++++++++++ backend/schemas.py | 8 ++ src/App.jsx | 244 ++++++++++++++++++++++++++++++++++++++++++--- src/styles.css | 67 +++++++++++++ 4 files changed, 415 insertions(+), 13 deletions(-) diff --git a/backend/main.py b/backend/main.py index 22bb0fc..4bc3240 100644 --- a/backend/main.py +++ b/backend/main.py @@ -167,4 +167,113 @@ def rename_session(session_id: str, req: schemas.GenerateTitleResponse, db: Sess db.commit() return {"ok": True} +@app.put("/sessions/{session_id}/messages/{index}") +def update_user_message(session_id: str, index: int, req: schemas.EditMessageRequest, db: Session = Depends(get_db)): + session = db.query(models.ChatSession).filter(models.ChatSession.session_id == session_id).first() + if not session: + raise HTTPException(status_code=404, detail="Session not found") + + msgs = ( + db.query(models.ChatMessage) + .filter(models.ChatMessage.session_pk == session.id) + .order_by(models.ChatMessage.created_at.asc()) + .all() + ) + + if index < 0 or index >= len(msgs): + raise HTTPException(status_code=404, detail="Message index out of range") + + # Only user messages can be edited per spec + if msgs[index].role != "user": + raise HTTPException(status_code=400, detail="Only user messages can be edited") + + # Update the content + msgs[index].content = req.message + + # Drop everything after the edited message + for m in msgs[index + 1:]: + db.delete(m) + + db.commit() + return {"ok": True} + +# ADD or REPLACE this whole function +@app.post("/sessions/{session_id}/regenerate") +async def regenerate(session_id: str, req: schemas.RegenerateRequest, db: Session = Depends(get_db)): + """ + Regenerate an assistant response for the conversation state at/before req.index. + If req.index points at an assistant message, we regenerate from the preceding user message. + """ + idx = req.index + model = req.model + stream = bool(req.stream) + + session = db.query(models.ChatSession).filter(models.ChatSession.session_id == session_id).first() + if not session: + raise HTTPException(status_code=404, detail="Session not found") + + msgs = ( + db.query(models.ChatMessage) + .filter(models.ChatMessage.session_pk == session.id) + .order_by(models.ChatMessage.created_at.asc()) + .all() + ) + + if idx < 0 or idx >= len(msgs): + raise HTTPException(status_code=400, detail="Invalid message index") + + # Find the last user message at/before idx + last_user_idx = idx + for i in range(idx, -1, -1): + if msgs[i].role == "user": + last_user_idx = i + break + + # Prune everything after last_user_idx + if last_user_idx < len(msgs) - 1: + for m in msgs[last_user_idx + 1:]: + db.delete(m) + db.commit() + + # Build the conversation up to & incl. the last user message + conversation = [{"role": m.role, "content": m.content} for m in msgs[: last_user_idx + 1]] + + # Avoid DetachedInstanceError during streaming + session_pk = session.id + + if stream: + async def stream_generator(): + full_reply = "" + try: + # ollama_chat_stream must already exist in your codebase (used by /chat) + async for chunk in ollama_chat_stream(model, conversation): + full_reply += chunk + yield chunk + except Exception as e: + yield f"Ollama error: {e}" + # Persist with a fresh DB session (streaming context) + try: + db_sess = SessionLocal() + db_sess.add(models.ChatMessage(session_pk=session_pk, role="assistant", content=full_reply)) + db_sess.commit() + finally: + try: + db_sess.close() + except Exception: + pass + + return StreamingResponse(stream_generator(), media_type="text/plain") + + # Non-streaming + try: + # ollama_chat must already exist in your codebase (used by /chat) + reply = await ollama_chat(model, conversation) + except Exception as e: + raise HTTPException(status_code=502, detail=f"Ollama error: {e}") + + db.add(models.ChatMessage(session_pk=session_pk, role="assistant", content=reply)) + db.commit() + return {"reply": reply} + + # To run standalone: python -m uvicorn backend.main:app --host 127.0.0.1 --port 8000 diff --git a/backend/schemas.py b/backend/schemas.py index 9bcb0c4..fdd53b9 100644 --- a/backend/schemas.py +++ b/backend/schemas.py @@ -40,3 +40,11 @@ class ChatSession(BaseModel): class SessionsResponse(BaseModel): sessions: List[ChatSession] + +class EditMessageRequest(BaseModel): + message: str + +class RegenerateRequest(BaseModel): + index: int + model: Optional[str] = None + stream: bool = True \ No newline at end of file diff --git a/src/App.jsx b/src/App.jsx index 67845e4..40ffb0a 100644 --- a/src/App.jsx +++ b/src/App.jsx @@ -45,9 +45,7 @@ function splitThinkBlocks(text) { // Renders assistant message with a collapsible "Thoughts" block (if present) function AssistantMessageContent({ content, streamOutput }) { const { think, answer } = splitThinkBlocks(content || ''); - const [open, setOpen] = React.useState(false); // Closed by default - - // If streaming, the button should appear as soon as 'think' content is detected + const [open, setOpen] = React.useState(false); const showThinkButton = !!think; return ( @@ -60,15 +58,24 @@ function AssistantMessageContent({ content, streamOutput }) { aria-expanded={open ? 'true' : 'false'} aria-controls="think-content" > - + Thoughts {open && ( -
+
)}
)} -
+
); } @@ -100,6 +107,177 @@ export default function App() { const [loading, setLoading] = useState(true); // Loading state for initial session fetch const [unreadSessions, setUnreadSessions] = useState([]); // Track unread messages const [scrollPositions, setScrollPositions] = useState({}); // Store scroll positions for each session + // Editing state for user messages + const [editingMessageIndex, setEditingMessageIndex] = useState(null); + const [editText, setEditText] = useState(''); + // Helpers + handlers for message copy/edit/regenerate (must live inside App) + function getVisibleTextForCopy(message) { + let raw = message.content || ''; + if (message.role === 'assistant') { + try { + const { think, answer } = splitThinkBlocks(raw); + raw = answer || raw; + } catch (_) {} + } + const html = markdownToHTML(raw); + const tempDiv = document.createElement('div'); + tempDiv.innerHTML = html; + tempDiv.querySelectorAll('br').forEach(br => br.replaceWith('\n')); + tempDiv.querySelectorAll('p,div,li,pre,code,h1,h2,h3,h4,h5,h6,table,tr').forEach(el => { + el.insertAdjacentText('beforeend', '\n'); + }); + return tempDiv.innerText.replace(/\n{3,}/g, '\n\n').trim(); + } + + async function handleCopyMessage(message) { + try { + await navigator.clipboard.writeText(getVisibleTextForCopy(message)); + } catch (err) { + console.error('Failed to copy message:', err); + } + } + + function startEditMessage(index, content) { + setEditingMessageIndex(index); + setEditText(content || ''); + } + function cancelEditMessage() { + setEditingMessageIndex(null); + setEditText(''); + } + + async function commitEditMessage(index) { + const original = (messages[index]?.content || '').trim(); + const next = (editText || '').trim(); + if (next === original) { + cancelEditMessage(); + return; + } + const sessionId = activeSessionId; + if (!sessionId) return; + + // Optimistically update UI: set edited content and prune following messages + setChatSessions(prev => + prev.map(s => { + if (s.session_id !== sessionId) return s; + const old = s.messages || []; + // keep up to index (inclusive) and update that item + const updated = old.slice(0, index + 1).map((m, j) => + j === index ? { ...m, content: next } : m + ); + return { ...s, messages: updated }; + }) + ); + + setEditingMessageIndex(null); + setEditText(''); + + // ⬇️ Scroll the chat frame to the bottom after the DOM updates + requestAnimationFrame(() => scrollToBottom('auto', sessionId)); // use 'smooth' if you prefer + + try { + const resp = await fetch(`${ollamaApiUrl}/sessions/${sessionId}/messages/${index}`, { + method: 'PUT', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ message: next }) + }); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + } catch (err) { + // Roll back to original content on failure + console.error('Failed to update message:', err); + setChatSessions(prev => + prev.map(s => { + if (s.session_id !== sessionId) return s; + const old = s.messages || []; + const restored = old.map((m, j) => + j === index ? { ...m, content: original } : m + ); + return { ...s, messages: restored }; + }) + ); + return; // don't regenerate on failure + } + + // Continue conversation from the edited message + await regenerateFromIndex(index); + } + + async function regenerateFromIndex(index) { + const sessionId = activeSessionId; + if (!sessionId || typeof index !== 'number') return; + + const msgs = (chatSessions.find(s => s.session_id === sessionId)?.messages) || []; + let lastUserIdx = index; + for (let i = index; i >= 0; i--) { + if (msgs[i]?.role === 'user') { lastUserIdx = i; break; } + } + + // Prune UI to lastUserIdx + setChatSessions(prev => + prev.map(s => s.session_id === sessionId + ? { ...s, messages: (s.messages || []).slice(0, lastUserIdx + 1) } + : s + ) + ); + + setIsSending(true); + + if (streamOutput) { + const assistantMsgId = `msg-${Date.now()}-${Math.random()}`; + setChatSessions(prev => + prev.map(s => s.session_id === sessionId + ? { ...s, messages: [...(s.messages || []), { id: assistantMsgId, role: 'assistant', content: '' }] } + : s + ) + ); + try { + const res = await fetch(`${ollamaApiUrl}/sessions/${sessionId}/regenerate`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ index, model, stream: true }) + }); + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + let full = ''; + while (true) { + const { value, done } = await reader.read(); + if (done) break; + const chunk = decoder.decode(value, { stream: true }); + full += chunk; + setChatSessions(prev => + prev.map(s => s.session_id === sessionId + ? { ...s, messages: (s.messages || []).map(m => m.id === assistantMsgId ? { ...m, content: full } : m) } + : s + ) + ); + } + } catch (e) { + console.error(e); + } finally { + setIsSending(false); + } + } else { + try { + const res = await fetch(`${ollamaApiUrl}/sessions/${sessionId}/regenerate`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ index, model, stream: false }) + }); + const data = await res.json(); + setChatSessions(prev => + prev.map(s => s.session_id === sessionId + ? { ...s, messages: [...(s.messages || []), { role: 'assistant', content: data.reply, id: `msg-${Date.now()}` }] } + : s + ) + ); + } catch (e) { + console.error(e); + } finally { + setIsSending(false); + } + } + } + // Persist userScrolledUp state per session + live ref for closures (streaming) const [userScrolledUpState, setUserScrolledUpState] = useState({}); @@ -891,13 +1069,53 @@ export default function App() {
{messages.map((m, i) => ( -
- {m.role === 'assistant' - ? - :
- } -
- ))} +
+ {m.role === 'assistant' ? ( +
+ + {!isSending && ( +
+ + +
+ )} +
+ ) : ( +
+ {editingMessageIndex === i ? ( + setEditText(e.target.value)} + onBlur={cancelEditMessage} + onKeyDown={(e) => { + if (e.key === 'Escape') { e.preventDefault(); cancelEditMessage(); } + if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); commitEditMessage(i); } + }} + autoFocus + maxRows={13} + /> + ) : ( +
+ )} + {!isSending && ( +
+ + +
+ )} +
+ )} +
+))}
{/* New message tip (active chat only) */} diff --git a/src/styles.css b/src/styles.css index 2820140..c6365ca 100644 --- a/src/styles.css +++ b/src/styles.css @@ -326,12 +326,59 @@ textarea.input { max-width: 80%; border: 1px solid var(--border); margin-right: 5px; + margin-bottom: 15px; } .msg.assistant { background: transparent; border: none; max-width: none; /* Remove max-width for assistant messages */ animation: fadeIn 0.3s ease-in-out; + margin-bottom: 30px; +} + +.user-message-wrapper { + display: flex; + flex-direction: column; + align-items: flex-end; /* Align content and options to the right */ + position: relative; +} + +.assistant-message-wrapper { + display: flex; + flex-direction: column; + align-items: flex-start; /* Align content and options to the left */ + position: relative; +} + +.message-options-bar { + display: flex; + gap: 2px; /* Smaller gap between icons */ + padding: 0; /* No padding for the container */ + border-radius: 8px; + background-color: transparent; /* No background for the container */ + border: none; /* No border for the container */ + position: absolute; + bottom: -30px; /* Moved further down */ + /* Always visible, no opacity or pointer-events changes */ +} + +.user-options { + right: -10px; /* Align to the right edge of the message content padding */ + bottom: -40px; +} + +.assistant-options { + left: 0px; /* Align to the left edge of the message content padding */ +} + +.message-options-bar .icon-button { + padding: 4px; /* Smaller padding for individual icons */ + border-radius: 4px; + border: none; /* No border for individual icons */ +} + +.message-options-bar .icon-button:hover { + background-color: var(--hover-bg); } /* new message pill */ @@ -681,3 +728,23 @@ input:checked + .slider:before { } .msg-content { /* container for assistant/user rendered markdown */ } + + +/* Editing textarea styled like user bubble */ +.edit-message-input { + background: var(--user-msg-bg); + border: 1px solid var(--border); + padding: 12px 14px; + border-radius: 12px; + margin-left: auto; + margin-right: 5px; + margin-bottom: 15px; + max-width: 80%; + width: 100%; + color: var(--text); + font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Inter, Helvetica, Arial; + line-height: 1.5; + resize: none; + overflow-y: auto; + display: block; +}