Add content policy checks and rate limiting

This commit is contained in:
Riccardo Giorato
2026-01-05 23:06:30 +01:00
parent 745b9d0ab5
commit 24c1f733b3
6 changed files with 117 additions and 48 deletions
+38 -19
View File
@@ -10,10 +10,12 @@ import {
getNextPageNumber,
getStoryWithPagesBySlug,
getLastPageImage,
deletePage,
} from "@/lib/db-actions";
import { freeTierRateLimit } from "@/lib/rate-limit";
import { uploadImageToS3 } from "@/lib/s3-upload";
import { buildComicPrompt } from "@/lib/prompt";
import { isContentPolicyViolation, getContentPolicyErrorMessage } from "@/lib/utils";
const NEW_MODEL = false;
@@ -63,25 +65,7 @@ export async function POST(request: NextRequest) {
return NextResponse.json({ error: "Unauthorized" }, { status: 403 });
}
// Apply rate limiting for free tier
const hasApiKey = request.headers.get("x-api-key");
if (!hasApiKey) {
const { success, reset } = await freeTierRateLimit.limit(userId);
if (!success) {
const resetDate = new Date(reset);
const timeUntilReset = Math.ceil(
(reset - Date.now()) / (1000 * 60 * 60 * 24)
);
return NextResponse.json(
{
error: `Free tier limit reached. You can generate 1 comic per week. Try again in ${timeUntilReset} day(s), or provide your own API key.`,
resetDate: resetDate.toISOString(),
isRateLimited: true,
},
{ status: 429 }
);
}
}
let page;
let pageNumber;
@@ -172,6 +156,30 @@ export async function POST(request: NextRequest) {
} catch (error) {
console.error("Together AI API error:", error);
// Clean up DB records if generation failed due to content policy
try {
if (error instanceof Error && error.message && error.message.includes("NO_IMAGE")) {
if (isRedraw) {
// For redraw, we don't delete the page, just don't update it
} else {
// For new page, delete the page that was created
await deletePage(page.id);
}
}
} catch (cleanupError) {
console.error("Error cleaning up DB on image generation failure:", cleanupError);
}
if (error instanceof Error && error.message && isContentPolicyViolation(error.message)) {
return NextResponse.json(
{
error: getContentPolicyErrorMessage(),
errorType: "content_policy",
},
{ status: 400 }
);
}
if (error instanceof Error && "status" in error) {
const status = (error as any).status;
if (status === 402) {
@@ -215,6 +223,17 @@ export async function POST(request: NextRequest) {
await updatePage(page.id, s3ImageUrl);
// Apply rate limiting for free tier after successful generation
const hasApiKey = request.headers.get("x-api-key");
if (!hasApiKey) {
try {
await freeTierRateLimit.limit(userId);
} catch (rateLimitError) {
console.error("Error applying rate limit after successful generation:", rateLimitError);
// Don't fail the request if rate limiting fails, just log it
}
}
return NextResponse.json({
imageUrl: s3ImageUrl,
pageId: page.id,
+38 -19
View File
@@ -10,11 +10,14 @@ import {
getStoryById,
getLastPageImage,
getStoryCharacterImages,
deletePage,
deleteStory,
} from "@/lib/db-actions";
import { freeTierRateLimit } from "@/lib/rate-limit";
import { COMIC_STYLES } from "@/lib/constants";
import { uploadImageToS3 } from "@/lib/s3-upload";
import { buildComicPrompt } from "@/lib/prompt";
import { isContentPolicyViolation, getContentPolicyErrorMessage } from "@/lib/utils";
const NEW_MODEL = false;
@@ -61,25 +64,6 @@ export async function POST(request: NextRequest) {
const isUsingFreeTier = !apiKey;
if (isUsingFreeTier) {
// Using free tier - apply rate limiting
const { success, reset } = await freeTierRateLimit.limit(userId);
if (!success) {
const resetDate = new Date(reset);
const timeUntilReset = Math.ceil(
(reset - Date.now()) / (1000 * 60 * 60 * 24)
); // days
return NextResponse.json(
{
error: `Free tier limit reached. You can generate 1 comic per week. Try again in ${timeUntilReset} day(s), or provide your own API key for unlimited access.`,
resetDate: resetDate.toISOString(),
isRateLimited: true,
},
{ status: 429 }
);
}
// Use default API key for free tier
finalApiKey = process.env.TOGETHER_API_KEY_DEFAULT;
if (!finalApiKey) {
@@ -92,6 +76,8 @@ export async function POST(request: NextRequest) {
}
}
let page;
let story;
let referenceImages: string[] = [];
@@ -251,6 +237,29 @@ Only return the JSON, no other text.`;
} catch (error) {
console.error("Together AI API error:", error);
// Clean up DB records if generation failed
try {
if (!storyId) {
// New story failed
await deleteStory(story!.id);
} else {
// Continuation failed
await deletePage(page.id);
}
} catch (cleanupError) {
console.error("Error cleaning up DB on image generation failure:", cleanupError);
}
if (error instanceof Error && error.message && isContentPolicyViolation(error.message)) {
return NextResponse.json(
{
error: getContentPolicyErrorMessage(),
errorType: "content_policy",
},
{ status: 400 }
);
}
if (error instanceof Error && "status" in error) {
const status = (error as any).status;
if (status === 402) {
@@ -334,6 +343,16 @@ Only return the JSON, no other text.`;
);
}
// Apply rate limiting for free tier after successful generation
if (isUsingFreeTier) {
try {
await freeTierRateLimit.limit(userId);
} catch (rateLimitError) {
console.error("Error applying rate limit after successful generation:", rateLimitError);
// Don't fail the request if rate limiting fails, just log it
}
}
const responseData = storyId
? { imageUrl: s3ImageUrl, pageId: page.id, pageNumber: page.pageNumber }
: {
+12 -5
View File
@@ -14,6 +14,7 @@ import { useToast } from "@/hooks/use-toast";
import { useKeyboardShortcut } from "@/hooks/use-keyboard-shortcut";
import { validateFileForUpload, generateFilePreview } from "@/lib/file-utils";
import { useS3Upload } from "next-s3-upload";
import { isContentPolicyViolation } from "@/lib/utils";
interface CharacterItem {
url: string;
@@ -268,16 +269,22 @@ export function GeneratePageModal({
});
} catch (error) {
console.error("Error generating page:", error);
const errorMessage =
error instanceof Error
? error.message
: "Failed to generate page. Please try again.";
let title = "Generation failed";
if (isContentPolicyViolation(errorMessage)) {
title = "Content policy violation";
}
toast({
title: "Generation failed",
description:
error instanceof Error
? error.message
: "Failed to generate page. Please try again.",
title,
description: errorMessage,
variant: "destructive",
duration: 4000,
});
setIsGenerating(false);
throw error; // Re-throw so the parent handler knows generation failed
}
};
+11 -5
View File
@@ -10,6 +10,7 @@ import { useAuth, SignInButton } from "@clerk/nextjs";
import { COMIC_STYLES } from "@/lib/constants";
import { useKeyboardShortcut } from "@/hooks/use-keyboard-shortcut";
import { useApiKey } from "@/hooks/use-api-key";
import { isContentPolicyViolation } from "@/lib/utils";
interface ComicCreationFormProps {
prompt: string;
@@ -171,12 +172,17 @@ export function ComicCreationForm({
router.push(`/story/${result.storySlug}`);
} catch (error) {
console.error("Error creating comic:", error);
const errorMessage =
error instanceof Error
? error.message
: "Failed to create comic. Please try again.";
let title = "Creation failed";
if (isContentPolicyViolation(errorMessage)) {
title = "Content policy violation";
}
toast({
title: "Creation failed",
description:
error instanceof Error
? error.message
: "Failed to create comic. Please try again.",
title,
description: errorMessage,
variant: "destructive",
duration: 4000,
});
+4
View File
@@ -145,4 +145,8 @@ export async function getNextPageNumber(storyId: string): Promise<number> {
export async function deletePage(pageId: string): Promise<void> {
await db.delete(pages).where(eq(pages.id, pageId));
}
export async function deleteStory(storyId: string): Promise<void> {
await db.delete(stories).where(eq(stories.id, storyId));
}
+14
View File
@@ -7,3 +7,17 @@ export function cn(...inputs: ClassValue[]) {
export const TOGETHER_LINK =
"https://togetherai.link/?utm_source=make-comics&utm_medium=referral&utm_campaign=example-app";
export function isContentPolicyViolation(errorMessage: string): boolean {
return (
errorMessage.includes("content policy") ||
errorMessage.includes("Invalid content detected") ||
errorMessage.includes("content moderation") ||
errorMessage.includes("flagged and rejected") ||
errorMessage.includes("NO_IMAGE")
);
}
export function getContentPolicyErrorMessage(): string {
return "Unable to generate image due to content policy. Please try a different prompt.";
}