Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 51 additions & 11 deletions apps/frontend/src/api/chat/get-completion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,9 @@ export async function getCompletion(
const { handleError } = useErrorStore.getState();
const {
updateMessage,
addMessageToChat,
createPendingMessageInMemory,
persistPendingMessageToDb,
removePendingMessageFromMemory,
selectedLlmModel,
selectedChatTools,
} = useChatsStore.getState();
Expand All @@ -79,6 +81,10 @@ export async function getCompletion(
externalChatTools.includes(tool),
);

// Id of the optimistic placeholder, set once created — lets the catch
// block clean it up if the stream errors out.
let messageIdForCleanup: number | undefined;

try {
// Abort any existing stream before starting a new one
abortStreaming();
Expand Down Expand Up @@ -144,7 +150,9 @@ export async function getCompletion(
return;
}

const messageId = await addMessageToChat(currentChat, {
// Optimistic placeholder — persisted only once the stream settles
// (see `onFinish` below).
const localMessageId = createPendingMessageInMemory(currentChat, {
content: "",
type: "text",
role: "assistant",
Expand All @@ -156,6 +164,7 @@ export async function getCompletion(
open_data_citations: null,
external_tool_context: isExternalToolContext,
});
messageIdForCleanup = localMessageId;

let currentText = "";
let documentCitations: number[] = [];
Expand All @@ -168,7 +177,7 @@ export async function getCompletion(
const writeMessage = () =>
updateMessage({
chat: currentChat,
messageId,
messageId: localMessageId,
content: currentText,
citations: documentCitations.length ? documentCitations : null,
web_citations: webCitations.length ? webCitations : null,
Expand Down Expand Up @@ -212,12 +221,43 @@ export async function getCompletion(
openDataCitations = sources;
writeMessage();
},
onFinish: () => {
onFinish: async (wasSuccessful) => {
setStatus("idle");
setStreamingAbortController(null);

// Only persist if the stream finished cleanly with non-whitespace
// content — otherwise drop the placeholder.
if (wasSuccessful && currentText.trim()) {
try {
await persistPendingMessageToDb(currentChat, localMessageId, {
content: currentText,
type: "text",
role: "assistant",
allowed_document_ids: allowedDocumentIds,
allowed_folder_ids: selectedFolderIds,
citations: documentCitations.length ? documentCitations : null,
web_citations: webCitations.length ? webCitations : null,
parla_citations: parlaCitations.length ? parlaCitations : null,
open_data_citations: openDataCitations.length
? openDataCitations
: null,
external_tool_context: isExternalToolContext,
});
} catch (error) {
removePendingMessageFromMemory(currentChat, localMessageId);
handleError(error, span);
}
} else {
removePendingMessageFromMemory(currentChat, localMessageId);
}
},
});
} catch (error) {
// Stream never settled — drop the placeholder if one was created.
if (messageIdForCleanup !== undefined) {
removePendingMessageFromMemory(currentChat, messageIdForCleanup);
}

// Only handle error if it's not an abort error
const isUserAbort = error instanceof Error && error.name === "AbortError";
if (isUserAbort) {
Expand All @@ -232,25 +272,25 @@ export async function getCompletion(
}
}

function processStreamLine(
async function processStreamLine(
line: string,
callbacks: {
onTextDelta: (delta: string) => void;
onCitations: (chunkIds: number[]) => void;
onWebCitations: (webCitationSources: WebCitationSource[]) => void;
onParlaCitations: (sources: ParlaCitationSource[]) => void;
onOpenDataCitations: (sources: OpenDataCitationSource[]) => void;
onFinish: () => void;
onFinish: (wasSuccessful: boolean) => void | Promise<void>;
},
): boolean {
): Promise<boolean> {
if (!line.startsWith("data: ")) {
return false;
}

const jsonStr = line.slice(6).trim();

if (jsonStr === "[DONE]") {
callbacks.onFinish();
await callbacks.onFinish(true);
return true;
}

Expand Down Expand Up @@ -299,7 +339,7 @@ async function parseStream(
onWebCitations: (webCitationSources: WebCitationSource[]) => void;
onParlaCitations: (sources: ParlaCitationSource[]) => void;
onOpenDataCitations: (sources: OpenDataCitationSource[]) => void;
onFinish: () => void;
onFinish: (wasSuccessful: boolean) => void | Promise<void>;
},
) {
const reader = body.getReader();
Expand All @@ -318,7 +358,7 @@ async function parseStream(
buffer = lines.pop() || "";

for (const line of lines) {
const isFinished = processStreamLine(line, callbacks);
const isFinished = await processStreamLine(line, callbacks);
if (isFinished) {
finishCalled = true;
}
Expand All @@ -333,6 +373,6 @@ async function parseStream(
"stream was done before reaching the the last streaming line ([DONE])",
),
);
callbacks.onFinish();
await callbacks.onFinish(false);
}
}
5 changes: 4 additions & 1 deletion apps/frontend/src/api/message/get-messages.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,5 +26,8 @@ export async function getMessages(chatId: number, signal: AbortSignal) {
* as `Jsonb | null` in the DB, which does not exist in Typescript.
* It actually is `number[] | null`, so we cast it here.
*/
return data as ChatMessage[];
return (data as ChatMessage[]).map((message) => ({
...message,
clientKey: message.id,
}));
}
3 changes: 3 additions & 0 deletions apps/frontend/src/api/message/insert-message.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,9 @@ export async function insertMessage(
allowed_document_ids: chatMessage.allowed_document_ids,
allowed_folder_ids: chatMessage.allowed_folder_ids,
citations: chatMessage.citations,
web_citations: chatMessage.web_citations,
parla_citations: chatMessage.parla_citations,
open_data_citations: chatMessage.open_data_citations,
external_tool_context: chatMessage.external_tool_context,
})
.select("*")
Expand Down
2 changes: 2 additions & 0 deletions apps/frontend/src/common.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,8 @@ export type ChatMessage = {
created_at: string;
external_tool_context: boolean;
id: number;
// Stable identity to avoid remounting the message's DOM node.
clientKey: number;
role: string;
type: string;
};
Expand Down
2 changes: 1 addition & 1 deletion apps/frontend/src/components/chat/chat-messages.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ export const ChatMessages: React.FC = () => {
className="flex w-full flex-col gap-y-1 lg:gap-y-3.5"
>
{messages.map((message) => (
<ChatMessage key={message.id} message={message} />
<ChatMessage key={message.clientKey} message={message} />
))}
{isWaitingForResponse && (
<div className="text-dunkelblau-50 flex gap-2 w-full items-center">
Expand Down
21 changes: 17 additions & 4 deletions apps/frontend/src/components/chat/hooks/use-chat-scrolling.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@ export function useChatScrolling(

const previousChatId = useRef(currentChatId);
const previousUserMessageCount = useRef(userMessageCount);
const latestUserMessageCount = useRef(userMessageCount);
latestUserMessageCount.current = userMessageCount;

//Jump to the bottom when a transient info message (tool deactivated / history scoped) appears.
useLayoutEffect(() => {
Expand All @@ -44,8 +46,11 @@ export function useChatScrolling(
}, [visibleInfoMessage, containerRef, scrollToBottom]);

/**
* Jump to the last message when a chat is opened or switched.
* Skipped when a search result asked us to scroll to a specific message.
* Jump to the last message when an existing chat is opened or switched.
* Skipped when a search result asked us to scroll to a specific message,
* and skipped when the chat id changed because it was *just created* by
* sending its first message — that case is a new message like any
* other and is handled by the "pin to top" effect below instead.
* The small timeout lets the newly selected chat's messages render first.
* Reads pendingScrollToMessage from getState so clearing it after a
* search-scroll does not re-trigger this effect.
Expand All @@ -54,6 +59,9 @@ export function useChatScrolling(
if (useChatScrollingStore.getState().pendingScrollToMessage !== null) {
return () => {};
}
if (latestUserMessageCount.current === 1) {
return () => {};
}
const timer = setTimeout(() => scrollToBottom("auto"), 1);
return () => clearTimeout(timer);
}, [currentChatId, scrollToBottom]);
Expand All @@ -71,7 +79,10 @@ export function useChatScrolling(

/**
* Scroll a newly sent user message to the top of the viewport.
* Skipped on a chat switch, where the effect above already jumps to the bottom.
* Skipped on a switch to an *existing* chat, where the effect above
* already jumps to the bottom — except when the chat id changed because
* it was just created by this very message (its first), which should
* still pin to top like any other new message.
*/
useLayoutEffect(() => {
const hasChatIdChanged = previousChatId.current !== currentChatId;
Expand All @@ -80,7 +91,9 @@ export function useChatScrolling(
previousChatId.current = currentChatId;
previousUserMessageCount.current = userMessageCount;

if (!hasChatIdChanged && hasNewUserMessage) {
const isFirstMessageInNewChat = hasChatIdChanged && userMessageCount === 1;

if ((!hasChatIdChanged || isFirstMessageInNewChat) && hasNewUserMessage) {
scrollNewMessageToTop();
}
}, [currentChatId, userMessageCount, scrollNewMessageToTop]);
Expand Down
Loading