add ability to stop streaming

This commit is contained in:
Josh Hawkins 2026-04-09 09:05:10 -05:00
parent 01519607de
commit a46d0a4c68
2 changed files with 58 additions and 15 deletions

View File

@ -1328,6 +1328,9 @@ When a user refers to a specific object they have seen or describe with identify
async def stream_body_llm(): async def stream_body_llm():
nonlocal conversation, stream_tool_calls, stream_iterations nonlocal conversation, stream_tool_calls, stream_iterations
while stream_iterations < max_iterations: while stream_iterations < max_iterations:
if await request.is_disconnected():
logger.debug("Client disconnected, stopping chat stream")
return
logger.debug( logger.debug(
f"Streaming LLM (iteration {stream_iterations + 1}/{max_iterations}) " f"Streaming LLM (iteration {stream_iterations + 1}/{max_iterations}) "
f"with {len(conversation)} message(s)" f"with {len(conversation)} message(s)"
@ -1337,6 +1340,9 @@ When a user refers to a specific object they have seen or describe with identify
tools=tools if tools else None, tools=tools if tools else None,
tool_choice="auto", tool_choice="auto",
): ):
if await request.is_disconnected():
logger.debug("Client disconnected, stopping chat stream")
return
kind, value = event kind, value = event
if kind == "content_delta": if kind == "content_delta":
yield ( yield (
@ -1366,6 +1372,11 @@ When a user refers to a specific object they have seen or describe with identify
msg.get("content"), pending msg.get("content"), pending
) )
) )
if await request.is_disconnected():
logger.debug(
"Client disconnected before tool execution"
)
return
( (
executed_calls, executed_calls,
tool_results, tool_results,

View File

@ -1,6 +1,6 @@
import { Button } from "@/components/ui/button"; import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input"; import { Input } from "@/components/ui/input";
import { FaArrowUpLong } from "react-icons/fa6"; import { FaArrowUpLong, FaStop } from "react-icons/fa6";
import { LuCircleAlert } from "react-icons/lu"; import { LuCircleAlert } from "react-icons/lu";
import { useTranslation } from "react-i18next"; import { useTranslation } from "react-i18next";
import { useState, useCallback, useRef, useEffect, useMemo } from "react"; import { useState, useCallback, useRef, useEffect, useMemo } from "react";
@ -28,6 +28,7 @@ export default function ChatPage() {
const [error, setError] = useState<string | null>(null); const [error, setError] = useState<string | null>(null);
const [attachedEventId, setAttachedEventId] = useState<string | null>(null); const [attachedEventId, setAttachedEventId] = useState<string | null>(null);
const scrollRef = useRef<HTMLDivElement>(null); const scrollRef = useRef<HTMLDivElement>(null);
const abortRef = useRef<AbortController | null>(null);
useEffect(() => { useEffect(() => {
document.title = t("documentTitle"); document.title = t("documentTitle");
@ -70,12 +71,24 @@ export default function ChatPage() {
...(axios.defaults.headers.common as Record<string, string>), ...(axios.defaults.headers.common as Record<string, string>),
}; };
await streamChatCompletion(url, headers, apiMessages, { const controller = new AbortController();
updateMessages: (updater) => setMessages(updater), abortRef.current = controller;
onError: (message) => setError(message),
onDone: () => setIsLoading(false), await streamChatCompletion(
defaultErrorMessage: t("error"), url,
}); headers,
apiMessages,
{
updateMessages: (updater) => setMessages(updater),
onError: (message) => setError(message),
onDone: () => {
abortRef.current = null;
setIsLoading(false);
},
defaultErrorMessage: t("error"),
},
controller.signal,
);
}, },
[isLoading, t], [isLoading, t],
); );
@ -106,6 +119,12 @@ export default function ChatPage() {
[attachedEventId, input, isLoading, messages, submitConversation], [attachedEventId, input, isLoading, messages, submitConversation],
); );
const stopGeneration = useCallback(() => {
abortRef.current?.abort();
abortRef.current = null;
setIsLoading(false);
}, []);
const handleEditSubmit = useCallback( const handleEditSubmit = useCallback(
(messageIndex: number, newContent: string) => { (messageIndex: number, newContent: string) => {
const newList: ChatMessage[] = [ const newList: ChatMessage[] = [
@ -237,6 +256,7 @@ export default function ChatPage() {
attachedEventId={attachedEventId} attachedEventId={attachedEventId}
onClearAttachment={handleClearAttachment} onClearAttachment={handleClearAttachment}
onAttach={setAttachedEventId} onAttach={setAttachedEventId}
onStop={stopGeneration}
recentEventIds={recentEventIds} recentEventIds={recentEventIds}
/> />
)} )}
@ -254,6 +274,7 @@ type ChatEntryProps = {
attachedEventId: string | null; attachedEventId: string | null;
onClearAttachment: () => void; onClearAttachment: () => void;
onAttach: (eventId: string) => void; onAttach: (eventId: string) => void;
onStop: () => void;
recentEventIds: string[]; recentEventIds: string[];
}; };
@ -266,6 +287,7 @@ function ChatEntry({
attachedEventId, attachedEventId,
onClearAttachment, onClearAttachment,
onAttach, onAttach,
onStop,
recentEventIds, recentEventIds,
}: ChatEntryProps) { }: ChatEntryProps) {
const handleKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => { const handleKeyDown = (e: React.KeyboardEvent<HTMLInputElement>) => {
@ -306,14 +328,24 @@ function ChatEntry({
onKeyDown={handleKeyDown} onKeyDown={handleKeyDown}
aria-busy={isLoading} aria-busy={isLoading}
/> />
<Button {isLoading ? (
variant="select" <Button
className="size-10 shrink-0 rounded-full" variant="destructive"
disabled={!input.trim() || isLoading} className="size-10 shrink-0 rounded-full"
onClick={() => sendMessage()} onClick={onStop}
> >
<FaArrowUpLong size="16" /> <FaStop className="size-3" />
</Button> </Button>
) : (
<Button
variant="select"
className="size-10 shrink-0 rounded-full"
disabled={!input.trim()}
onClick={() => sendMessage()}
>
<FaArrowUpLong className="size-4" />
</Button>
)}
</div> </div>
</div> </div>
); );