diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts new file mode 100644 index 0000000..58539be --- /dev/null +++ b/app/api/chat/route.ts @@ -0,0 +1,53 @@ +import { NextRequest } from 'next/server'; +import ollama from '@/lib/ollama'; + +export async function POST(request: NextRequest) { + try { + const { model, messages } = await request.json(); + + if (!model || !messages) { + return new Response(JSON.stringify({ error: 'Model and messages are required' }), { + status: 400, + headers: { 'Content-Type': 'application/json' }, + }); + } + + const response = await ollama.chat({ + model, + messages, + stream: true, + }); + + // Create a readable stream from the Ollama response + const stream = new ReadableStream({ + async start(controller) { + const encoder = new TextEncoder(); + + try { + for await (const chunk of response) { + const text = chunk.message?.content || ''; + if (text) { + controller.enqueue(encoder.encode(text)); + } + } + controller.close(); + } catch (error) { + controller.error(error); + } + }, + }); + + return new Response(stream, { + headers: { + 'Content-Type': 'text/plain; charset=utf-8', + 'Transfer-Encoding': 'chunked', + }, + }); + } catch (error: any) { + console.error('Chat stream error:', error); + return new Response(JSON.stringify({ error: error.message || 'Failed to stream response' }), { + status: 500, + headers: { 'Content-Type': 'application/json' }, + }); + } +} diff --git a/components/Chat/ChatLayout.tsx b/components/Chat/ChatLayout.tsx index d96c14c..03f3f5b 100644 --- a/components/Chat/ChatLayout.tsx +++ b/components/Chat/ChatLayout.tsx @@ -30,7 +30,6 @@ import { useMantineTheme, } from '@mantine/core'; import { useDisclosure } from '@mantine/hooks'; -import { chat, type ChatMessage } from '@/app/actions/chat'; import { getInstalledModels, type OllamaModel } from '@/app/actions/ollama'; import { useThemeContext } from '@/components/DynamicThemeProvider'; import { SettingsModal } from '@/components/Settings/SettingsModal'; @@ -91,7 +90,8 @@ export default function ChatLayout() { // Model State const [models, setModels] = useState([]); const [selectedModel, setSelectedModel] = useState(null); - const [isGenerating, setIsGenerating] = useState(false); + const [_isGenerating, setIsGenerating] = useState(false); + const [streamingMessageId, setStreamingMessageId] = useState(null); // Fetch chats and models on load useEffect(() => { @@ -157,7 +157,9 @@ export default function ChatLayout() { }; const handleSendMessage = async () => { - if (!inputValue.trim() || !selectedModel) return; + if (!inputValue.trim() || !selectedModel) { + return; + } const userMessage: Message = { id: Date.now().toString(), @@ -165,78 +167,114 @@ export default function ChatLayout() { content: inputValue, }; - // Optimistic update + // Optimistic update - add user message and empty assistant message for streaming + const assistantMessageId = (Date.now() + 1).toString(); const newMessages = [...messages, userMessage]; - setMessages(newMessages); + setMessages([...newMessages, { id: assistantMessageId, role: 'assistant', content: '' }]); setInputValue(''); setIsGenerating(true); + setStreamingMessageId(assistantMessageId); try { - // Convert to format expected by server action - const chatHistory: ChatMessage[] = newMessages.map((m) => ({ + // Convert to format expected by API + const chatHistory = newMessages.map((m) => ({ role: m.role as 'user' | 'assistant', content: m.content, })); - // Call Ollama via Server Action - const result = await chat(selectedModel, chatHistory); + // Call streaming API + const response = await fetch('/api/chat', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: selectedModel, + messages: chatHistory, + }), + }); - if (result.success && result.message) { - const responseMessage: Message = { - id: (Date.now() + 1).toString(), - role: 'assistant', - content: result.message.content, - }; - const finalMessages = [...newMessages, responseMessage]; - setMessages(finalMessages); + if (!response.ok) { + const errorData = await response.json(); + throw new Error(errorData.error || 'Failed to get response'); + } - // Save both user message and assistant response to database - try { - // Save user message - const userSaveRes = await fetch('/api/chats', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - chatId: activeChatId, - messages: [userMessage], - }), - }); - const userSaveData = await userSaveRes.json(); - const savedChatId = userSaveData.chatId; + // Read the stream + const reader = response.body?.getReader(); + if (!reader) { + throw new Error('No response body'); + } - // Update activeChatId if this was a new chat - if (!activeChatId && savedChatId) { - setActiveChatId(savedChatId); - } + const decoder = new TextDecoder(); + let fullContent = ''; - // Save assistant response - await fetch('/api/chats', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - chatId: savedChatId, - messages: [responseMessage], - }), - }); - - // Refresh chat list - fetchChats(); - } catch (saveError) { - console.error('Failed to save messages:', saveError); + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; } - } else { - // Error handling - const errorMessage: Message = { - id: (Date.now() + 1).toString(), - role: 'assistant', - content: `Error: ${result.error}`, - }; - setMessages([...newMessages, errorMessage]); + + const chunk = decoder.decode(value, { stream: true }); + fullContent += chunk; + + // Update the assistant message with accumulated content + setMessages((prev) => + prev.map((m) => (m.id === assistantMessageId ? { ...m, content: fullContent } : m)) + ); + } + + // Create final response message for saving + const responseMessage: Message = { + id: assistantMessageId, + role: 'assistant', + content: fullContent, + }; + + // Save both user message and assistant response to database + try { + // Save user message + const userSaveRes = await fetch('/api/chats', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + chatId: activeChatId, + messages: [userMessage], + }), + }); + const userSaveData = await userSaveRes.json(); + const savedChatId = userSaveData.chatId; + + // Update activeChatId if this was a new chat + if (!activeChatId && savedChatId) { + setActiveChatId(savedChatId); + } + + // Save assistant response + await fetch('/api/chats', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + chatId: savedChatId, + messages: [responseMessage], + }), + }); + + // Refresh chat list + fetchChats(); + } catch (saveError) { + console.error('Failed to save messages:', saveError); } } catch (e) { console.error('Failed to send message', e); + // Update assistant message with error + setMessages((prev) => + prev.map((m) => + m.id === assistantMessageId + ? { ...m, content: `Error: ${e instanceof Error ? e.message : 'Unknown error'}` } + : m + ) + ); } finally { setIsGenerating(false); + setStreamingMessageId(null); } }; @@ -373,7 +411,10 @@ export default function ChatLayout() { }} > {message.role === 'assistant' ? ( - + ) : ( {message.content} diff --git a/components/Chat/MarkdownMessage.module.css b/components/Chat/MarkdownMessage.module.css index b642127..7286410 100644 --- a/components/Chat/MarkdownMessage.module.css +++ b/components/Chat/MarkdownMessage.module.css @@ -7,6 +7,28 @@ margin-bottom: 0; } +/* Streaming cursor that blinks at the end */ +.streamingCursor { + display: inline-block; + width: 0.5em; + height: 1em; + background-color: currentColor; + margin-left: 2px; + animation: blink 1s step-end infinite; + vertical-align: text-bottom; + opacity: 0.7; +} + +@keyframes blink { + 0%, + 100% { + opacity: 0.7; + } + 50% { + opacity: 0; + } +} + .preWrapper { margin: 0.5em 0; } diff --git a/components/Chat/MarkdownMessage.tsx b/components/Chat/MarkdownMessage.tsx index 2842e15..6a563e8 100644 --- a/components/Chat/MarkdownMessage.tsx +++ b/components/Chat/MarkdownMessage.tsx @@ -4,9 +4,10 @@ import classes from './MarkdownMessage.module.css'; interface MarkdownMessageProps { content: string; + isStreaming?: boolean; } -export function MarkdownMessage({ content }: MarkdownMessageProps) { +export function MarkdownMessage({ content, isStreaming = false }: MarkdownMessageProps) { return (
), - code: ({ className, children, ...props }) => { + code: ({ className, children }) => { const isInline = !className; if (isInline) { return {children}; @@ -74,6 +75,7 @@ export function MarkdownMessage({ content }: MarkdownMessageProps) { > {content} + {isStreaming && }
); }