diff --git a/__tests__/ChatScreen.test.tsx b/__tests__/ChatScreen.test.tsx index ebb943d..f7d08a2 100644 --- a/__tests__/ChatScreen.test.tsx +++ b/__tests__/ChatScreen.test.tsx @@ -165,7 +165,7 @@ describe('ChatScreen', () => { expect(apiMessages[0].role).toBe('system'); }); - it('edits without regenerating when no model is selected', () => { + it('edits and explains that a model is needed when none is selected', () => { const id = setup(); act(() => { useAppStore.setState({ activeModelId: '' }); @@ -177,8 +177,9 @@ describe('ChatScreen', () => { jest.advanceTimersByTime(100); }); - expect(state().messages).toHaveLength(1); + expect(state().messages).toHaveLength(2); expect(state().messages[0]).toMatchObject({ id, content: 'new question' }); + expect(state().messages[1].content).toContain('No model selected'); expect(streamMock).not.toHaveBeenCalled(); }); diff --git a/__tests__/useLMStudio.test.ts b/__tests__/useLMStudio.test.ts index 555b9ff..5e8a542 100644 --- a/__tests__/useLMStudio.test.ts +++ b/__tests__/useLMStudio.test.ts @@ -177,22 +177,34 @@ describe('regenerateResponse', () => { expect(state().isStreaming).toBe(true); }); - it('does nothing without a user message, while streaming, or without a model', async () => { + it('keeps the current answer and asks for a model when none is selected', async () => { + state().addMessage({ role: 'user', content: 'q' }); + state().addMessage({ role: 'assistant', content: 'a' }); + useAppStore.setState({ activeModelId: '' }); + const { result } = renderHook(() => useLMStudio()); + + await act(async () => { + await result.current.regenerateResponse(); + }); + expect(mockApi.chatCompletionStream).not.toHaveBeenCalled(); + expect(state().messages.map((m) => m.content)).toEqual([ + 'q', + 'a', + expect.stringContaining('No model selected'), + ]); + }); + + it('does nothing without a user message or while streaming', async () => { const { result, rerender } = renderHook(() => useLMStudio()); await act(async () => { await result.current.regenerateResponse(); }); expect(mockApi.chatCompletionStream).not.toHaveBeenCalled(); + expect(state().messages).toEqual([]); act(() => { state().addMessage({ role: 'user', content: 'q' }); - useAppStore.setState({ activeModelId: '' }); }); - await act(async () => { - await result.current.regenerateResponse(); - }); - expect(mockApi.chatCompletionStream).not.toHaveBeenCalled(); - act(() => { useAppStore.setState({ activeModelId: 'model-a', isStreaming: true }); }); diff --git a/hooks/useLMStudio.ts b/hooks/useLMStudio.ts index 13e0af5..1d36b70 100644 --- a/hooks/useLMStudio.ts +++ b/hooks/useLMStudio.ts @@ -2,7 +2,16 @@ import { useCallback, useRef } from 'react'; import { SYSTEM_PROMPT } from '../constants/SystemPrompt'; import * as api from '../services/lmStudioApi'; import { ChatMessage as ApiMessage } from '../services/types'; -import { useAppStore } from '../store/useAppStore'; +import { ChatMessage, useAppStore } from '../store/useAppStore'; + +const NO_MODEL_MESSAGE = '⚠️ No model selected. Please select a model first.'; + +function buildApiMessages(messages: ChatMessage[]): ApiMessage[] { + return [ + { role: 'system', content: SYSTEM_PROMPT }, + ...messages.map((m) => ({ role: m.role, content: m.content })), + ]; +} export function useLMStudio() { const serverConfig = useAppStore((s) => s.serverConfig); @@ -44,95 +53,16 @@ export function useLMStudio() { } }, [activeModelId, setAvailableModels, setIsConnected, setActiveModelId]); - const sendMessage = useCallback(async (text: string) => { - if (isStreaming) return; - - addMessage({ role: 'user', content: text }); - - const currentMessages = useAppStore.getState().messages; - const apiMessages: ApiMessage[] = [ - { role: 'system', content: SYSTEM_PROMPT }, - ...currentMessages.map((m) => ({ - role: m.role, - content: m.content, - })) - ]; - - const modelId = useAppStore.getState().activeModelId; - if (!modelId) { - addMessage({ - role: 'assistant', - content: '⚠️ No model selected. Please select a model first.', - }); - return; - } - + // Adds an empty assistant message and streams the model's answer into it. + const streamAnswer = useCallback((modelId: string, apiMessages: ApiMessage[]) => { const assistantId = addMessage({ role: 'assistant', content: '' }); setIsStreaming(true); - try { - abortRef.current = api.chatCompletionStream( - serverConfig, - modelId, - apiMessages, - (chunk) => { - appendToMessage(assistantId, chunk); - }, - () => { - setIsStreaming(false); - abortRef.current = null; - useAppStore.getState().saveCurrentChat(); - }, - (error) => { - const errorMsg = `\n\n⚠️ Error: ${error.message}`; - appendToMessage(assistantId, errorMsg); - setIsStreaming(false); - abortRef.current = null; - useAppStore.getState().saveCurrentChat(); - }, - ); - } catch (err) { - const errorMsg = `⚠️ Connection failed: ${(err as Error).message}`; - appendToMessage(assistantId, errorMsg); + const finish = () => { setIsStreaming(false); + abortRef.current = null; useAppStore.getState().saveCurrentChat(); - } - }, [isStreaming, serverConfig, addMessage, appendToMessage, setIsStreaming]); - - const regenerateResponse = useCallback(async () => { - if (isStreaming) return; - - const msgs = useAppStore.getState().messages; - // Find the last user message - let lastUserIdx = -1; - for (let i = msgs.length - 1; i >= 0; i--) { - if (msgs[i].role === 'user') { - lastUserIdx = i; - break; - } - } - if (lastUserIdx === -1) return; - - const lastUserMsg = msgs[lastUserIdx]; - // Remove messages after the last user message - const store = useAppStore.getState(); - store.deleteMessagesAfter(lastUserMsg.id); - - // Build API messages from remaining - const remaining = useAppStore.getState().messages; - const apiMessages: ApiMessage[] = [ - { role: 'system', content: SYSTEM_PROMPT }, - ...remaining.map((m) => ({ - role: m.role, - content: m.content, - })) - ]; - - const modelId = useAppStore.getState().activeModelId; - if (!modelId) return; - - const assistantId = addMessage({ role: 'assistant', content: '' }); - setIsStreaming(true); + }; try { abortRef.current = api.chatCompletionStream( @@ -140,24 +70,51 @@ export function useLMStudio() { modelId, apiMessages, (chunk) => appendToMessage(assistantId, chunk), - () => { - setIsStreaming(false); - abortRef.current = null; - useAppStore.getState().saveCurrentChat(); - }, + finish, (error) => { - appendToMessage(assistantId, `\n\n⚠️ Error: ${error.message}`); - setIsStreaming(false); - abortRef.current = null; - useAppStore.getState().saveCurrentChat(); + appendToMessage(assistantId, ` + +⚠️ Error: ${error.message}`); + finish(); }, ); } catch (err) { appendToMessage(assistantId, `⚠️ Connection failed: ${(err as Error).message}`); - setIsStreaming(false); - useAppStore.getState().saveCurrentChat(); + finish(); } - }, [isStreaming, serverConfig, addMessage, appendToMessage, setIsStreaming]); + }, [serverConfig, addMessage, appendToMessage, setIsStreaming]); + + const sendMessage = useCallback(async (text: string) => { + if (isStreaming) return; + + addMessage({ role: 'user', content: text }); + + const modelId = useAppStore.getState().activeModelId; + if (!modelId) { + addMessage({ role: 'assistant', content: NO_MODEL_MESSAGE }); + return; + } + + streamAnswer(modelId, buildApiMessages(useAppStore.getState().messages)); + }, [isStreaming, addMessage, streamAnswer]); + + const regenerateResponse = useCallback(async () => { + if (isStreaming) return; + + const { messages: msgs, activeModelId: modelId, deleteMessagesAfter } = useAppStore.getState(); + const lastUserMsg = [...msgs].reverse().find((m) => m.role === 'user'); + if (!lastUserMsg) return; + + // Checked before deleting anything, so the current answer is not lost for nothing. + if (!modelId) { + addMessage({ role: 'assistant', content: NO_MODEL_MESSAGE }); + return; + } + + // Drop the previous answer(s) and ask again from the last user message. + deleteMessagesAfter(lastUserMsg.id); + streamAnswer(modelId, buildApiMessages(useAppStore.getState().messages)); + }, [isStreaming, addMessage, streamAnswer]); const stopStreaming = useCallback(() => { if (abortRef.current) {