Skip to content

Commit 99d9cd0

Browse files
uinstinctuinstinct
committed
do not split lines again in processgeminiresponse
1 parent baad1e2 commit 99d9cd0

File tree

3 files changed

+87
-123
lines changed

3 files changed

+87
-123
lines changed

‎core/llm/llms/Gemini.ts‎

Lines changed: 45 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -312,75 +312,57 @@ class Gemini extends BaseLLM {
312312
}
313313

314314
public async *processGeminiResponse(
315-
stream: AsyncIterable<string>,
315+
response: Response,
316316
): AsyncGenerator<ChatMessage> {
317-
let buffer = "";
318-
for await (const chunk of stream) {
319-
buffer += chunk;
320-
321-
const parts = buffer.split("\n,");
322-
323-
let foundIncomplete = false;
324-
for (let i = 0; i < parts.length; i++) {
325-
const part = parts[i];
326-
let data: GeminiChatResponse;
327-
try {
328-
data = JSON.parse(part) as GeminiChatResponse;
329-
} catch (e) {
330-
foundIncomplete = true;
331-
continue; // yo!
332-
}
317+
for await (const chunk of streamSse(response)) {
318+
let data: GeminiChatResponse;
319+
try {
320+
data = JSON.parse(chunk) as GeminiChatResponse;
321+
} catch (e) {
322+
continue;
323+
}
333324

334-
if ("error" in data) {
335-
throw new Error(data.error.message);
336-
}
325+
if ("error" in data) {
326+
throw new Error(data.error.message);
327+
}
337328

338-
// In case of max tokens reached, gemini will sometimes return content with no parts, even though that doesn't match the API spec
339-
const contentParts = data?.candidates?.[0]?.content?.parts;
340-
if (contentParts) {
341-
const textParts: MessagePart[] = [];
342-
const toolCalls: ToolCallDelta[] = [];
343-
344-
for (const part of contentParts) {
345-
if ("text" in part) {
346-
textParts.push({ type: "text", text: part.text });
347-
} else if ("functionCall" in part) {
348-
toolCalls.push({
349-
type: "function",
350-
id: part.functionCall.id ?? uuidv4(),
351-
function: {
352-
name: part.functionCall.name,
353-
arguments:
354-
typeof part.functionCall.args === "string"
355-
? part.functionCall.args
356-
: JSON.stringify(part.functionCall.args),
357-
},
358-
});
359-
} else {
360-
// Note: function responses shouldn't be streamed, images not supported
361-
console.warn("Unsupported gemini part type received", part);
362-
}
329+
const contentParts = data?.candidates?.[0]?.content?.parts;
330+
if (contentParts) {
331+
const textParts: MessagePart[] = [];
332+
const toolCalls: ToolCallDelta[] = [];
333+
334+
for (const part of contentParts) {
335+
if ("text" in part) {
336+
textParts.push({ type: "text", text: part.text });
337+
} else if ("functionCall" in part) {
338+
toolCalls.push({
339+
type: "function",
340+
id: part.functionCall.id ?? uuidv4(),
341+
function: {
342+
name: part.functionCall.name,
343+
arguments:
344+
typeof part.functionCall.args === "string"
345+
? part.functionCall.args
346+
: JSON.stringify(part.functionCall.args),
347+
},
348+
});
349+
} else {
350+
console.warn("Unsupported gemini part type received", part);
363351
}
352+
}
364353

365-
const assistantMessage: AssistantChatMessage = {
366-
role: "assistant",
367-
content: textParts.length ? textParts : "",
368-
};
369-
if (toolCalls.length > 0) {
370-
assistantMessage.toolCalls = toolCalls;
371-
}
372-
if (textParts.length || toolCalls.length) {
373-
yield assistantMessage;
374-
}
375-
} else {
376-
// Handle the case where the expected data structure is not found
377-
console.warn("Unexpected response format:", data);
354+
const assistantMessage: AssistantChatMessage = {
355+
role: "assistant",
356+
content: textParts.length ? textParts : "",
357+
};
358+
if (toolCalls.length > 0) {
359+
assistantMessage.toolCalls = toolCalls;
360+
}
361+
if (textParts.length || toolCalls.length) {
362+
yield assistantMessage;
378363
}
379-
}
380-
if (foundIncomplete) {
381-
buffer = parts[parts.length - 1];
382364
} else {
383-
buffer = "";
365+
console.warn("Unexpected response format:", data);
384366
}
385367
}
386368
}
@@ -406,7 +388,7 @@ class Gemini extends BaseLLM {
406388
signal,
407389
});
408390

409-
for await (const chunk of this.processGeminiResponse(streamSse(response))) {
391+
for await (const chunk of this.processGeminiResponse(response)) {
410392
yield chunk;
411393
}
412394
}

‎core/llm/llms/VertexAI.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ class VertexAI extends BaseLLM {
287287
body: JSON.stringify(body),
288288
signal,
289289
});
290-
yield* this.geminiInstance.processGeminiResponse(streamSse(response));
290+
yield* this.geminiInstance.processGeminiResponse(response);
291291
}
292292

293293
private async *streamChatBison(

‎packages/openai-adapters/src/apis/Gemini.ts‎

Lines changed: 41 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -284,76 +284,58 @@ export class GeminiApi implements BaseLlmApi {
284284
}
285285

286286
async *handleStreamResponse(response: any, model: string) {
287-
let buffer = "";
288287
let usage: UsageInfo | undefined = undefined;
289288
for await (const chunk of streamSse(response as any)) {
290-
buffer += chunk;
291-
292-
const parts = buffer.split("\n,");
293-
294-
let foundIncomplete = false;
295-
for (let i = 0; i < parts.length; i++) {
296-
const part = parts[i];
297-
let data;
298-
try {
299-
data = JSON.parse(part);
300-
} catch (e) {
301-
foundIncomplete = true;
302-
continue; // yo!
303-
}
304-
if (data.error) {
305-
throw new Error(data.error.message);
306-
}
289+
let data;
290+
try {
291+
data = JSON.parse(chunk);
292+
} catch (e) {
293+
continue;
294+
}
295+
if (data.error) {
296+
throw new Error(data.error.message);
297+
}
307298

308-
// Check for usage metadata
309-
if (data.usageMetadata) {
310-
usage = {
311-
prompt_tokens: data.usageMetadata.promptTokenCount || 0,
312-
completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
313-
total_tokens: data.usageMetadata.totalTokenCount || 0,
314-
};
315-
}
299+
if (data.usageMetadata) {
300+
usage = {
301+
prompt_tokens: data.usageMetadata.promptTokenCount || 0,
302+
completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
303+
total_tokens: data.usageMetadata.totalTokenCount || 0,
304+
};
305+
}
316306

317-
// In case of max tokens reached, gemini will sometimes return content with no parts, even though that doesn't match the API spec
318-
const contentParts = data?.candidates?.[0]?.content?.parts;
319-
if (contentParts) {
320-
for (const part of contentParts) {
321-
if ("text" in part) {
322-
yield chatChunk({
323-
content: part.text,
324-
model,
325-
});
326-
} else if ("functionCall" in part) {
327-
yield chatChunkFromDelta({
328-
model,
329-
delta: {
330-
tool_calls: [
331-
{
332-
index: 0,
333-
id: part.functionCall.id ?? uuidv4(),
334-
type: "function",
335-
function: {
336-
name: part.functionCall.name,
337-
arguments: JSON.stringify(part.functionCall.args),
338-
},
307+
const contentParts = data?.candidates?.[0]?.content?.parts;
308+
if (contentParts) {
309+
for (const part of contentParts) {
310+
if ("text" in part) {
311+
yield chatChunk({
312+
content: part.text,
313+
model,
314+
});
315+
} else if ("functionCall" in part) {
316+
yield chatChunkFromDelta({
317+
model,
318+
delta: {
319+
tool_calls: [
320+
{
321+
index: 0,
322+
id: part.functionCall.id ?? uuidv4(),
323+
type: "function",
324+
function: {
325+
name: part.functionCall.name,
326+
arguments: JSON.stringify(part.functionCall.args),
339327
},
340-
],
341-
},
342-
});
343-
}
328+
},
329+
],
330+
},
331+
});
344332
}
345-
} else {
346-
console.warn("Unexpected response format:", data);
347333
}
348-
}
349-
if (foundIncomplete) {
350-
buffer = parts[parts.length - 1];
351334
} else {
352-
buffer = "";
335+
console.warn("Unexpected response format:", data);
353336
}
354337
}
355338

356-
// Emit usage at the end if we have it
357339
if (usage) {
358340
yield usageChatChunk({
359341
model,

0 commit comments

Comments
 (0)