UPDATE use stream custom hook, implement newly added custom hooks in it

Author: blazedglimmer

hooks/use-stream.ts

@@ -1,15 +1,30 @@
 'use client';
 
-import { useState } from 'react';
+import { useState, useRef } from 'react';
 import { Message } from '@/types/chat';
 import { ModelConfig } from '@/hooks/use-model-config';
+import { useMetrics } from '@/hooks/use-metrics';
+import { useConversationCache } from '@/hooks/use-conversation-cache';
 
 export function useStream() {
   const [messages, setMessages] = useState<Message[]>([]);
   const [isLoading, setIsLoading] = useState(false);
+  const abortController = useRef<AbortController | null>(null);
+  const { metrics, updateMetrics, resetMetrics } = useMetrics();
+  const { saveMessages } = useConversationCache();
+
+  const cancelGeneration = () => {
+    if (abortController.current) {
+      abortController.current.abort();
+      abortController.current = null;
+      setIsLoading(false);
+    }
+  };
 
   const streamMessage = async (content: string, config: ModelConfig) => {
     setIsLoading(true);
+    resetMetrics();
+
     const userMessage: Message = {
       role: 'user',
       content,
@@ -19,6 +34,8 @@ export function useStream() {
     setMessages(prev => [...prev, userMessage]);
 
     try {
+      abortController.current = new AbortController();
+
       const response = await fetch('/api/chat', {
         method: 'POST',
         headers: {
@@ -28,6 +45,7 @@ export function useStream() {
           messages: [...messages, userMessage],
           config,
         }),
+        signal: abortController.current.signal,
       });
 
       if (!response.ok) throw new Error('Failed to send message');
@@ -47,15 +65,31 @@ export function useStream() {
 
         const text = new TextDecoder().decode(value);
         assistantMessage.content += text;
+        updateMetrics(text);
 
         setMessages(prev => [...prev.slice(0, -1), { ...assistantMessage }]);
       }
+
+      // Save to IndexedDB after completion
+      const updatedMessages = [...messages, userMessage, assistantMessage];
+      await saveMessages(updatedMessages);
     } catch (error) {
-      console.error('Error streaming message:', error);
+      if (error instanceof Error && error.name === 'AbortError') {
+        console.log('Message generation cancelled');
+      } else {
+        console.error('Error streaming message:', error);
+      }
     } finally {
       setIsLoading(false);
+      abortController.current = null;
     }
   };
 
-  return { messages, isLoading, streamMessage };
+  return {
+    messages,
+    isLoading,
+    streamMessage,
+    cancelGeneration,
+    metrics,
+  };
 }