File size: 2,373 Bytes
a1c5622
a86b547
a1c5622
0fd8446
ca7a659
a1c5622
 
 
 
 
 
 
4af6326
 
ae074fc
a86b547
5ec491a
a1c5622
 
4af6326
d553ae5
a1c5622
 
 
 
 
a86b547
92f037b
a86b547
 
 
 
 
5d7d435
a1c5622
 
 
 
4af6326
84c9f51
a1c5622
a86b547
a1c5622
a86b547
 
ca7a659
 
 
a86b547
 
5ec491a
92f037b
5ec491a
c232e44
84c9f51
c232e44
 
 
 
 
 
 
84c9f51
 
 
 
a86b547
0fd8446
ae074fc
 
 
a1c5622
 
 
ae074fc
a1c5622
 
a86b547
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
import { useChat } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useRef, useState } from 'react';
import { ChatWithMessages, MessageUI, MessageUserInput } from '../types';
import {
  dbPostCreateMessage,
  dbPostUpdateMessageResponse,
} from '../db/functions';
import {
  convertAssistantUIMessageToDBMessageResponse,
  convertDBMessageToUIMessage,
} from '../utils/message';
import { useSetAtom } from 'jotai';
import { selectedMessageId } from '@/state/chat';
import { Message } from '@prisma/client';

const useVisionAgent = (chat: ChatWithMessages) => {
  const { messages: dbMessages, id, mediaUrl } = chat;
  const latestDbMessage = dbMessages[dbMessages.length - 1];
  const setMessageId = useSetAtom(selectedMessageId);

  // Temporary solution for now while single we have to pass mediaUrl separately outside of the messages
  const currMediaUrl = useRef<string>(mediaUrl);
  const currMessageId = useRef<string>(latestDbMessage?.id);

  const { messages, append, isLoading, reload } = useChat({
    api: '/api/vision-agent',
    streamMode: 'text',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: async message => {
      await dbPostUpdateMessageResponse(
        currMessageId.current,
        convertAssistantUIMessageToDBMessageResponse(message),
      );
      setMessageId(currMessageId.current);
    },
    initialMessages: convertDBMessageToUIMessage(dbMessages),
    body: {
      mediaUrl: currMediaUrl.current,
      id,
    },
    onError: err => {
      err && toast.error(err.message);
    },
  });

  /**
   * If case this is first time user navigated with init message, we need to reload the chat for the first response
   */
  const once = useRef(true);
  useEffect(() => {
    if (
      !isLoading &&
      messages.length === 1 &&
      messages[0].role === 'user' &&
      once.current
    ) {
      once.current = false;
      reload();
    }
  }, [isLoading, messages, reload]);

  return {
    messages: messages as MessageUI[],
    append: (message: Message) => {
      currMediaUrl.current = message.mediaUrl;
      currMessageId.current = message.id;
      append({
        id,
        role: 'user',
        content: message.prompt,
      });
    },
    reload,
    isLoading,
  };
};

export default useVisionAgent;