Refactor AI chat components for improved context handling and error management

- Introduced useAgentChatRequestBody hook to streamline request body construction for chat messages.
- Updated AgentChatProvider to utilize new chat configuration logic, enhancing loading state management.
- Enhanced AIChatErrorMessage to support context-specific error handling with AIChatErrorMessageWithRecordsContext.
- Refactored SendMessageButton and SendMessageWithRecordsContextButton to accept agentId, improving component reusability.
- Simplified useAgentChat hook by removing unnecessary parameters and integrating chat context directly.
This commit is contained in:
Abdul Rahman
2025-10-13 23:13:09 +05:30
parent 1d0b193d74
commit afc518a056
10 changed files with 204 additions and 131 deletions
@@ -1,4 +1,7 @@
import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
import { useAgentChatRequestBody } from '@/ai/hooks/useAgentChatRequestBody';
import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
import { useChat } from '@ai-sdk/react';
import { useTheme } from '@emotion/react';
import styled from '@emotion/styled';
import { t } from '@lingui/core/macro';
@@ -43,14 +46,24 @@ const StyledErrorMessage = styled.div`
type AIChatErrorMessageProps = {
error: Error;
isRetrying?: boolean;
records?: ObjectRecord[];
};
export const AIChatErrorMessage = ({
error,
isRetrying = false,
records,
}: AIChatErrorMessageProps) => {
const theme = useTheme();
const { handleRetry } = useAgentChatContextOrThrow();
const { chat } = useAgentChatContextOrThrow();
const { buildRequestBody } = useAgentChatRequestBody();
const { regenerate } = useChat({ chat });
const handleRetry = () => {
regenerate({
body: buildRequestBody(records),
});
};
return (
<StyledErrorContainer>
@@ -8,8 +8,11 @@ import { AgentChatMessageRole } from '@/ai/constants/AgentChatMessageRole';
import { AIChatAssistantMessageRenderer } from '@/ai/components/AIChatAssistantMessageRenderer';
import { AIChatErrorMessage } from '@/ai/components/AIChatErrorMessage';
import { AIChatErrorMessageWithRecordsContext } from '@/ai/components/internal/AIChatErrorMessageWithRecordsContext';
import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
import { LightCopyIconButton } from '@/object-record/record-field/ui/components/LightCopyIconButton';
import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
import { isDefined } from 'twenty-shared/utils';
import { dateLocaleState } from '~/localization/states/dateLocaleState';
import { beautifyPastDateRelativeToNow } from '~/utils/date-utils';
@@ -152,6 +155,10 @@ export const AIChatMessage = ({
const theme = useTheme();
const { localeCatalog } = useRecoilValue(dateLocaleState);
const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
contextStoreCurrentObjectMetadataItemIdComponentState,
);
const showError =
isDefined(error) && message.role === AgentChatMessageRole.ASSISTANT;
@@ -195,9 +202,15 @@ export const AIChatMessage = ({
))}
</StyledFilesContainer>
)}
{showError && (
<AIChatErrorMessage error={error} isRetrying={isRetrying} />
)}
{showError &&
(contextStoreCurrentObjectMetadataItemId ? (
<AIChatErrorMessageWithRecordsContext
error={error}
isRetrying={isRetrying}
/>
) : (
<AIChatErrorMessage error={error} isRetrying={isRetrying} />
))}
{message.parts.length > 0 && message.metadata?.createdAt && (
<StyledMessageFooter className="message-footer">
<span>
@@ -17,7 +17,7 @@ import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
import { SendMessageWithRecordsContextButton } from '@/ai/components/internal/SendMessageWithRecordsContextButton';
import { AI_CHAT_INPUT_ID } from '@/ai/constants/AiChatInputId';
import { useAIChatFileUpload } from '@/ai/hooks/useAIChatFileUpload';
import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
import { useAgentChat } from '@/ai/hooks/useAgentChat';
import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
import { t } from '@lingui/core/macro';
@@ -70,7 +70,7 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
messages,
isStreaming,
error,
} = useAgentChatContextOrThrow();
} = useAgentChat(agentId);
const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
contextStoreCurrentObjectMetadataItemIdComponentState,
@@ -146,9 +146,9 @@ export const AIChatTab = ({ agentId }: { agentId: string }) => {
/>
<AgentChatFileUploadButton />
{contextStoreCurrentObjectMetadataItemId ? (
<SendMessageWithRecordsContextButton />
<SendMessageWithRecordsContextButton agentId={agentId} />
) : (
<SendMessageButton />
<SendMessageButton agentId={agentId} />
)}
</StyledButtonsContainer>
</StyledInputArea>
@@ -1,12 +1,27 @@
import { AgentChatContext } from '@/ai/contexts/AgentChatContext';
import { useAgentChat } from '@/ai/hooks/useAgentChat';
import { useAgentChatData } from '@/ai/hooks/useAgentChatData';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { currentWorkspaceState } from '@/auth/states/currentWorkspaceState';
import { useIsFeatureEnabled } from '@/workspace/hooks/useIsFeatureEnabled';
import { Chat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { Suspense } from 'react';
import { useRecoilValue } from 'recoil';
import { FeatureFlagKey } from '~/generated/graphql';
const createLoadingChat = () =>
new Chat<UIMessageWithMetadata>({
transport: new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/stream`,
headers: () => ({}),
}),
messages: [],
id: 'loading',
});
const AgentChatProviderContent = ({
agentId,
children,
@@ -15,15 +30,24 @@ const AgentChatProviderContent = ({
children: React.ReactNode;
}) => {
const { uiMessages, isLoading } = useAgentChatData(agentId);
const chatState = useAgentChat(agentId, uiMessages);
const combinedIsLoading = chatState.isLoading || isLoading;
const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
const chatConfig = isLoading
? createLoadingChat()
: new Chat<UIMessageWithMetadata>({
transport: new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/stream`,
headers: () => ({
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
}),
}),
messages: uiMessages,
id: `${currentAIChatThread}-${uiMessages.length}`,
});
return (
<AgentChatContext.Provider
value={{
...chatState,
isLoading: combinedIsLoading,
}}
value={{ chat: chatConfig, isLoadingData: isLoading }}
>
{children}
</AgentChatContext.Provider>
@@ -41,14 +65,30 @@ export const AgentChatProvider = ({
if (!isAiEnabled || !agentId) {
return (
<AgentChatContext.Provider value={null}>
<AgentChatContext.Provider
value={{
chat: createLoadingChat(),
isLoadingData: false,
}}
>
{children}
</AgentChatContext.Provider>
);
}
return (
<Suspense fallback={null}>
<Suspense
fallback={
<AgentChatContext.Provider
value={{
chat: createLoadingChat(),
isLoadingData: true,
}}
>
{children}
</AgentChatContext.Provider>
}
>
<AgentChatProviderContent agentId={agentId}>
{children}
</AgentChatProviderContent>
@@ -0,0 +1,22 @@
import { AIChatErrorMessage } from '@/ai/components/AIChatErrorMessage';
import { useFindManyRecordsSelectedInContextStore } from '@/context-store/hooks/useFindManyRecordsSelectedInContextStore';
export const AIChatErrorMessageWithRecordsContext = ({
error,
isRetrying,
}: {
error: Error;
isRetrying?: boolean;
}) => {
const { records } = useFindManyRecordsSelectedInContextStore({
limit: 10,
});
return (
<AIChatErrorMessage
error={error}
isRetrying={isRetrying}
records={records}
/>
);
};
@@ -1,24 +1,57 @@
import { AI_CHAT_INPUT_ID } from '@/ai/constants/AiChatInputId';
import { useAgentChat } from '@/ai/hooks/useAgentChat';
import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
import { useAgentChatRequestBody } from '@/ai/hooks/useAgentChatRequestBody';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
import { useHotkeysOnFocusedElement } from '@/ui/utilities/hotkey/hooks/useHotkeysOnFocusedElement';
import { useChat } from '@ai-sdk/react';
import { t } from '@lingui/core/macro';
import { useRecoilState } from 'recoil';
import { Key } from 'ts-key-enum';
import { Button } from 'twenty-ui/input';
export const SendMessageButton = ({
agentId,
records,
}: {
agentId: string;
records?: ObjectRecord[];
}) => {
const { handleSendMessage, isLoading, input } = useAgentChatContextOrThrow();
const { input, isLoading, handleInputChange } = useAgentChat(agentId);
const { chat } = useAgentChatContextOrThrow();
const { buildRequestBody } = useAgentChatRequestBody();
const { sendMessage } = useChat({ chat });
const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useRecoilState(
agentChatUploadedFilesState,
);
const handleSendMessage = () => {
if (input.trim() === '' || isLoading) {
return;
}
sendMessage(
{
text: input,
files: agentChatUploadedFiles,
},
{
body: buildRequestBody(records),
},
);
handleInputChange('');
setAgentChatUploadedFiles([]);
};
useHotkeysOnFocusedElement({
keys: [Key.Enter],
callback: (event: KeyboardEvent) => {
if (!event.ctrlKey && !event.metaKey) {
event.preventDefault();
handleSendMessage(records);
handleSendMessage();
}
},
focusId: AI_CHAT_INPUT_ID,
@@ -31,7 +64,7 @@ export const SendMessageButton = ({
return (
<Button
hotkeys={input && !isLoading ? ['⏎'] : undefined}
onClick={() => handleSendMessage(records)}
onClick={handleSendMessage}
disabled={!input || isLoading}
variant="primary"
accent="blue"
@@ -1,10 +1,14 @@
import { SendMessageButton } from '@/ai/components/internal/SendMessageButton';
import { useFindManyRecordsSelectedInContextStore } from '@/context-store/hooks/useFindManyRecordsSelectedInContextStore';
export const SendMessageWithRecordsContextButton = () => {
export const SendMessageWithRecordsContextButton = ({
agentId,
}: {
agentId: string;
}) => {
const { records } = useFindManyRecordsSelectedInContextStore({
limit: 10,
});
return <SendMessageButton records={records} />;
return <SendMessageButton agentId={agentId} records={records} />;
};
@@ -1,25 +1,10 @@
import { type AIChatObjectMetadataAndRecordContext } from '@/ai/states/agentChatContextState';
import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
import { type Chat } from '@ai-sdk/react';
import { createContext } from 'react';
import { type ObjectRecord } from '../../object-record/types/ObjectRecord';
export type AgentChatContextValue = {
messages: UIMessageWithMetadata[];
isStreaming: boolean;
isLoading: boolean;
error?: Error;
input: string;
handleInputChange: (value: string) => void;
handleSendMessage: (records?: ObjectRecord[]) => Promise<void>;
handleSetContext: (
items: Array<AIChatObjectMetadataAndRecordContext>,
) => Promise<void>;
scrollWrapperId: string;
context: Array<AIChatObjectMetadataAndRecordContext>;
handleRetry: () => void;
chat: Chat<UIMessageWithMetadata>;
isLoadingData: boolean;
};
export const AgentChatContext = createContext<AgentChatContextValue | null>(
@@ -1,39 +1,17 @@
import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
import { useRecoilState, useRecoilValue } from 'recoil';
import { useAgentChatContextOrThrow } from '@/ai/hooks/useAgentChatContextOrThrow';
import {
type AIChatObjectMetadataAndRecordContext,
agentChatContextState,
} from '@/ai/states/agentChatContextState';
import { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { isAgentChatCurrentContextActiveState } from '@/ai/states/isAgentChatCurrentContextActiveState';
import { type UIMessageWithMetadata } from '@/ai/types/UIMessageWithMetadata';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
import { useGetObjectMetadataItemById } from '@/object-metadata/hooks/useGetObjectMetadataItemById';
import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
import { useScrollWrapperHTMLElement } from '@/ui/utilities/scroll/hooks/useScrollWrapperHTMLElement';
import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { isDefined } from 'twenty-shared/utils';
import { REST_API_BASE_URL } from '../../apollo/constant/rest-api-base-url';
import { agentChatInputState } from '../states/agentChatInputState';
export const useAgentChat = (
agentId: string,
uiMessages: UIMessageWithMetadata[],
) => {
const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
contextStoreCurrentObjectMetadataItemIdComponentState,
);
const isAgentChatCurrentContextActive = useRecoilValue(
isAgentChatCurrentContextActiveState,
);
export const useAgentChat = (agentId: string) => {
const { chat, isLoadingData } = useAgentChatContextOrThrow();
const agentChatSelectedFiles = useRecoilValue(agentChatSelectedFilesState);
@@ -43,80 +21,22 @@ export const useAgentChat = (
const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useRecoilState(
agentChatUploadedFilesState,
);
const [agentChatInput, setAgentChatInput] =
useRecoilState(agentChatInputState);
const scrollWrapperId = `scroll-wrapper-ai-chat-${agentId}`;
const { scrollWrapperHTMLElement } =
useScrollWrapperHTMLElement(scrollWrapperId);
const { sendMessage, messages, status, error, regenerate } = useChat({
transport: new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/stream`,
headers: () => ({
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
}),
}),
messages: uiMessages,
id: `${currentAIChatThread}-${uiMessages.length}`,
const { messages, status, error } = useChat({
chat,
});
const isStreaming = status === 'streaming';
const scrollToBottom = () => {
scrollWrapperHTMLElement?.scroll({
top: scrollWrapperHTMLElement.scrollHeight,
behavior: 'smooth',
});
};
const isLoading =
!currentAIChatThread || isStreaming || agentChatSelectedFiles.length > 0;
const handleSendMessage = async (records?: ObjectRecord[]) => {
if (agentChatInput.trim() === '' || isLoading === true) {
return;
}
const content = agentChatInput.trim();
setAgentChatInput('');
const recordIdsByObjectMetadataNameSingular = [];
if (
isAgentChatCurrentContextActive === true &&
isDefined(records) &&
isDefined(contextStoreCurrentObjectMetadataItemId)
) {
recordIdsByObjectMetadataNameSingular.push({
objectMetadataNameSingular: getObjectMetadataItemById(
contextStoreCurrentObjectMetadataItemId,
).nameSingular,
recordIds: records.map(({ id }) => id),
});
}
sendMessage(
{
text: content,
files: agentChatUploadedFiles,
},
{
body: {
threadId: currentAIChatThread,
recordIdsByObjectMetadataNameSingular,
},
},
);
setAgentChatUploadedFiles([]);
setTimeout(scrollToBottom, 100);
};
isLoadingData ||
!currentAIChatThread ||
isStreaming ||
agentChatSelectedFiles.length > 0;
const handleSetContext = async (
items: Array<AIChatObjectMetadataAndRecordContext>,
@@ -130,11 +50,9 @@ export const useAgentChat = (
input: agentChatInput,
context: agentChatContext,
handleSetContext,
handleSendMessage,
isLoading,
scrollWrapperId,
isStreaming,
error,
handleRetry: regenerate,
};
};
@@ -0,0 +1,45 @@
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { isAgentChatCurrentContextActiveState } from '@/ai/states/isAgentChatCurrentContextActiveState';
import { contextStoreCurrentObjectMetadataItemIdComponentState } from '@/context-store/states/contextStoreCurrentObjectMetadataItemIdComponentState';
import { useGetObjectMetadataItemById } from '@/object-metadata/hooks/useGetObjectMetadataItemById';
import { type ObjectRecord } from '@/object-record/types/ObjectRecord';
import { useRecoilComponentValue } from '@/ui/utilities/state/component-state/hooks/useRecoilComponentValue';
import { useRecoilValue } from 'recoil';
import { isDefined } from 'twenty-shared/utils';
export const useAgentChatRequestBody = () => {
const currentAIChatThread = useRecoilValue(currentAIChatThreadState);
const { getObjectMetadataItemById } = useGetObjectMetadataItemById();
const contextStoreCurrentObjectMetadataItemId = useRecoilComponentValue(
contextStoreCurrentObjectMetadataItemIdComponentState,
);
const isAgentChatCurrentContextActive = useRecoilValue(
isAgentChatCurrentContextActiveState,
);
const buildRequestBody = (records?: ObjectRecord[]) => {
const recordIdsByObjectMetadataNameSingular = [];
if (
isAgentChatCurrentContextActive === true &&
isDefined(records) &&
isDefined(contextStoreCurrentObjectMetadataItemId)
) {
recordIdsByObjectMetadataNameSingular.push({
objectMetadataNameSingular: getObjectMetadataItemById(
contextStoreCurrentObjectMetadataItemId,
).nameSingular,
recordIds: records.map(({ id }) => id),
});
}
return {
threadId: currentAIChatThread,
recordIdsByObjectMetadataNameSingular,
};
};
return { buildRequestBody };
};