initial commit

server/utils/auth.ts

@@ -0,0 +1,15 @@
+import { H3Event } from 'h3';
+import { auth } from '~~/lib/auth';
+
+export const protectRoute = async (event: H3Event) => {
+    const sessionData = await auth.api.getSession(event);
+
+    if (sessionData === null) {
+        throw createError({
+            statusCode: 401,
+            statusMessage: 'Unauthorized',
+        });
+    }
+
+    event.context.user = sessionData.user;
+};

server/utils/drizzle.ts

@@ -0,0 +1,11 @@
+import { drizzle } from "drizzle-orm/node-postgres";
+import * as schema from "~~/db/schema";
+
+export const useDrizzle = () => {
+    return drizzle(process.env.DATABASE_URL!)
+}
+
+export const tables = schema;
+
+export const UserInsert = schema.user.$inferInsert;
+export type UserRegisterType = Omit<typeof UserInsert, "createdAt" | "updatedAt" | "id" | "emailVerified">;

server/utils/generation.ts

@@ -0,0 +1,230 @@
+import { useDrizzle } from '~~/server/utils/drizzle';
+import { generations, messages as messages_drizzle } from '~~/db/schema';
+import { type GenerationStreamEvent, type ChatMessage, type MessageType } from '~~/server/types/chat';
+import { eq } from 'drizzle-orm';
+
+interface ActiveGeneration {
+    userId: string;
+    topicId: string;
+    messages: ChatMessage[];
+    content: string;
+    clients: Set<ReadableStreamDefaultController<Uint8Array>>;
+    complete: boolean;
+}
+
+interface PendingGeneration {
+    generationId: string;
+    userId: string;
+    topicId: string;
+    messages: ChatMessage[];
+    timeout: NodeJS.Timeout;
+    expired: boolean;
+}
+
+const activeGenerations = new Map<string, ActiveGeneration>();
+const pendingGenerations = new Map<string, PendingGeneration>();
+
+export const getActiveGeneration = (generationId: string): ActiveGeneration | undefined => {
+    return activeGenerations.get(generationId);
+};
+
+export const getPendingGeneration = (generationId: string): PendingGeneration | undefined => {
+    return pendingGenerations.get(generationId);
+};
+
+export const isGenerationActive = (generationId: string): boolean => {
+    return activeGenerations.has(generationId);
+};
+
+export const isGenerationPending = (generationId: string): boolean => {
+    return pendingGenerations.has(generationId);
+};
+
+export const addClientToGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): boolean => {
+    const generation = activeGenerations.get(generationId);
+    if (!generation) {
+        return false;
+    }
+    generation.clients.add(controller);
+    return true;
+};
+
+export const removeClientFromGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): void => {
+    const generation = activeGenerations.get(generationId);
+    if (generation) {
+        generation.clients.delete(controller);
+        if (generation.clients.size === 0 && generation.complete) {
+            activeGenerations.delete(generationId);
+        }
+    }
+};
+
+const sendToClients = (generation: ActiveGeneration, event: GenerationStreamEvent): void => {
+    for (const client of generation.clients) {
+        sendToClient(client, event);
+    }
+};
+
+export const sendToClient = (client: ReadableStreamDefaultController<Uint8Array>, event: GenerationStreamEvent): void => {
+    const data = JSON.stringify(event);
+    const encoder = new TextEncoder();
+    try {
+        client.enqueue(encoder.encode(`${data}\n`));
+    } catch (error) {
+        console.error('Failed to send to client:', error);
+    }
+}
+
+const buildPrompt = (messages: ChatMessage[]): string => {
+    return messages
+        .map((msg: ChatMessage) => {
+            const roleMap: Record<MessageType, string> = {
+                system: 'System',
+                user: 'User',
+                agent: 'Assistant'
+            };
+            return `${roleMap[msg.type]}: ${msg.message}`;
+        })
+        .join('\n\n');
+};
+
+const db = useDrizzle();
+
+export const startGeneration = async (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): Promise<void> => {
+    const pending = pendingGenerations.get(generationId);
+
+    if (!pending) {
+        console.error(`Generation ${generationId} not found in pending generations`);
+        return;
+    }
+
+    clearTimeout(pending.timeout);
+    pendingGenerations.delete(generationId);
+
+    const { userId, topicId, messages } = pending;
+    const prompt = buildPrompt(messages);
+
+    const generation: ActiveGeneration = {
+        userId,
+        topicId,
+        messages,
+        content: '',
+        clients: new Set([controller]),
+        complete: false
+    };
+
+    await db.insert(generations).values({
+        id: generationId,
+        userId,
+        topicId,
+        messageId: null
+    });
+
+    activeGenerations.set(generationId, generation);
+
+    try {
+        const dummyResponse = generateDummyResponse(prompt, messages);
+        const tokens = dummyResponse.split(' ');
+
+        for (const token of tokens) {
+            await new Promise(resolve => setTimeout(resolve, 50));
+
+            generation.content += token + ' ';
+
+            sendToClients(generation, {
+                type: 'token',
+                data: token + ' '
+            });
+        }
+
+        const [message] = await db.insert(messages_drizzle).values({
+            topicId: generation.topicId,
+            userId: generation.userId,
+            isUser: false,
+            content: generation.content.trim(),
+            model: 'dummy-model-v1',
+            tokensGenerated: tokens.length,
+            tokensUsedThinking: 0
+        }).returning();
+
+        generation.complete = true;
+
+        sendToClients(generation, {
+            type: 'complete',
+            data: message
+        });
+
+        if (generation.clients.size === 0) {
+            activeGenerations.delete(generationId);
+        } else {
+            generation.clients.forEach((client) => {
+                try {
+                    client.close();
+                } catch {
+                }
+            });
+            activeGenerations.delete(generationId);
+        }
+
+        await db.update(generations).set({
+            messageId: message.id
+        }).where(eq(generations.id, generationId));
+
+    } catch (error) {
+        console.error('Generation failed:', error);
+
+        sendToClients(generation, {
+            type: 'error',
+            data: error instanceof Error ? error.message : 'Unknown error'
+        });
+
+        await db.delete(generations).where(eq(generations.id, generationId));
+
+        activeGenerations.delete(generationId);
+    }
+};
+
+export const registerPendingGeneration = (userId: string, generationId: string, topicId: string, messages: ChatMessage[]): void => {
+    const timeout = setTimeout(() => {
+        const gen = pendingGenerations.get(generationId)
+        if (gen) gen.expired = true;
+        console.log(`Generation ${generationId} expired - no client connected within 60 seconds`);
+    }, 60000);
+
+    pendingGenerations.set(generationId, {
+        generationId,
+        userId,
+        topicId,
+        messages,
+        timeout,
+        expired: false
+    });
+
+    console.log(`Registered pending generation ${generationId}, waiting for client connection...`);
+};
+
+const generateDummyResponse = (prompt: string, messages: ChatMessage[]): string => {
+    const responses = [
+        "This is a simulated response to your prompt. In a real implementation, this would be generated by an AI model like GPT-4 or Claude."
+        + " I'm processing your message about: " + prompt.substring(0, 50) + "... "
+        + "This dummy generation demonstrates the streaming and background save functionality.",
+
+        "I understand your query. This is a placeholder response that simulates AI-generated content."
+        + " The system will continue generating this response even if you close the tab, and it will"
+        + " automatically save to the database when complete.",
+
+        "Here's a simulated AI response. This demonstrates two key features:"
+        + " 1) The generation continues in the background even if you disconnect,"
+        + " 2) The complete response is automatically saved to the database without requiring"
+        + " a separate update request from the client."
+    ];
+
+    const lastUserMessage = messages[messages.length - 1]?.message.toLowerCase() || '';
+
+    if (lastUserMessage.includes('hello') || lastUserMessage.includes('hi')) {
+        return "Hello! I'm a dummy AI assistant. This is a simulated response to your greeting."
+            + " In production, this would be replaced with actual AI-generated content from an LLM provider.";
+    }
+
+    return responses[Math.floor(Math.random() * responses.length)];
+};