resolves server-side error when a Gemini prompt is blocked due to safety

This commit is contained in:
nai-degen 2024-12-07 00:54:13 -06:00
parent 6d54cbc785
commit 28447d0811
2 changed files with 25 additions and 3 deletions

View File

@ -264,7 +264,15 @@ export function getCompletionFromBody(req: Request, body: Record<string, any>) {
if ("choices" in body) {
return body.choices[0].message.content;
}
return body.candidates[0].content.parts[0].text;
const text = body.candidates[0].content?.parts?.[0]?.text;
if (!text) {
req.log.warn(
{ body: JSON.stringify(body) },
"Received empty Google AI text completion"
);
return "";
}
return text;
case "openai-image":
return body.data?.map((item: any) => item.url).join("\n");
default:

View File

@ -9,7 +9,7 @@ const log = logger.child({
type GoogleAIStreamEvent = {
candidates: {
content: { parts: { text: string }[]; role: string };
content?: { parts?: { text: string }[]; role: string };
finishReason?: "STOP" | "MAX_TOKENS" | "SAFETY" | "RECITATION" | "OTHER";
index: number;
tokenCount?: number;
@ -34,9 +34,15 @@ export const googleAIToOpenAI: StreamingCompletionTransformer = (params) => {
return { position: -1 };
}
const parts = completionEvent.candidates[0].content.parts;
const parts = completionEvent.candidates[0].content?.parts || [];
let content = parts[0]?.text ?? "";
if (isSafetyStop(completionEvent)) {
content = `[Proxy Warning] Gemini safety filter triggered: ${JSON.stringify(
completionEvent.candidates[0].safetyRatings
)}`;
}
// If this is the first chunk, try stripping speaker names from the response
// e.g. "John: Hello" -> "Hello"
if (index === 0) {
@ -60,6 +66,14 @@ export const googleAIToOpenAI: StreamingCompletionTransformer = (params) => {
return { position: -1, event: newEvent };
};
function isSafetyStop(completion: GoogleAIStreamEvent) {
const isSafetyStop = ["SAFETY", "OTHER"].includes(
completion.candidates[0].finishReason ?? ""
);
const hasNoContent = completion.candidates[0].content?.parts?.length === 0;
return isSafetyStop && hasNoContent;
}
function asCompletion(event: ServerSentEvent): GoogleAIStreamEvent | null {
try {
const parsed = JSON.parse(event.data) as GoogleAIStreamEvent;