99 lines
2.2 KiB
TypeScript
99 lines
2.2 KiB
TypeScript
import {
|
|
type AIStreamCallbacksAndOptions,
|
|
createCallbacksTransformer,
|
|
createStreamDataTransformer
|
|
} from 'ai'
|
|
import { StreamingTextResponse } from 'ai'
|
|
import { useEffect } from 'react'
|
|
|
|
const serverUrl = process.env.NEXT_PUBLIC_TABBY_SERVER_URL || ''
|
|
|
|
export function usePatchFetch() {
|
|
useEffect(() => {
|
|
const fetch = window.fetch
|
|
|
|
window.fetch = async function (url, options) {
|
|
if (url !== '/api/chat') {
|
|
return fetch(url, options)
|
|
}
|
|
|
|
const res = await fetch(`${serverUrl}/v1beta/chat/completions`, {
|
|
...options,
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json'
|
|
}
|
|
})
|
|
|
|
const stream = StreamAdapter(res, undefined)
|
|
return new StreamingTextResponse(stream)
|
|
}
|
|
}, [])
|
|
}
|
|
|
|
const utf8Decoder = new TextDecoder('utf-8')
|
|
|
|
async function processLines(
|
|
lines: string[],
|
|
controller: ReadableStreamDefaultController<string>
|
|
) {
|
|
for (const line of lines) {
|
|
const { content } = JSON.parse(line)
|
|
controller.enqueue(content)
|
|
}
|
|
}
|
|
|
|
async function readAndProcessLines(
|
|
reader: ReadableStreamDefaultReader<Uint8Array>,
|
|
controller: ReadableStreamDefaultController<string>
|
|
) {
|
|
let segment = ''
|
|
|
|
while (true) {
|
|
const { value: chunk, done } = await reader.read()
|
|
if (done) {
|
|
break
|
|
}
|
|
|
|
segment += utf8Decoder.decode(chunk, { stream: true })
|
|
|
|
const linesArray = segment.split(/\r\n|\n|\r/g)
|
|
segment = linesArray.pop() || ''
|
|
|
|
await processLines(linesArray, controller)
|
|
}
|
|
|
|
if (segment) {
|
|
const linesArray = [segment]
|
|
await processLines(linesArray, controller)
|
|
}
|
|
|
|
controller.close()
|
|
}
|
|
|
|
function createParser(res: Response) {
|
|
const reader = res.body?.getReader()
|
|
|
|
return new ReadableStream<string>({
|
|
async start(controller): Promise<void> {
|
|
if (!reader) {
|
|
controller.close()
|
|
return
|
|
}
|
|
|
|
await readAndProcessLines(reader, controller)
|
|
}
|
|
})
|
|
}
|
|
|
|
function StreamAdapter(
|
|
reader: Response,
|
|
callbacks?: AIStreamCallbacksAndOptions
|
|
): ReadableStream {
|
|
return createParser(reader)
|
|
.pipeThrough(createCallbacksTransformer(callbacks))
|
|
.pipeThrough(
|
|
createStreamDataTransformer(callbacks?.experimental_streamData)
|
|
)
|
|
}
|