scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-09-07 08:26:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: adding sources at the end of the stream to preserve

Browse source
This commit is contained in:
Ton Hoang Nguyen (Bill) 2025-07-14 21:21:56 +01:00
parent 22eb360bc2
commit 7e9d2527f5
1 changed files with 25 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
src/api/providers/gemini.ts
View file

@ -93,9 +93,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
const result = await this.client.models.generateContentStream(params)
let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
let accumulatedText = ""
let pendingGroundingMetadata: GroundingMetadata | undefined
let hasGroundingEnabled = this.options.enableGrounding
for await (const chunk of result) {
// Process candidates and their parts to separate thoughts from content
@ -116,11 +114,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
} else {
// This is regular content
if (part.text) {
accumulatedText += part.text
if (!hasGroundingEnabled) {
yield { type: "text", text: part.text }
}
yield { type: "text", text: part.text }
}
}
}
@ -129,11 +123,7 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
// Fallback to the original text property if no candidates structure
else if (chunk.text) {
accumulatedText += chunk.text
if (!hasGroundingEnabled) {
yield { type: "text", text: chunk.text }
}
yield { type: "text", text: chunk.text }
}
if (chunk.usageMetadata) {
@ -141,14 +131,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
}
}
if (hasGroundingEnabled && accumulatedText) {
let finalText = accumulatedText
if (pendingGroundingMetadata) {
finalText = this.processGroundingCitations(accumulatedText, pendingGroundingMetadata)
if (pendingGroundingMetadata) {
const citations = this.extractCitationsOnly(pendingGroundingMetadata)
if (citations) {
yield { type: "text", text: `\n\nSources: ${citations}` }
}
yield { type: "text", text: finalText }
}
if (lastUsageMetadata) {
@ -188,39 +175,28 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
return { id: id.endsWith(":thinking") ? id.replace(":thinking", "") : id, info, ...params }
}
private processGroundingCitations(text: string, groundingMetadata?: GroundingMetadata): string {
const supports = groundingMetadata?.groundingSupports
private extractCitationsOnly(groundingMetadata?: GroundingMetadata): string | null {
const chunks = groundingMetadata?.groundingChunks
if (!supports || !chunks) {
return text
if (!chunks) {
return null
}
const sortedSupports = [...supports].sort((a, b) => (b.segment?.endIndex ?? 0) - (a.segment?.endIndex ?? 0))
const citationLinks = chunks
.map((chunk, i) => {
const uri = chunk.web?.uri
if (uri) {
return `[${i + 1}](${uri})`
}
return null
})
.filter((link): link is string => link !== null)
for (const support of sortedSupports) {
const endIndex = support.segment?.endIndex
if (endIndex === undefined || !support.groundingChunkIndices?.length) {
continue
}
const citationLinks = support.groundingChunkIndices
.map((i) => {
const uri = chunks[i]?.web?.uri
if (uri) {
return `[${i + 1}](${uri})`
}
return null
})
.filter(Boolean)
if (citationLinks.length > 0) {
const citationString = citationLinks.join(", ")
text = text.slice(0, endIndex) + citationString + text.slice(endIndex)
}
if (citationLinks.length > 0) {
return citationLinks.join(", ")
}
return text
return null
}
async completePrompt(prompt: string): Promise<string> {
@ -256,7 +232,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
const candidate = result.candidates?.[0]
if (candidate?.groundingMetadata) {
text = this.processGroundingCitations(text, candidate.groundingMetadata)
const citations = this.extractCitationsOnly(candidate.groundingMetadata)
if (citations) {
text += `\n\nSources: ${citations}`
}
}
return text