konnosDashboard
queryAi.ts69 lines · main
1import type OpenAI from 'openai'
2import { SSE } from 'sse.js'
3
4import { BASE_PATH } from '../shared/constants'
5import { type Message } from './utils'
6
7/**
8 * Perform a one-off query to AI based on a snapshot of messages
9 */
10export function queryAi(messages: Message[], timeout = 0) {
11 return new Promise<string>((resolve, reject) => {
12 const eventSource = new SSE(`${BASE_PATH}/api/ai/docs`, {
13 headers: {
14 apikey: process.env.NEXT_PUBLIC_BRIVEN_ANON_KEY ?? '',
15 Authorization: `Bearer ${process.env.NEXT_PUBLIC_BRIVEN_ANON_KEY}`,
16 'Content-Type': 'application/json',
17 },
18 payload: JSON.stringify({
19 messages: messages.map(({ role, content }) => ({ role, content })),
20 }),
21 })
22
23 let timeoutId: number | undefined
24
25 function handleError<T>(err: T) {
26 if (timeoutId) {
27 clearTimeout(timeoutId)
28 }
29 console.error(err)
30 reject(err)
31 }
32
33 if (timeout > 0) {
34 timeoutId = window.setTimeout(() => {
35 handleError(new Error('AI query timed out'))
36 }, timeout)
37 }
38
39 let answer = ''
40
41 eventSource.addEventListener('error', handleError)
42 eventSource.addEventListener('message', (e: MessageEvent) => {
43 try {
44 if (e.data === '[DONE]') {
45 if (timeoutId) {
46 clearTimeout(timeoutId)
47 }
48 resolve(answer)
49 return
50 }
51
52 const completionChunk: OpenAI.Chat.Completions.ChatCompletionChunk = JSON.parse(e.data)
53 const [
54 {
55 delta: { content },
56 },
57 ] = completionChunk.choices
58
59 if (content) {
60 answer += content
61 }
62 } catch (err) {
63 handleError(err)
64 }
65 })
66
67 eventSource.stream()
68 })
69}