From 55a1d867c3da44460a6872b76416cd5cf3ebd373 Mon Sep 17 00:00:00 2001 From: ManishMadan2882 Date: Thu, 19 Jun 2025 02:16:03 +0530 Subject: [PATCH] (refactor/converstationSlice) separate preview --- frontend/src/agents/AgentPreview.tsx | 25 +- frontend/src/agents/SharedAgent.tsx | 4 +- frontend/src/agents/agentPreviewSlice.ts | 319 ++++++++++++++++ .../src/conversation/conversationSlice.ts | 345 +++++++++--------- frontend/src/store.ts | 2 + 5 files changed, 501 insertions(+), 194 deletions(-) create mode 100644 frontend/src/agents/agentPreviewSlice.ts diff --git a/frontend/src/agents/AgentPreview.tsx b/frontend/src/agents/AgentPreview.tsx index cd4fdc18..8798a155 100644 --- a/frontend/src/agents/AgentPreview.tsx +++ b/frontend/src/agents/AgentPreview.tsx @@ -6,24 +6,23 @@ import ConversationMessages from '../conversation/ConversationMessages'; import { Query } from '../conversation/conversationModels'; import { addQuery, - fetchAnswer, - handleAbort, + fetchPreviewAnswer, + handlePreviewAbort, resendQuery, - resetConversation, - selectQueries, - selectStatus, -} from '../conversation/conversationSlice'; + resetPreview, + selectPreviewQueries, + selectPreviewStatus, +} from './agentPreviewSlice'; import { selectSelectedAgent } from '../preferences/preferenceSlice'; import { AppDispatch } from '../store'; export default function AgentPreview() { const dispatch = useDispatch(); - const queries = useSelector(selectQueries); - const status = useSelector(selectStatus); + const queries = useSelector(selectPreviewQueries); + const status = useSelector(selectPreviewStatus); const selectedAgent = useSelector(selectSelectedAgent); - const [input, setInput] = useState(''); const [lastQueryReturnedErr, setLastQueryReturnedErr] = useState(false); const fetchStream = useRef(null); @@ -31,7 +30,7 @@ export default function AgentPreview() { const handleFetchAnswer = useCallback( ({ question, index }: { question: string; index?: number }) => { fetchStream.current = dispatch( - fetchAnswer({ question, indx: index, isPreview: true }), + fetchPreviewAnswer({ question, indx: index }), ); }, [dispatch], @@ -95,11 +94,11 @@ export default function AgentPreview() { }; useEffect(() => { - dispatch(resetConversation()); + dispatch(resetPreview()); return () => { if (fetchStream.current) fetchStream.current.abort(); - handleAbort(); - dispatch(resetConversation()); + handlePreviewAbort(); + dispatch(resetPreview()); }; }, [dispatch]); diff --git a/frontend/src/agents/SharedAgent.tsx b/frontend/src/agents/SharedAgent.tsx index d1d2d370..1982b697 100644 --- a/frontend/src/agents/SharedAgent.tsx +++ b/frontend/src/agents/SharedAgent.tsx @@ -57,9 +57,7 @@ export default function SharedAgent() { const handleFetchAnswer = useCallback( ({ question, index }: { question: string; index?: number }) => { - fetchStream.current = dispatch( - fetchAnswer({ question, indx: index, isPreview: false }), - ); + fetchStream.current = dispatch(fetchAnswer({ question, indx: index })); }, [dispatch], ); diff --git a/frontend/src/agents/agentPreviewSlice.ts b/frontend/src/agents/agentPreviewSlice.ts new file mode 100644 index 00000000..c6e628f6 --- /dev/null +++ b/frontend/src/agents/agentPreviewSlice.ts @@ -0,0 +1,319 @@ +import { createAsyncThunk, createSlice, PayloadAction } from '@reduxjs/toolkit'; +import { + Answer, + ConversationState, + Query, + Status, +} from '../conversation/conversationModels'; +import { + handleFetchAnswer, + handleFetchAnswerSteaming, +} from '../conversation/conversationHandlers'; +import { + selectCompletedAttachments, + clearAttachments, +} from '../upload/uploadSlice'; +import store from '../store'; + +const initialState: ConversationState = { + queries: [], + status: 'idle', + conversationId: null, +}; + +const API_STREAMING = import.meta.env.VITE_API_STREAMING === 'true'; + +let abortController: AbortController | null = null; +export function handlePreviewAbort() { + if (abortController) { + abortController.abort(); + abortController = null; + } +} + +export const fetchPreviewAnswer = createAsyncThunk< + Answer, + { question: string; indx?: number } +>( + 'agentPreview/fetchAnswer', + async ({ question, indx }, { dispatch, getState }) => { + if (abortController) abortController.abort(); + abortController = new AbortController(); + const { signal } = abortController; + + const state = getState() as RootState; + const attachmentIds = selectCompletedAttachments(state) + .filter((a) => a.id) + .map((a) => a.id) as string[]; + + if (attachmentIds.length > 0) { + dispatch(clearAttachments()); + } + + if (state.preference) { + if (API_STREAMING) { + await handleFetchAnswerSteaming( + question, + signal, + state.preference.token, + state.preference.selectedDocs!, + state.agentPreview.queries, + null, // No conversation ID for previews + state.preference.prompt.id, + state.preference.chunks, + state.preference.token_limit, + (event) => { + const data = JSON.parse(event.data); + const targetIndex = indx ?? state.agentPreview.queries.length - 1; + + if (data.type === 'end') { + dispatch(agentPreviewSlice.actions.setStatus('idle')); + } else if (data.type === 'thought') { + dispatch( + updateThought({ + index: targetIndex, + query: { thought: data.thought }, + }), + ); + } else if (data.type === 'source') { + dispatch( + updateStreamingSource({ + index: targetIndex, + query: { sources: data.source ?? [] }, + }), + ); + } else if (data.type === 'tool_call') { + dispatch( + updateToolCall({ + index: targetIndex, + tool_call: data.data, + }), + ); + } else if (data.type === 'error') { + dispatch(agentPreviewSlice.actions.setStatus('failed')); + dispatch( + agentPreviewSlice.actions.raiseError({ + index: targetIndex, + message: data.error, + }), + ); + } else { + dispatch( + updateStreamingQuery({ + index: targetIndex, + query: { response: data.answer }, + }), + ); + } + }, + indx, + state.preference.selectedAgent?.id, + attachmentIds, + false, // Don't save preview conversations + ); + } else { + // Non-streaming implementation + const answer = await handleFetchAnswer( + question, + signal, + state.preference.token, + state.preference.selectedDocs!, + state.agentPreview.queries, + null, // No conversation ID for previews + state.preference.prompt.id, + state.preference.chunks, + state.preference.token_limit, + state.preference.selectedAgent?.id, + attachmentIds, + false, // Don't save preview conversations + ); + + if (answer) { + const sourcesPrepped = answer.sources.map( + (source: { title: string }) => { + if (source && source.title) { + const titleParts = source.title.split('/'); + return { + ...source, + title: titleParts[titleParts.length - 1], + }; + } + return source; + }, + ); + + const targetIndex = indx ?? state.agentPreview.queries.length - 1; + + dispatch( + updateQuery({ + index: targetIndex, + query: { + response: answer.answer, + thought: answer.thought, + sources: sourcesPrepped, + tool_calls: answer.toolCalls, + }, + }), + ); + dispatch(agentPreviewSlice.actions.setStatus('idle')); + } + } + } + + return { + conversationId: null, + title: null, + answer: '', + query: question, + result: '', + thought: '', + sources: [], + tool_calls: [], + }; + }, +); + +export const agentPreviewSlice = createSlice({ + name: 'agentPreview', + initialState, + reducers: { + addQuery(state, action: PayloadAction) { + state.queries.push(action.payload); + }, + resendQuery( + state, + action: PayloadAction<{ index: number; prompt: string; query?: Query }>, + ) { + state.queries = [ + ...state.queries.splice(0, action.payload.index), + action.payload, + ]; + }, + updateStreamingQuery( + state, + action: PayloadAction<{ + index: number; + query: Partial; + }>, + ) { + const { index, query } = action.payload; + if (state.status === 'idle') return; + + if (query.response != undefined) { + state.queries[index].response = + (state.queries[index].response || '') + query.response; + } + }, + updateThought( + state, + action: PayloadAction<{ + index: number; + query: Partial; + }>, + ) { + const { index, query } = action.payload; + if (query.thought != undefined) { + state.queries[index].thought = + (state.queries[index].thought || '') + query.thought; + } + }, + updateStreamingSource( + state, + action: PayloadAction<{ + index: number; + query: Partial; + }>, + ) { + const { index, query } = action.payload; + if (!state.queries[index].sources) { + state.queries[index].sources = query?.sources; + } else if (query.sources) { + state.queries[index].sources!.push(...query.sources); + } + }, + updateToolCall(state, action) { + const { index, tool_call } = action.payload; + + if (!state.queries[index].tool_calls) { + state.queries[index].tool_calls = []; + } + + const existingIndex = state.queries[index].tool_calls.findIndex( + (call) => call.call_id === tool_call.call_id, + ); + + if (existingIndex !== -1) { + const existingCall = state.queries[index].tool_calls[existingIndex]; + state.queries[index].tool_calls[existingIndex] = { + ...existingCall, + ...tool_call, + }; + } else state.queries[index].tool_calls.push(tool_call); + }, + updateQuery( + state, + action: PayloadAction<{ index: number; query: Partial }>, + ) { + const { index, query } = action.payload; + state.queries[index] = { + ...state.queries[index], + ...query, + }; + }, + setStatus(state, action: PayloadAction) { + state.status = action.payload; + }, + raiseError( + state, + action: PayloadAction<{ + index: number; + message: string; + }>, + ) { + const { index, message } = action.payload; + state.queries[index].error = message; + }, + resetPreview: (state) => { + state.queries = initialState.queries; + state.status = initialState.status; + state.conversationId = initialState.conversationId; + handlePreviewAbort(); + }, + }, + extraReducers(builder) { + builder + .addCase(fetchPreviewAnswer.pending, (state) => { + state.status = 'loading'; + }) + .addCase(fetchPreviewAnswer.rejected, (state, action) => { + if (action.meta.aborted) { + state.status = 'idle'; + return state; + } + state.status = 'failed'; + state.queries[state.queries.length - 1].error = 'Something went wrong'; + }); + }, +}); + +type RootState = ReturnType; + +export const selectPreviewQueries = (state: RootState) => + state.agentPreview.queries; +export const selectPreviewStatus = (state: RootState) => + state.agentPreview.status; + +export const { + addQuery, + updateQuery, + resendQuery, + updateStreamingQuery, + updateThought, + updateStreamingSource, + updateToolCall, + setStatus, + raiseError, + resetPreview, +} = agentPreviewSlice.actions; + +export default agentPreviewSlice.reducer; diff --git a/frontend/src/conversation/conversationSlice.ts b/frontend/src/conversation/conversationSlice.ts index c500532a..708119f3 100644 --- a/frontend/src/conversation/conversationSlice.ts +++ b/frontend/src/conversation/conversationSlice.ts @@ -32,200 +32,189 @@ export function handleAbort() { export const fetchAnswer = createAsyncThunk< Answer, - { question: string; indx?: number; isPreview?: boolean } ->( - 'fetchAnswer', - async ({ question, indx, isPreview = false }, { dispatch, getState }) => { - if (abortController) abortController.abort(); - abortController = new AbortController(); - const { signal } = abortController; + { question: string; indx?: number } +>('fetchAnswer', async ({ question, indx }, { dispatch, getState }) => { + if (abortController) abortController.abort(); + abortController = new AbortController(); + const { signal } = abortController; - let isSourceUpdated = false; - const state = getState() as RootState; - const attachmentIds = selectCompletedAttachments(state) - .filter((a) => a.id) - .map((a) => a.id) as string[]; + let isSourceUpdated = false; + const state = getState() as RootState; + const attachmentIds = selectCompletedAttachments(state) + .filter((a) => a.id) + .map((a) => a.id) as string[]; - if (attachmentIds.length > 0) { - dispatch(clearAttachments()); - } + if (attachmentIds.length > 0) { + dispatch(clearAttachments()); + } - const currentConversationId = state.conversation.conversationId; - const conversationIdToSend = isPreview ? null : currentConversationId; - const save_conversation = isPreview ? false : true; + const currentConversationId = state.conversation.conversationId; - if (state.preference) { - if (API_STREAMING) { - await handleFetchAnswerSteaming( - question, - signal, - state.preference.token, - state.preference.selectedDocs!, - state.conversation.queries, - conversationIdToSend, - state.preference.prompt.id, - state.preference.chunks, - state.preference.token_limit, - (event) => { - const data = JSON.parse(event.data); - const targetIndex = indx ?? state.conversation.queries.length - 1; + if (state.preference) { + if (API_STREAMING) { + await handleFetchAnswerSteaming( + question, + signal, + state.preference.token, + state.preference.selectedDocs!, + state.conversation.queries, + currentConversationId, + state.preference.prompt.id, + state.preference.chunks, + state.preference.token_limit, + (event) => { + const data = JSON.parse(event.data); + const targetIndex = indx ?? state.conversation.queries.length - 1; - // Only process events if they match the current conversation - if (currentConversationId === state.conversation.conversationId) { - if (data.type === 'end') { - dispatch(conversationSlice.actions.setStatus('idle')); - if (!isPreview) { - getConversations(state.preference.token) - .then((fetchedConversations) => { - dispatch(setConversations(fetchedConversations)); - }) - .catch((error) => { - console.error('Failed to fetch conversations: ', error); - }); - } - if (!isSourceUpdated) { - dispatch( - updateStreamingSource({ - conversationId: currentConversationId, - index: targetIndex, - query: { sources: [] }, - }), - ); - } - } else if (data.type === 'id') { - if (!isPreview) { - // Only update the conversationId if it's currently null - const currentState = getState() as RootState; - if (currentState.conversation.conversationId === null) { - dispatch( - updateConversationId({ - query: { conversationId: data.id }, - }), - ); - } - } - } else if (data.type === 'thought') { - const result = data.thought; - dispatch( - updateThought({ - conversationId: currentConversationId, - index: targetIndex, - query: { thought: result }, - }), - ); - } else if (data.type === 'source') { - isSourceUpdated = true; + // Only process events if they match the current conversation + if (currentConversationId === state.conversation.conversationId) { + if (data.type === 'end') { + dispatch(conversationSlice.actions.setStatus('idle')); + getConversations(state.preference.token) + .then((fetchedConversations) => { + dispatch(setConversations(fetchedConversations)); + }) + .catch((error) => { + console.error('Failed to fetch conversations: ', error); + }); + if (!isSourceUpdated) { dispatch( updateStreamingSource({ conversationId: currentConversationId, index: targetIndex, - query: { sources: data.source ?? [] }, - }), - ); - } else if (data.type === 'tool_call') { - dispatch( - updateToolCall({ - index: targetIndex, - tool_call: data.data as ToolCallsType, - }), - ); - } else if (data.type === 'error') { - // set status to 'failed' - dispatch(conversationSlice.actions.setStatus('failed')); - dispatch( - conversationSlice.actions.raiseError({ - conversationId: currentConversationId, - index: targetIndex, - message: data.error, - }), - ); - } else { - dispatch( - updateStreamingQuery({ - conversationId: currentConversationId, - index: targetIndex, - query: { response: data.answer }, + query: { sources: [] }, }), ); } + } else if (data.type === 'id') { + // Only update the conversationId if it's currently null + const currentState = getState() as RootState; + if (currentState.conversation.conversationId === null) { + dispatch( + updateConversationId({ + query: { conversationId: data.id }, + }), + ); + } + } else if (data.type === 'thought') { + const result = data.thought; + dispatch( + updateThought({ + conversationId: currentConversationId, + index: targetIndex, + query: { thought: result }, + }), + ); + } else if (data.type === 'source') { + isSourceUpdated = true; + dispatch( + updateStreamingSource({ + conversationId: currentConversationId, + index: targetIndex, + query: { sources: data.source ?? [] }, + }), + ); + } else if (data.type === 'tool_call') { + dispatch( + updateToolCall({ + index: targetIndex, + tool_call: data.data as ToolCallsType, + }), + ); + } else if (data.type === 'error') { + // set status to 'failed' + dispatch(conversationSlice.actions.setStatus('failed')); + dispatch( + conversationSlice.actions.raiseError({ + conversationId: currentConversationId, + index: targetIndex, + message: data.error, + }), + ); + } else { + dispatch( + updateStreamingQuery({ + conversationId: currentConversationId, + index: targetIndex, + query: { response: data.answer }, + }), + ); } - }, - indx, - state.preference.selectedAgent?.id, - attachmentIds, - save_conversation, - ); - } else { - const answer = await handleFetchAnswer( - question, - signal, - state.preference.token, - state.preference.selectedDocs!, - state.conversation.queries, - state.conversation.conversationId, - state.preference.prompt.id, - state.preference.chunks, - state.preference.token_limit, - state.preference.selectedAgent?.id, - attachmentIds, - save_conversation, - ); - if (answer) { - let sourcesPrepped = []; - sourcesPrepped = answer.sources.map((source: { title: string }) => { - if (source && source.title) { - const titleParts = source.title.split('/'); - return { - ...source, - title: titleParts[titleParts.length - 1], - }; - } - return source; - }); - - const targetIndex = indx ?? state.conversation.queries.length - 1; - - dispatch( - updateQuery({ - index: targetIndex, - query: { - response: answer.answer, - thought: answer.thought, - sources: sourcesPrepped, - tool_calls: answer.toolCalls, - }, - }), - ); - if (!isPreview) { - dispatch( - updateConversationId({ - query: { conversationId: answer.conversationId }, - }), - ); - getConversations(state.preference.token) - .then((fetchedConversations) => { - dispatch(setConversations(fetchedConversations)); - }) - .catch((error) => { - console.error('Failed to fetch conversations: ', error); - }); } - dispatch(conversationSlice.actions.setStatus('idle')); - } + }, + indx, + state.preference.selectedAgent?.id, + attachmentIds, + true, // Always save conversation + ); + } else { + const answer = await handleFetchAnswer( + question, + signal, + state.preference.token, + state.preference.selectedDocs!, + state.conversation.queries, + state.conversation.conversationId, + state.preference.prompt.id, + state.preference.chunks, + state.preference.token_limit, + state.preference.selectedAgent?.id, + attachmentIds, + true, // Always save conversation + ); + if (answer) { + let sourcesPrepped = []; + sourcesPrepped = answer.sources.map((source: { title: string }) => { + if (source && source.title) { + const titleParts = source.title.split('/'); + return { + ...source, + title: titleParts[titleParts.length - 1], + }; + } + return source; + }); + + const targetIndex = indx ?? state.conversation.queries.length - 1; + + dispatch( + updateQuery({ + index: targetIndex, + query: { + response: answer.answer, + thought: answer.thought, + sources: sourcesPrepped, + tool_calls: answer.toolCalls, + }, + }), + ); + dispatch( + updateConversationId({ + query: { conversationId: answer.conversationId }, + }), + ); + getConversations(state.preference.token) + .then((fetchedConversations) => { + dispatch(setConversations(fetchedConversations)); + }) + .catch((error) => { + console.error('Failed to fetch conversations: ', error); + }); + dispatch(conversationSlice.actions.setStatus('idle')); } } - return { - conversationId: null, - title: null, - answer: '', - query: question, - result: '', - thought: '', - sources: [], - tool_calls: [], - }; - }, -); + } + return { + conversationId: null, + title: null, + answer: '', + query: question, + result: '', + thought: '', + sources: [], + tool_calls: [], + }; +}); export const conversationSlice = createSlice({ name: 'conversation', diff --git a/frontend/src/store.ts b/frontend/src/store.ts index 1f25606d..76a34e71 100644 --- a/frontend/src/store.ts +++ b/frontend/src/store.ts @@ -8,6 +8,7 @@ import { prefSlice, } from './preferences/preferenceSlice'; import uploadReducer from './upload/uploadSlice'; +import agentPreviewReducer from './agents/agentPreviewSlice'; const key = localStorage.getItem('DocsGPTApiKey'); const prompt = localStorage.getItem('DocsGPTPrompt'); @@ -54,6 +55,7 @@ const store = configureStore({ conversation: conversationSlice.reducer, sharedConversation: sharedConversationSlice.reducer, upload: uploadReducer, + agentPreview: agentPreviewReducer, }, middleware: (getDefaultMiddleware) => getDefaultMiddleware().concat(prefListenerMiddleware.middleware),