Skip to content

Commit d3fb59c

Browse files
0xGingi0xGingi
committed
update image gen to fix some backend bug fixes
1 parent e7dda46 commit d3fb59c

3 files changed

Lines changed: 141 additions & 64 deletions

File tree

‎src/routes/api/db/messages/+server.ts‎

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,51 @@ import {
1313
import { getConversationById } from '$lib/db/queries/conversations';
1414
import { getAuthenticatedUserId } from '$lib/backend/auth-utils';
1515

16+
function extractMarkdownImages(content: string): Array<{ url: string; fileName?: string }> {
17+
const images: Array<{ url: string; fileName?: string }> = [];
18+
const regex = /!\[([^\]]*)\]\(([^)]+)\)/g;
19+
for (const match of content.matchAll(regex)) {
20+
const alt = match[1]?.trim();
21+
const rawUrl = match[2]?.trim();
22+
if (!rawUrl) continue;
23+
const url = rawUrl.split(/\s+/)[0]?.replace(/^<|>$/g, '');
24+
if (!url) continue;
25+
images.push({ url, fileName: alt || undefined });
26+
}
27+
return images;
28+
}
29+
30+
function extractStorageId(url: string): string | null {
31+
const match = url.match(/\/api\/storage\/([a-zA-Z0-9_-]+)/);
32+
return match ? match[1] : null;
33+
}
34+
35+
function withInlineImages<
36+
T extends {
37+
id?: string;
38+
content: string;
39+
images?: Array<{ url: string; storage_id: string; fileName?: string }>;
40+
}
41+
>(
42+
messages: T[]
43+
): T[] {
44+
return messages.map((message) => {
45+
if (message.images && message.images.length > 0) return message;
46+
if (!message.content) return message;
47+
48+
const inlineImages = extractMarkdownImages(message.content);
49+
if (inlineImages.length === 0) return message;
50+
51+
const images = inlineImages.map((image, index) => ({
52+
url: image.url,
53+
storage_id: extractStorageId(image.url) ?? `${message.id ?? 'inline'}-${index}`,
54+
fileName: image.fileName,
55+
}));
56+
57+
return { ...message, images };
58+
});
59+
}
60+
1661
// GET - get messages for a conversation
1762
export const GET: RequestHandler = async ({ request, url }) => {
1863
const conversationId = url.searchParams.get('conversationId');
@@ -27,12 +72,12 @@ export const GET: RequestHandler = async ({ request, url }) => {
2772
if (!messages) {
2873
return error(404, 'Conversation not found or not public');
2974
}
30-
return json(messages);
75+
return json(withInlineImages(messages));
3176
}
3277

3378
const userId = await getAuthenticatedUserId(request);
3479
const messages = await getConversationMessages(conversationId, userId);
35-
return json(messages);
80+
return json(withInlineImages(messages));
3681
};
3782

3883
// POST - create or update message

‎src/routes/api/generate-message/+server.ts‎

Lines changed: 62 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1878,25 +1878,9 @@ export const POST: RequestHandler = async ({ request }) => {
18781878
const imageCost = response.cost ?? 0;
18791879
log(`Image generation cost: $${imageCost}`, startTime);
18801880

1881-
const image = response.data?.[0];
1882-
if (!image?.b64_json && !image?.url) {
1883-
throw new Error('No image data returned details: ' + JSON.stringify(image));
1884-
}
1885-
1886-
let buffer: Buffer;
1887-
let mimeType = 'image/png';
1888-
1889-
if (image.b64_json) {
1890-
buffer = Buffer.from(image.b64_json, 'base64');
1891-
} else if (image.url) {
1892-
// Fallback download
1893-
const imgRes = await fetch(image.url);
1894-
const arrayBuffer = await imgRes.arrayBuffer();
1895-
buffer = Buffer.from(arrayBuffer);
1896-
const contentType = imgRes.headers.get('content-type');
1897-
if (contentType) mimeType = contentType;
1898-
} else {
1899-
throw new Error('No image data');
1881+
const responseImages = response.data ?? [];
1882+
if (!Array.isArray(responseImages) || responseImages.length === 0) {
1883+
throw new Error('No image data returned details: ' + JSON.stringify(response.data));
19001884
}
19011885

19021886
// Ensure upload dir exists
@@ -1905,37 +1889,76 @@ export const POST: RequestHandler = async ({ request }) => {
19051889
mkdirSync(UPLOAD_DIR, { recursive: true });
19061890
}
19071891

1908-
const storageId = generateId();
1909-
const extension = (mimeType.split('/')[1] || 'png').replace(/[^a-zA-Z0-9]/g, '') || 'png';
1910-
const filename = `${storageId}.${extension}`;
1911-
const filepath = join(UPLOAD_DIR, filename);
1892+
const generatedImages: Array<{ url: string; storage_id: string; fileName?: string }> = [];
19121893

1913-
// Prevent path traversal
1914-
if (!resolve(filepath).startsWith(resolve(UPLOAD_DIR))) {
1915-
throw new Error('Invalid file path');
1916-
}
1894+
for (const [index, image] of responseImages.entries()) {
1895+
if (!image?.b64_json && !image?.url) {
1896+
console.warn('Skipping image without data:', image);
1897+
continue;
1898+
}
19171899

1918-
writeFileSync(filepath, buffer);
1900+
let buffer: Buffer;
1901+
let mimeType = 'image/png';
1902+
1903+
if (image.b64_json) {
1904+
buffer = Buffer.from(image.b64_json, 'base64');
1905+
} else if (image.url) {
1906+
// Fallback download
1907+
const imgRes = await fetch(image.url);
1908+
const arrayBuffer = await imgRes.arrayBuffer();
1909+
buffer = Buffer.from(arrayBuffer);
1910+
const contentType = imgRes.headers.get('content-type');
1911+
if (contentType) mimeType = contentType;
1912+
} else {
1913+
continue;
1914+
}
19191915

1920-
await db.insert(storage).values({
1921-
id: storageId,
1922-
userId,
1923-
filename,
1924-
mimeType,
1925-
size: buffer.byteLength,
1926-
path: filepath,
1927-
createdAt: new Date(),
1928-
});
1916+
const storageId = generateId();
1917+
const extension =
1918+
(mimeType.split('/')[1] || 'png').replace(/[^a-zA-Z0-9]/g, '') || 'png';
1919+
const filename = `${storageId}.${extension}`;
1920+
const filepath = join(UPLOAD_DIR, filename);
1921+
1922+
// Prevent path traversal
1923+
if (!resolve(filepath).startsWith(resolve(UPLOAD_DIR))) {
1924+
throw new Error('Invalid file path');
1925+
}
1926+
1927+
writeFileSync(filepath, buffer);
1928+
1929+
await db.insert(storage).values({
1930+
id: storageId,
1931+
userId,
1932+
filename,
1933+
mimeType,
1934+
size: buffer.byteLength,
1935+
path: filepath,
1936+
createdAt: new Date(),
1937+
});
1938+
1939+
const imageUrl = `/api/storage/${storageId}`;
1940+
generatedImages.push({
1941+
url: imageUrl,
1942+
storage_id: storageId,
1943+
fileName: `generated-image-${index + 1}.${extension}`,
1944+
});
1945+
}
1946+
1947+
if (generatedImages.length === 0) {
1948+
throw new Error('No valid image data returned from API');
1949+
}
19291950

1930-
const imageUrl = `/api/storage/${storageId}`;
1931-
const markdownContent = `![Generated Image](${imageUrl})`;
1951+
const markdownContent = generatedImages
1952+
.map((img, idx) => `![Generated Image ${idx + 1}](${img.url})`)
1953+
.join('\n\n');
19321954

19331955
await db
19341956
.update(messages)
19351957
.set({
19361958
content: markdownContent,
19371959
contentHtml: null,
19381960
tokenCount: 0,
1961+
images: generatedImages,
19391962
costUsd: imageCost,
19401963
})
19411964
.where(eq(messages.id, assistantMessageId));

‎src/routes/chat/[id]/message.svelte‎

Lines changed: 32 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,15 @@
5555
5656
let { message, childMessageId }: Props = $props();
5757
58+
const safeContent = $derived.by(() => (typeof message.content === 'string' ? message.content : ''));
59+
const safeImages = $derived.by(() =>
60+
Array.isArray(message.images)
61+
? message.images.filter(
62+
(image) => image && typeof image.url === 'string' && typeof image.storage_id === 'string'
63+
)
64+
: []
65+
);
66+
5867
let imageModal = $state<{ open: boolean; imageUrl: string; fileName: string }>({
5968
open: false,
6069
imageUrl: '',
@@ -127,8 +136,8 @@
127136
128137
// Detect Video URL in content
129138
const videoUrl = $derived.by(() => {
130-
if (!message.content) return null;
131-
const match = message.content.match(/\[Video Result\]\((.*?)\)/);
139+
if (!safeContent) return null;
140+
const match = safeContent.match(/\[Video Result\]\((.*?)\)/);
132141
return match ? match[1] : null;
133142
});
134143
@@ -205,7 +214,7 @@
205214
});
206215
207216
function startEditing() {
208-
editedContent = message.content;
217+
editedContent = safeContent;
209218
isEditing = true;
210219
}
211220
@@ -246,7 +255,7 @@
246255
}
247256
248257
async function saveMessage() {
249-
if (editedContent === message.content) {
258+
if (editedContent === safeContent) {
250259
cancelEditing();
251260
return;
252261
}
@@ -324,7 +333,7 @@
324333
}
325334
</script>
326335

327-
{#if message.role !== 'system' && !(message.role === 'assistant' && message.content.length === 0 && message.reasoning?.length === 0 && !message.error)}
336+
{#if message.role !== 'system' && !(message.role === 'assistant' && safeContent.length === 0 && message.reasoning?.length === 0 && !message.error)}
328337
<div
329338
class={cn('group flex flex-col gap-1', { 'max-w-[80%] self-end ': message.role === 'user' })}
330339
{@attach (node) => {
@@ -349,9 +358,9 @@
349358
</div>
350359
{/if}
351360

352-
{#if message.images && message.images.length > 0}
361+
{#if safeImages.length > 0}
353362
<div class="mb-2 flex flex-wrap gap-2">
354-
{#each message.images as image (image.storage_id)}
363+
{#each safeImages as image (image.storage_id)}
355364
<button
356365
type="button"
357366
onclick={() => openImageModal(image.url, image.fileName || 'image')}
@@ -382,7 +391,7 @@
382391
<ChevronRightIcon
383392
class={cn('size-4 transition-transform duration-200', { 'rotate-90': showReasoning })}
384393
/>
385-
{#if message.content.length === 0}
394+
{#if safeContent.length === 0}
386395
<ShinyText class="font-medium">Thinking...</ShinyText>
387396
{:else}
388397
<span class="font-medium">Reasoning</span>
@@ -444,7 +453,7 @@
444453
{@html sanitizeHtml(message.contentHtml)}
445454
{:else}
446455
<svelte:boundary>
447-
<MarkdownRenderer content={message.content} />
456+
<MarkdownRenderer content={safeContent} />
448457

449458
{#snippet failed(error)}
450459
<div class="text-destructive">
@@ -537,7 +546,7 @@
537546
{message.role === 'user' ? 'Branch and regenerate message' : 'Branch off this message'}
538547
</Tooltip>
539548

540-
{#if message.role === 'assistant' && message.content.length > 0}
549+
{#if message.role === 'assistant' && safeContent.length > 0}
541550
<Tooltip>
542551
{#snippet trigger(tooltip)}
543552
<Button
@@ -548,7 +557,7 @@
548557
if (audioPlayer.isPlaying && audioPlayer.currentMessageId === message.id) {
549558
audioPlayer.stop();
550559
} else {
551-
audioPlayer.play(message.content, message.id);
560+
audioPlayer.play(safeContent, message.id);
552561
}
553562
}}
554563
{...tooltip.trigger}
@@ -567,7 +576,7 @@
567576
: 'Read aloud'}
568577
</Tooltip>
569578
{/if}
570-
{#if message.role === 'assistant' && message.content.length > 0 && !message.error}
579+
{#if message.role === 'assistant' && safeContent.length > 0 && !message.error}
571580
<Tooltip>
572581
{#snippet trigger(tooltip)}
573582
<Button
@@ -589,15 +598,15 @@
589598
</Tooltip>
590599
{/if}
591600

592-
{#if message.content.length > 0}
593-
<Tooltip>
594-
{#snippet trigger(tooltip)}
595-
<CopyButton
596-
class={cn('order-1 size-7', { 'order-2': message.role === 'user' })}
597-
text={message.content}
598-
onclick={() => logInteraction('copy')}
599-
{...tooltip.trigger}
600-
/>
601+
{#if safeContent.length > 0}
602+
<Tooltip>
603+
{#snippet trigger(tooltip)}
604+
<CopyButton
605+
class={cn('order-1 size-7', { 'order-2': message.role === 'user' })}
606+
text={safeContent}
607+
onclick={() => logInteraction('copy')}
608+
{...tooltip.trigger}
609+
/>
601610
{/snippet}
602611
Copy
603612
</Tooltip>
@@ -685,14 +694,14 @@
685694
</div>
686695
{/if}
687696
</div>
688-
{#if message.role === 'assistant' && message.content.length > 0 && !message.error}
697+
{#if message.role === 'assistant' && safeContent.length > 0 && !message.error}
689698
<div class="mt-2">
690699
<MessageRating messageId={message.id} onRate={handleRating} />
691700
</div>
692701
{/if}
693702
</div>
694703

695-
{#if message.images && message.images.length > 0}
704+
{#if safeImages.length > 0}
696705
<ImageModal
697706
bind:open={imageModal.open}
698707
imageUrl={imageModal.imageUrl}

0 commit comments

Comments
 (0)