scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-10-06 02:47:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat: adding citation for streams (generateContextStream)

Browse source
This commit is contained in:
Ton Hoang Nguyen (Bill) 2025-07-08 16:17:12 +01:00
parent bf0161831e
commit 63c7b25f0d
1 changed files with 28 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -93,11 +93,19 @@ 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
if (chunk.candidates && chunk.candidates.length > 0) {
const candidate = chunk.candidates[0]
if (candidate.groundingMetadata) {
pendingGroundingMetadata = candidate.groundingMetadata
}
if (candidate.content && candidate.content.parts) {
for (const part of candidate.content.parts) {
if (part.thought) {
@ -108,7 +116,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
} else {
// This is regular content
if (part.text) {
yield { type: "text", text: part.text }
accumulatedText += part.text
if (!hasGroundingEnabled) {
yield { type: "text", text: part.text }
}
}
}
}
@ -117,7 +129,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
// Fallback to the original text property if no candidates structure
else if (chunk.text) {
yield { type: "text", text: chunk.text }
accumulatedText += chunk.text
if (!hasGroundingEnabled) {
yield { type: "text", text: chunk.text }
}
}
if (chunk.usageMetadata) {
@ -125,6 +141,16 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
}
}
if (hasGroundingEnabled && accumulatedText) {
let finalText = accumulatedText
if (pendingGroundingMetadata) {
finalText = this.processGroundingCitations(accumulatedText, pendingGroundingMetadata)
}
yield { type: "text", text: finalText }
}
if (lastUsageMetadata) {
const inputTokens = lastUsageMetadata.promptTokenCount ?? 0
const outputTokens = lastUsageMetadata.candidatesTokenCount ?? 0