context_management: handle summarization in UI (#2377)
This commit is contained in:
Lily Delalande
@@ -0,0 +1,77 @@
|
||||
import React, { useState, useRef, useEffect } from 'react';
|
||||
import { Message } from '../../types/message';
|
||||
import { useChatContextManager } from './ContextManager';
|
||||
|
||||
interface ContextLengthExceededHandlerProps {
|
||||
messages: Message[];
|
||||
messageId: string;
|
||||
}
|
||||
|
||||
export const ContextLengthExceededHandler: React.FC<ContextLengthExceededHandlerProps> = ({
|
||||
messages,
|
||||
messageId,
|
||||
}) => {
|
||||
const { fetchSummary, summaryContent, isLoadingSummary, errorLoadingSummary, openSummaryModal } =
|
||||
useChatContextManager();
|
||||
|
||||
const [hasFetchStarted, setHasFetchStarted] = useState(false);
|
||||
|
||||
// Find the relevant message to check if it's the latest
|
||||
const isCurrentMessageLatest =
|
||||
messageId === messages[messages.length - 1].id ||
|
||||
messageId === messages[messages.length - 1].created.toString();
|
||||
|
||||
// Only allow interaction for the most recent context length exceeded event
|
||||
const shouldAllowSummaryInteraction = isCurrentMessageLatest;
|
||||
|
||||
// Use a ref to track if we've started the fetch
|
||||
const fetchStartedRef = useRef(false);
|
||||
|
||||
useEffect(() => {
|
||||
// Automatically fetch summary if conditions are met
|
||||
if (
|
||||
!summaryContent &&
|
||||
!hasFetchStarted &&
|
||||
shouldAllowSummaryInteraction &&
|
||||
!fetchStartedRef.current
|
||||
) {
|
||||
setHasFetchStarted(true);
|
||||
fetchStartedRef.current = true;
|
||||
fetchSummary(messages);
|
||||
}
|
||||
}, [fetchSummary, hasFetchStarted, messages, shouldAllowSummaryInteraction, summaryContent]);
|
||||
|
||||
// Handle retry
|
||||
const handleRetry = () => {
|
||||
if (!shouldAllowSummaryInteraction) return;
|
||||
fetchSummary(messages);
|
||||
};
|
||||
|
||||
// Render the notification UI
|
||||
return (
|
||||
<div className="flex flex-col items-start mt-1 pl-4">
|
||||
{isLoadingSummary && shouldAllowSummaryInteraction ? (
|
||||
// Only show loading indicator during loading state
|
||||
<div className="flex items-center text-xs text-gray-400">
|
||||
<span className="mr-2">Preparing summary...</span>
|
||||
<span className="animate-spin h-3 w-3 border-2 border-gray-400 rounded-full border-t-transparent"></span>
|
||||
</div>
|
||||
) : (
|
||||
// Show different UI based on whether it's already handled
|
||||
<>
|
||||
<span className="text-xs text-gray-400 italic">{'Session summarized'}</span>
|
||||
|
||||
{/* Only show the button if its last message */}
|
||||
{shouldAllowSummaryInteraction && (
|
||||
<button
|
||||
onClick={() => (errorLoadingSummary ? handleRetry() : openSummaryModal())}
|
||||
className="text-xs text-textStandard hover:text-textSubtle transition-colors mt-1 flex items-center"
|
||||
>
|
||||
{errorLoadingSummary ? 'Retry loading summary' : 'View or edit summary'}
|
||||
</button>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
@@ -0,0 +1,170 @@
|
||||
import React, { createContext, useContext, useState } from 'react';
|
||||
import { Message } from '../../types/message';
|
||||
import { manageContextFromBackend, convertApiMessageToFrontendMessage } from './index';
|
||||
|
||||
// Define the context management interface
|
||||
interface ContextManagerState {
|
||||
summaryContent: string;
|
||||
summarizedThread: Message[];
|
||||
isSummaryModalOpen: boolean;
|
||||
isLoadingSummary: boolean;
|
||||
errorLoadingSummary: boolean;
|
||||
}
|
||||
|
||||
interface ContextManagerActions {
|
||||
fetchSummary: (messages: Message[]) => Promise<void>;
|
||||
updateSummary: (newSummaryContent: string) => void;
|
||||
resetMessagesWithSummary: (
|
||||
messages: Message[],
|
||||
setMessages: (messages: Message[]) => void
|
||||
) => void;
|
||||
openSummaryModal: () => void;
|
||||
closeSummaryModal: () => void;
|
||||
hasContextLengthExceededContent: (message: Message) => boolean;
|
||||
}
|
||||
|
||||
// Create the context
|
||||
const ContextManagerContext = createContext<
|
||||
(ContextManagerState & ContextManagerActions) | undefined
|
||||
>(undefined);
|
||||
|
||||
// Create the provider component
|
||||
export const ContextManagerProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
|
||||
const [summaryContent, setSummaryContent] = useState<string>('');
|
||||
const [summarizedThread, setSummarizedThread] = useState<Message[]>([]);
|
||||
const [isSummaryModalOpen, setIsSummaryModalOpen] = useState<boolean>(false);
|
||||
const [isLoadingSummary, setIsLoadingSummary] = useState<boolean>(false);
|
||||
const [errorLoadingSummary, setErrorLoadingSummary] = useState<boolean>(false);
|
||||
|
||||
const fetchSummary = async (messages: Message[]) => {
|
||||
setIsLoadingSummary(true);
|
||||
setErrorLoadingSummary(false);
|
||||
|
||||
try {
|
||||
const response = await manageContextFromBackend({
|
||||
messages: messages,
|
||||
manageAction: 'summarize',
|
||||
});
|
||||
|
||||
// Convert API messages to frontend messages
|
||||
const convertedMessages = response.messages.map((apiMessage) =>
|
||||
convertApiMessageToFrontendMessage(apiMessage)
|
||||
);
|
||||
|
||||
// Extract the summary text from the first message
|
||||
const summaryMessage = convertedMessages[0].content[0];
|
||||
if (summaryMessage.type === 'text') {
|
||||
const summary = summaryMessage.text;
|
||||
setSummaryContent(summary);
|
||||
setSummarizedThread(convertedMessages);
|
||||
}
|
||||
|
||||
setIsLoadingSummary(false);
|
||||
} catch (err) {
|
||||
console.error('Error fetching summary:', err);
|
||||
setErrorLoadingSummary(true);
|
||||
setIsLoadingSummary(false);
|
||||
}
|
||||
};
|
||||
|
||||
const updateSummary = (newSummaryContent: string) => {
|
||||
// Update the summary content
|
||||
setSummaryContent(newSummaryContent);
|
||||
|
||||
// Update the thread if it exists
|
||||
if (summarizedThread.length > 0) {
|
||||
// Create a deep copy of the thread
|
||||
const updatedThread = [...summarizedThread];
|
||||
|
||||
// Create a copy of the first message
|
||||
const firstMessage = { ...updatedThread[0] };
|
||||
|
||||
// Create a copy of the content array
|
||||
const updatedContent = [...firstMessage.content];
|
||||
|
||||
// Update the summary text in the first content item
|
||||
if (updatedContent[0] && updatedContent[0].type === 'text') {
|
||||
updatedContent[0] = {
|
||||
...updatedContent[0],
|
||||
text: newSummaryContent,
|
||||
};
|
||||
}
|
||||
|
||||
// Update the message with the new content
|
||||
firstMessage.content = updatedContent;
|
||||
updatedThread[0] = firstMessage;
|
||||
|
||||
// Update the thread
|
||||
setSummarizedThread(updatedThread);
|
||||
}
|
||||
};
|
||||
|
||||
const resetMessagesWithSummary = (
|
||||
messages: Message[],
|
||||
setMessages: (messages: Message[]) => void
|
||||
) => {
|
||||
// Update summarizedThread with metadata
|
||||
const updatedSummarizedThread = summarizedThread.map((msg) => ({
|
||||
...msg,
|
||||
display: false,
|
||||
sendToLLM: true,
|
||||
}));
|
||||
|
||||
// Update list of messages with other metadata
|
||||
const updatedMessages = messages.map((msg) => ({
|
||||
...msg,
|
||||
display: true,
|
||||
sendToLLM: false,
|
||||
}));
|
||||
|
||||
// Make a copy that combines both
|
||||
const newMessages = [...updatedMessages, ...updatedSummarizedThread];
|
||||
|
||||
// Update the messages state
|
||||
setMessages(newMessages);
|
||||
|
||||
// Clear the summarized thread and content
|
||||
setSummarizedThread([]);
|
||||
setSummaryContent('');
|
||||
};
|
||||
|
||||
const hasContextLengthExceededContent = (message: Message): boolean => {
|
||||
return message.content.some((content) => content.type === 'contextLengthExceeded');
|
||||
};
|
||||
|
||||
const openSummaryModal = () => {
|
||||
setIsSummaryModalOpen(true);
|
||||
};
|
||||
|
||||
const closeSummaryModal = () => {
|
||||
setIsSummaryModalOpen(false);
|
||||
};
|
||||
|
||||
const value = {
|
||||
// State
|
||||
summaryContent,
|
||||
summarizedThread,
|
||||
isSummaryModalOpen,
|
||||
isLoadingSummary,
|
||||
errorLoadingSummary,
|
||||
|
||||
// Actions
|
||||
fetchSummary,
|
||||
updateSummary,
|
||||
resetMessagesWithSummary,
|
||||
openSummaryModal,
|
||||
closeSummaryModal,
|
||||
hasContextLengthExceededContent,
|
||||
};
|
||||
|
||||
return <ContextManagerContext.Provider value={value}>{children}</ContextManagerContext.Provider>;
|
||||
};
|
||||
|
||||
// Create a hook to use the context
|
||||
export const useChatContextManager = () => {
|
||||
const context = useContext(ContextManagerContext);
|
||||
if (context === undefined) {
|
||||
throw new Error('useContextManager must be used within a ContextManagerProvider');
|
||||
}
|
||||
return context;
|
||||
};
|
||||
@@ -1,37 +1,119 @@
|
||||
import { Message } from '../../types/message';
|
||||
import { getApiUrl, getSecretKey } from '../../config';
|
||||
import {
|
||||
Message as FrontendMessage,
|
||||
Content as FrontendContent,
|
||||
MessageContent as FrontendMessageContent,
|
||||
ToolCallResult,
|
||||
ToolCall,
|
||||
} from '../../types/message';
|
||||
import {
|
||||
ContextManageRequest,
|
||||
ContextManageResponse,
|
||||
manageContext,
|
||||
Message as ApiMessage,
|
||||
MessageContent as ApiMessageContent,
|
||||
} from '../../api';
|
||||
import { generateId } from 'ai';
|
||||
|
||||
export async function manageContext({
|
||||
export async function manageContextFromBackend({
|
||||
messages,
|
||||
manageAction,
|
||||
}: {
|
||||
messages: Message[];
|
||||
manageAction: 'trunction' | 'summarize';
|
||||
}) {
|
||||
const response = await fetch(getApiUrl('/context/manage'), {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'X-Secret-Key': getSecretKey(),
|
||||
},
|
||||
body: JSON.stringify({
|
||||
messages,
|
||||
manageAction,
|
||||
}),
|
||||
});
|
||||
if (!response.ok) {
|
||||
if (!response.ok) {
|
||||
// Get the status text or a default message
|
||||
const errorText = await response.text().catch(() => 'Unknown error');
|
||||
// log error with status and details
|
||||
console.error(
|
||||
`Context management failed: ${response.status} ${response.statusText} - ${errorText}`
|
||||
);
|
||||
throw new Error(
|
||||
`Context management failed: ${response.status} ${response.statusText} - ${errorText}\n\nStart a new session.`
|
||||
);
|
||||
messages: FrontendMessage[];
|
||||
manageAction: 'truncation' | 'summarize';
|
||||
}): Promise<ContextManageResponse> {
|
||||
try {
|
||||
const contextManagementRequest = { manageAction, messages };
|
||||
|
||||
// Cast to the API-expected type
|
||||
const result = await manageContext({
|
||||
body: contextManagementRequest as unknown as ContextManageRequest,
|
||||
});
|
||||
|
||||
// Check for errors in the result
|
||||
if (result.error) {
|
||||
throw new Error(`Context management failed: ${result.error}`);
|
||||
}
|
||||
|
||||
// Extract the actual data from the result
|
||||
if (!result.data) {
|
||||
throw new Error('Context management returned no data');
|
||||
}
|
||||
|
||||
return result.data;
|
||||
} catch (error) {
|
||||
console.error(`Context management failed: ${error || 'Unknown error'}`);
|
||||
throw new Error(
|
||||
`Context management failed: ${error || 'Unknown error'}\n\nStart a new session.`
|
||||
);
|
||||
}
|
||||
const data = await response.json();
|
||||
return data;
|
||||
}
|
||||
|
||||
// Function to convert API Message to frontend Message
|
||||
export function convertApiMessageToFrontendMessage(apiMessage: ApiMessage): FrontendMessage {
|
||||
return {
|
||||
display: false,
|
||||
sendToLLM: false,
|
||||
id: generateId(),
|
||||
role: apiMessage.role,
|
||||
created: apiMessage.created,
|
||||
content: apiMessage.content
|
||||
.map((apiContent) => mapApiContentToFrontendMessageContent(apiContent))
|
||||
.filter((content): content is FrontendMessageContent => content !== null),
|
||||
};
|
||||
}
|
||||
|
||||
// Function to convert API MessageContent to frontend MessageContent
|
||||
function mapApiContentToFrontendMessageContent(
|
||||
apiContent: ApiMessageContent
|
||||
): FrontendMessageContent | null {
|
||||
// Handle each content type specifically based on its "type" property
|
||||
if (apiContent.type === 'text') {
|
||||
return {
|
||||
type: 'text',
|
||||
text: apiContent.text,
|
||||
annotations: apiContent.annotations as Record<string, unknown> | undefined,
|
||||
};
|
||||
} else if (apiContent.type === 'image') {
|
||||
return {
|
||||
type: 'image',
|
||||
data: apiContent.data,
|
||||
mimeType: apiContent.mimeType,
|
||||
annotations: apiContent.annotations as Record<string, unknown> | undefined,
|
||||
};
|
||||
} else if (apiContent.type === 'toolRequest') {
|
||||
// Ensure the toolCall has the correct type structure
|
||||
const toolCall = apiContent.toolCall as unknown as ToolCallResult<ToolCall>;
|
||||
|
||||
return {
|
||||
type: 'toolRequest',
|
||||
id: apiContent.id,
|
||||
toolCall: toolCall,
|
||||
};
|
||||
} else if (apiContent.type === 'toolResponse') {
|
||||
// Ensure the toolResult has the correct type structure
|
||||
const toolResult = apiContent.toolResult as unknown as ToolCallResult<FrontendContent[]>;
|
||||
|
||||
return {
|
||||
type: 'toolResponse',
|
||||
id: apiContent.id,
|
||||
toolResult: toolResult,
|
||||
};
|
||||
} else if (apiContent.type === 'toolConfirmationRequest') {
|
||||
return {
|
||||
type: 'toolConfirmationRequest',
|
||||
id: apiContent.id,
|
||||
toolName: apiContent.toolName,
|
||||
arguments: apiContent.arguments as Record<string, unknown>,
|
||||
prompt: apiContent.prompt === null ? undefined : apiContent.prompt,
|
||||
};
|
||||
} else if (apiContent.type === 'contextLengthExceeded') {
|
||||
return {
|
||||
type: 'contextLengthExceeded',
|
||||
msg: apiContent.msg,
|
||||
};
|
||||
}
|
||||
|
||||
// For types that exist in API but not in frontend, either skip or convert
|
||||
console.warn(`Skipping unsupported content type: ${apiContent.type}`);
|
||||
return null;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user