Northern_Lights/cherry-studio
1
0
Fork 0
mirror of https://github.com/CherryHQ/cherry-studio.git synced 2025-12-25 11:20:07 +08:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

refactor(GeminiProvider): streamline abort signal handling and improve stream processing (#5276)

Browse Source
This commit is contained in:
SuYao 2025-04-24 13:52:29 +08:00 committed by GitHub
parent 25b1e309ed
commit 1f588d242e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 47 additions and 37772 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37698
.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch vendored
View File

File diff suppressed because one or more lines are too long

2
package.json
View File

@ -120,7 +120,7 @@
"@emotion/is-prop-valid": "^1.3.1",
"@eslint-react/eslint-plugin": "^1.36.1",
"@eslint/js": "^9.22.0",
"@google/genai": "patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch",
"@google/genai": "^0.10.0",
"@hello-pangea/dnd": "^16.6.0",
"@kangfenmao/keyv-storage": "^0.1.0",
"@modelcontextprotocol/sdk": "^1.9.0",

95
src/renderer/src/providers/AiProvider/GeminiProvider.ts
View File

@ -320,32 +320,12 @@ export default class GeminiProvider extends BaseProvider {
const start_time_millsec = new Date().getTime()
const { cleanup, abortController } = this.createAbortController(userLastMessage?.id, true)
const signalProxy = {
_originalSignal: abortController.signal,
addEventListener: (eventName: string, listener: () => void) => {
if (eventName === 'abort') {
abortController.signal.addEventListener('abort', listener)
}
},
removeEventListener: (eventName: string, listener: () => void) => {
if (eventName === 'abort') {
abortController.signal.removeEventListener('abort', listener)
}
},
get aborted() {
return abortController.signal.aborted
}
}
if (!streamOutput) {
const response = await chat.sendMessage({
message: messageContents as PartUnion,
config: {
...generateContentConfig,
httpOptions: {
signal: signalProxy as any
}
abortSignal: abortController.signal
}
})
const time_completion_millsec = new Date().getTime() - start_time_millsec
@ -371,9 +351,7 @@ export default class GeminiProvider extends BaseProvider {
message: messageContents as PartUnion,
config: {
...generateContentConfig,
httpOptions: {
signal: signalProxy as any
}
abortSignal: abortController.signal
}
})
let time_first_token_millsec = 0
@ -399,9 +377,7 @@ export default class GeminiProvider extends BaseProvider {
message: flatten(toolResults.map((ts) => (ts as Content).parts)) as PartUnion,
config: {
...generateContentConfig,
httpOptions: {
signal: signalProxy as any
}
abortSignal: abortController.signal
}
})
await processStream(newStream, idx + 1)
@ -410,38 +386,43 @@ export default class GeminiProvider extends BaseProvider {
const processStream = async (stream: AsyncGenerator<GenerateContentResponse>, idx: number) => {
let content = ''
for await (const chunk of stream) {
if (window.keyv.get(EVENT_NAMES.CHAT_COMPLETION_PAUSED)) break
try {
for await (const chunk of stream) {
if (window.keyv.get(EVENT_NAMES.CHAT_COMPLETION_PAUSED)) break
if (time_first_token_millsec == 0) {
time_first_token_millsec = new Date().getTime() - start_time_millsec
if (time_first_token_millsec == 0) {
time_first_token_millsec = new Date().getTime() - start_time_millsec
}
const time_completion_millsec = new Date().getTime() - start_time_millsec
if (chunk.text !== undefined) {
content += chunk.text
}
await processToolUses(content, idx)
const generateImage = this.processGeminiImageResponse(chunk)
onChunk({
text: chunk.text !== undefined ? chunk.text : '',
usage: {
prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
thoughts_tokens: chunk.usageMetadata?.thoughtsTokenCount || 0,
total_tokens: chunk.usageMetadata?.totalTokenCount || 0
},
metrics: {
completion_tokens: chunk.usageMetadata?.candidatesTokenCount,
time_completion_millsec,
time_first_token_millsec
},
search: chunk.candidates?.[0]?.groundingMetadata,
mcpToolResponse: toolResponses,
generateImage: generateImage
})
}
const time_completion_millsec = new Date().getTime() - start_time_millsec
if (chunk.text !== undefined) {
content += chunk.text
}
await processToolUses(content, idx)
const generateImage = this.processGeminiImageResponse(chunk)
onChunk({
text: chunk.text !== undefined ? chunk.text : '',
usage: {
prompt_tokens: chunk.usageMetadata?.promptTokenCount || 0,
completion_tokens: chunk.usageMetadata?.candidatesTokenCount || 0,
thoughts_tokens: chunk.usageMetadata?.thoughtsTokenCount || 0,
total_tokens: chunk.usageMetadata?.totalTokenCount || 0
},
metrics: {
completion_tokens: chunk.usageMetadata?.candidatesTokenCount,
time_completion_millsec,
time_first_token_millsec
},
search: chunk.candidates?.[0]?.groundingMetadata,
mcpToolResponse: toolResponses,
generateImage: generateImage
})
} catch (error) {
console.error('Error processing stream chunk:', error)
throw error
}
}

24
yarn.lock
View File

@ -1489,23 +1489,15 @@ __metadata:
languageName: node
linkType: hard
"@google/genai@npm:0.8.0":
version: 0.8.0
resolution: "@google/genai@npm:0.8.0"
"@google/genai@npm:^0.10.0":
version: 0.10.0
resolution: "@google/genai@npm:0.10.0"
dependencies:
google-auth-library: "npm:^9.14.2"
ws: "npm:^8.18.0"
checksum: 10c0/8a26a7dd1ab26aeeef5b5610612965ab271142460912c31b12f201cf6e00f5a4965910b195033992bdee1a7ee2b88c55f55d3a2727e09e4cd8d30ecbd0d655d0
languageName: node
linkType: hard
"@google/genai@patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch":
version: 0.8.0
resolution: "@google/genai@patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch::version=0.8.0&hash=e4f82c"
dependencies:
google-auth-library: "npm:^9.14.2"
ws: "npm:^8.18.0"
checksum: 10c0/b2082599cbdb9ffe18020461a2290ac60b0912f21106254bb618a1a1faa6e400caf16f92479fdb877928fd9670617ec0131f5787b4e3db2e5918daf2c7b87def
zod: "npm:^3.22.4"
zod-to-json-schema: "npm:^3.22.4"
checksum: 10c0/89a80eebdd053f2643fb977b939c3de1ece6889549f7d030b81a21dae0e0e0750233e453dafe665f8d337c642b4ba7a2f19b06bc270f3e2de8aeb803af1851d1
languageName: node
linkType: hard
@ -4315,7 +4307,7 @@ __metadata:
"@emotion/is-prop-valid": "npm:^1.3.1"
"@eslint-react/eslint-plugin": "npm:^1.36.1"
"@eslint/js": "npm:^9.22.0"
"@google/genai": "patch:@google/genai@npm%3A0.8.0#~/.yarn/patches/@google-genai-npm-0.8.0-450d0d9a7d.patch"
"@google/genai": "npm:^0.10.0"
"@hello-pangea/dnd": "npm:^16.6.0"
"@kangfenmao/keyv-storage": "npm:^0.1.0"
"@langchain/community": "npm:^0.3.36"
@ -18032,7 +18024,7 @@ __metadata:
languageName: node
linkType: hard
"zod-to-json-schema@npm:^3.22.3, zod-to-json-schema@npm:^3.22.5, zod-to-json-schema@npm:^3.24.1":
"zod-to-json-schema@npm:^3.22.3, zod-to-json-schema@npm:^3.22.4, zod-to-json-schema@npm:^3.22.5, zod-to-json-schema@npm:^3.24.1":
version: 3.24.5
resolution: "zod-to-json-schema@npm:3.24.5"
peerDependencies: