File size: 2,885 Bytes
84c9f51
a86b547
 
5ec491a
96ac62a
0d6f04b
 
 
 
 
0a08909
d553ae5
5ec491a
 
5d7d435
 
 
 
 
 
 
0a08909
 
 
5d7d435
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a86b547
5ec491a
 
d553ae5
a86b547
 
84c9f51
a86b547
92f037b
a86b547
 
92f037b
 
a86b547
 
 
 
 
5d7d435
3ebf44a
 
 
 
84c9f51
a86b547
 
5ec491a
a86b547
 
 
 
5ec491a
92f037b
5ec491a
84c9f51
92f037b
84c9f51
 
 
 
 
5ec491a
 
 
 
84c9f51
 
 
a86b547
5ec491a
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import { useChat, type Message, UseChatHelpers } from 'ai/react';
import { toast } from 'react-hot-toast';
import { useEffect, useState } from 'react';
import { MessageBase, SignedPayload } from '../types';
import { fetcher, nanoid } from '../utils';
import {
  getCleanedUpMessages,
  generateAnswersImageMarkdown,
  generateInputImageMarkdown,
} from '../messageUtils';
import { CLEANED_SEPARATOR } from '../constants';
import { useSearchParams } from 'next/navigation';
import { ChatWithMessages, MessageRaw } from '../db/types';
import { dbPostCreateMessage } from '../db/functions';

const uploadBase64 = async (
  base64: string,
  messageId: string,
  chatId: string,
  index: number,
) => {
  const res = await fetch(
    'data:image/png;base64,' + base64.replace('base:64', ''),
  );
  const blob = await res.blob();
  const { signedUrl, publicUrl, fields } = await fetcher<SignedPayload>(
    '/api/sign',
    {
      method: 'POST',
      body: JSON.stringify({
        id: `${chatId}/${messageId}`,
        fileType: blob.type,
        fileName: `answer-${index}.${blob.type.split('/')[1]}`,
      }),
    },
  );
  const formData = new FormData();
  Object.entries(fields).forEach(([key, value]) => {
    formData.append(key, value as string);
  });
  formData.append('file', blob);

  const uploadResponse = await fetch(signedUrl, {
    method: 'POST',
    body: formData,
  });
  if (uploadResponse.ok) {
    return publicUrl;
  } else {
    throw new Error('Upload failed');
  }
};

const useVisionAgent = (chat: ChatWithMessages) => {
  const { messages: initialMessages, id, mediaUrl } = chat;

  const {
    messages,
    append: appendRaw,
    isLoading,
    reload,
  } = useChat({
    api: '/api/vision-agent',
    // @ts-ignore https://sdk.vercel.ai/docs/troubleshooting/common-issues/use-chat-failed-to-parse-stream
    streamMode: 'text',
    onResponse(response) {
      if (response.status !== 200) {
        toast.error(response.statusText);
      }
    },
    onFinish: async message => {
      await dbPostCreateMessage(id, {
        role: message.role as 'user' | 'assistant',
        content: message.content,
      });
    },
    initialMessages: initialMessages,
    body: {
      mediaUrl,
      id,
    },
  });

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

  const append: UseChatHelpers['append'] = async message => {
    dbPostCreateMessage(id, {
      role: message.role as 'user' | 'assistant',
      content: message.content,
    });
    return appendRaw(message);
  };

  return {
    messages: messages as MessageBase[],
    append,
    reload,
    isLoading,
  };
};

export default useVisionAgent;