sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(ai): preserve google providers raw stop reason

Browse source
This commit is contained in:
David Brailovsky 2026-07-29 11:31:53 +00:00
parent 027a584790
commit 23cb385b67
4 changed files with 117 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

6
packages/ai/src/api/google-generative-ai.ts
View file

@ -212,6 +212,7 @@ export const stream: StreamFunction<"google-generative-ai", GoogleOptions> = (
}
if (candidate?.finishReason) {
output.rawStopReason = candidate.finishReason;
output.stopReason = mapStopReason(candidate.finishReason);
if (output.content.some((b) => b.type === "toolCall")) {
output.stopReason = "toolUse";
@ -266,7 +267,10 @@ export const stream: StreamFunction<"google-generative-ai", GoogleOptions> = (
throw new Error("Google stream ended without a finish reason");
}
if (output.stopReason === "aborted" || output.stopReason === "error") {
throw new Error("An unknown error occurred");
const errorMessage = output.rawStopReason
? `Gemini stopped with ${output.rawStopReason}`
: "An unknown error occurred";
throw new Error(errorMessage);
}
stream.push({ type: "done", reason: output.stopReason, message: output });

6
packages/ai/src/api/google-vertex.ts
View file

@ -229,6 +229,7 @@ export const stream: StreamFunction<"google-vertex", GoogleVertexOptions> = (
}
if (candidate?.finishReason) {
output.rawStopReason = candidate.finishReason;
output.stopReason = mapStopReason(candidate.finishReason);
if (output.content.some((b) => b.type === "toolCall")) {
output.stopReason = "toolUse";
@ -283,7 +284,10 @@ export const stream: StreamFunction<"google-vertex", GoogleVertexOptions> = (
throw new Error("Google Vertex stream ended without a finish reason");
}
if (output.stopReason === "aborted" || output.stopReason === "error") {
throw new Error("An unknown error occurred");
const errorMessage = output.rawStopReason
? `Gemini stopped with ${output.rawStopReason}`
: "An unknown error occurred";
throw new Error(errorMessage);
}
stream.push({ type: "done", reason: output.stopReason, message: output });

1
packages/ai/src/types.ts
View file

@ -408,6 +408,7 @@ export interface AssistantMessage {
usage: Usage;
stopReason: StopReason;
errorMessage?: string;
rawStopReason?: string;
timestamp: number; // Unix timestamp in milliseconds
}

106
packages/ai/test/google-raw-stop-reason.test.ts Normal file
View file

@ -0,0 +1,106 @@
import { describe, expect, it, vi } from "vitest";
const googleGenAiMock = vi.hoisted(() => ({
finishReason: "MALFORMED_FUNCTION_CALL",
}));
vi.mock("@google/genai", () => {
class GoogleGenAI {
models = {
generateContentStream: async function* () {
yield {
responseId: "google-response-id",
candidates: [
{
finishReason: googleGenAiMock.finishReason,
},
],
usageMetadata: {
promptTokenCount: 1,
candidatesTokenCount: 0,
totalTokenCount: 1,
},
};
},
};
}
return {
FinishReason: {
STOP: "STOP",
MAX_TOKENS: "MAX_TOKENS",
BLOCKLIST: "BLOCKLIST",
PROHIBITED_CONTENT: "PROHIBITED_CONTENT",
SPII: "SPII",
SAFETY: "SAFETY",
IMAGE_SAFETY: "IMAGE_SAFETY",
IMAGE_PROHIBITED_CONTENT: "IMAGE_PROHIBITED_CONTENT",
IMAGE_RECITATION: "IMAGE_RECITATION",
IMAGE_OTHER: "IMAGE_OTHER",
RECITATION: "RECITATION",
FINISH_REASON_UNSPECIFIED: "FINISH_REASON_UNSPECIFIED",
OTHER: "OTHER",
LANGUAGE: "LANGUAGE",
MALFORMED_FUNCTION_CALL: "MALFORMED_FUNCTION_CALL",
UNEXPECTED_TOOL_CALL: "UNEXPECTED_TOOL_CALL",
NO_IMAGE: "NO_IMAGE",
},
FunctionCallingConfigMode: {
AUTO: "AUTO",
NONE: "NONE",
ANY: "ANY",
VALIDATED: "VALIDATED",
},
GoogleGenAI,
ResourceScope: {
COLLECTION: "COLLECTION",
},
ThinkingLevel: {
THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED",
MINIMAL: "MINIMAL",
LOW: "LOW",
MEDIUM: "MEDIUM",
HIGH: "HIGH",
},
};
});
import { stream as streamGoogleGenerativeAi } from "../src/api/google-generative-ai.ts";
import { stream as streamGoogleVertex } from "../src/api/google-vertex.ts";
import { getModel } from "../src/compat.ts";
import type { Context } from "../src/types.ts";
const context: Context = {
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
};
describe("Google raw stop reasons", () => {
it("preserves raw Gemini finish reasons for Google Generative AI errors", async () => {
googleGenAiMock.finishReason = "MALFORMED_FUNCTION_CALL";
const stream = streamGoogleGenerativeAi(getModel("google", "gemini-2.5-flash"), context, {
apiKey: "test-api-key",
});
const message = await stream.result();
expect(message.stopReason).toBe("error");
expect(message.rawStopReason).toBe("MALFORMED_FUNCTION_CALL");
expect(message.errorMessage).toBe("Gemini stopped with MALFORMED_FUNCTION_CALL");
});
it("preserves raw Gemini finish reasons for Google Vertex errors", async () => {
googleGenAiMock.finishReason = "SAFETY";
const stream = streamGoogleVertex(getModel("google-vertex", "gemini-3-flash-preview"), context, {
project: "test-project",
location: "us-central1",
});
const message = await stream.result();
expect(message.stopReason).toBe("error");
expect(message.rawStopReason).toBe("SAFETY");
expect(message.errorMessage).toBe("Gemini stopped with SAFETY");
});
});