feat: add google chat stream support

src/providers/google/api.ts

@@ -1,15 +1,16 @@
 export interface GoogleFetchPayload {
  apiKey: string
+ stream: boolean
  body: Record<string, any>
  model?: string
 }
 
 export const fetchChatCompletion = async(payload: GoogleFetchPayload) => {
- const { apiKey, body, model } = payload || {}
+ const { apiKey, body, model, stream } = payload || {}
  const initOptions = {
  headers: { 'Content-Type': 'application/json' },
  method: 'POST',
  body: JSON.stringify({ ...body }),
  }
- return fetch(`https://generativelanguage.googleapis.com/v1/models/${model}:generateContent?key=${apiKey}`, initOptions)
+ return fetch(`https://generativelanguage.googleapis.com/v1beta/models/${model}:streamGenerateContent?${stream ? 'alt=sse&' : ''}key=${apiKey}`, initOptions)
 }

src/providers/google/handler.ts

@@ -1,5 +1,5 @@
 import { fetchChatCompletion } from './api'
-import { parseMessageList } from './parser'
+import { parseMessageList, parseStream } from './parser'
 import type { Message } from '@/types/message'
 import type { HandlerPayload, Provider } from '@/types/provider'
 
@@ -51,15 +51,19 @@ export const handleChatCompletion = async(payload: HandlerPayload, signal?: Abor
  messages.unshift(m)
  }
 
+ const stream = payload.globalSettings.stream as boolean ?? true
  const response = await fetchChatCompletion({
  apiKey: payload.globalSettings.apiKey as string,
+ stream,
  body: {
  contents: parseMessageList(messages),
  },
  model: payload.globalSettings.model as string,
  })
 
  if (response.ok) {
+ if (stream)
+ return parseStream(response)
  const json = await response.json()
  // console.log('json', json)
  const output = json.candidates[0].content.parts[0].text || json

src/providers/google/parser.ts

@@ -1,4 +1,6 @@
+import { createParser } from 'eventsource-parser'
 import type { Message } from '@/types/message'
+import type { ParsedEvent, ReconnectInterval } from 'eventsource-parser'
 
 export const parseMessageList = (rawList: Message[]) => {
  interface GoogleGeminiMessage {
@@ -33,3 +35,41 @@ export const parseMessageList = (rawList: Message[]) => {
  }
  return parsedList
 }
+
+export const parseStream = (rawResponse: Response) => {
+ const encoder = new TextEncoder()
+ const decoder = new TextDecoder()
+ const rb = rawResponse.body as ReadableStream
+
+ return new ReadableStream({
+ async start(controller) {
+ const streamParser = (event: ParsedEvent | ReconnectInterval) => {
+ if (event.type === 'event') {
+ const data = event.data
+ try {
+ const json = JSON.parse(data)
+ const text = json.candidates[0].content.parts[0].text || ''
+ const queue = encoder.encode(text)
+ controller.enqueue(queue)
+ if (json.candidates[0].finishReason === 'STOP')
+ controller.close()
+ } catch (e) {
+ controller.error(e)
+ }
+ }
+ }
+ const reader = rb.getReader()
+ const parser = createParser(streamParser)
+ let done = false
+ while (!done) {
+ const { done: isDone, value } = await reader.read()
+ if (isDone) {
+ done = true
+ controller.close()
+ return
+ }
+ parser.feed(decoder.decode(value))
+ }
+ },
+ })
+}