feat: add message editing and regeneration endpoints

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.
2025-08-25 23:56:26 +02:00
parent d28d88d1f2
commit 08d27f0007
4 changed files with 415 additions and 13 deletions
--- 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
--- 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
--- 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"
          >
-            <span className="think-toggle-icon" aria-hidden="true">{open ? '▾' : '▸'}</span>
+            <span className="think-toggle-icon" aria-hidden="true">
+              {open ? '▾' : '▸'}
+            </span>
            Thoughts
          </button>
          {open && (
-            <div id="think-content" className="think-content" dangerouslySetInnerHTML={{ __html: markdownToHTML(think) }} />
+            <div
+              id="think-content"
+              className="think-content"
+              dangerouslySetInnerHTML={{ __html: markdownToHTML(think) }}
+            />
          )}
        </div>
      )}
-      <div className="msg-content" dangerouslySetInnerHTML={{ __html: markdownToHTML(answer || content || '') }} />
+      <div
+        className="msg-content"
+        dangerouslySetInnerHTML={{ __html: markdownToHTML(answer || content || '') }}
+      />
    </div>
  );
 }
@@ -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() {

            <div key={activeSessionId} className="chat" ref={chatRef} onClick={handleChatFrameClick}>
              {messages.map((m, i) => (
-                <div key={m.id || i} id={m.id} className={'msg ' + (m.role === 'user' ? 'user' : 'assistant')}>
-                  {m.role === 'assistant'
-                    ? <AssistantMessageContent content={m.content} streamOutput={streamOutput} />
-                    : <div className="msg-content" dangerouslySetInnerHTML={{ __html: markdownToHTML(m.content) }} />
-                  }
-                </div>
-              ))}
+  <div key={m.id || i} id={m.id} className={'msg ' + (m.role === 'user' ? 'user' : 'assistant')}>
+    {m.role === 'assistant' ? (
+      <div className="assistant-message-wrapper">
+        <AssistantMessageContent content={m.content} streamOutput={streamOutput} />
+        {!isSending && (
+          <div className="message-options-bar assistant-options">
+            <button className="icon-button" title="Copy message" onClick={() => handleCopyMessage(m)}>
+              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>
+            </button>
+            <button className="icon-button" title="Regenerate response" onClick={() => regenerateFromIndex(i)}>
+              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><path d="M21.5 2v6h-6M2.5 22v-6h6M2 11.5a10 10 0 0 1 18.8-4.3M22 12.5a10 10 0 0 1-18.8 4.3"></path></svg>
+            </button>
+          </div>
+        )}
+      </div>
+    ) : (
+      <div className="user-message-wrapper">
+        {editingMessageIndex === i ? (
+          <TextareaAutosize
+            className="edit-message-input"
+            value={editText}
+            onChange={(e) => 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}
+          />
+        ) : (
+          <div className="msg-content" dangerouslySetInnerHTML={{ __html: markdownToHTML(m.content) }} />
+        )}
+        {!isSending && (
+          <div className="message-options-bar user-options">
+            <button className="icon-button" title="Edit message" onClick={() => startEditMessage(i, m.content)}>
+              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><path d="M12 20h9"></path><path d="M16.5 3.5a2.121 2.121 0 0 1 3 3L7 19l-4 1 1-4L16.5 3.5z"></path></svg>
+            </button>
+            <button className="icon-button" title="Copy message" onClick={() => handleCopyMessage(m)}>
+              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>
+            </button>
+          </div>
+        )}
+      </div>
+    )}
+  </div>
+))}
            </div>

            {/* New message tip (active chat only) */}
--- 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;
+}