|
|
|
@@ -1,305 +0,0 @@
|
|
|
|
|
import { useDrizzle } from '~~/server/utils/drizzle';
|
|
|
|
|
import { generations, messages as messages_drizzle, messagesRelations } from '~~/db/schema';
|
|
|
|
|
import { type GenerationStreamEvent, type ChatMessage, type GenerationStatus } from '~~/server/types/chat';
|
|
|
|
|
import { eq } from 'drizzle-orm';
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Streaming generation state - only stores active stream controllers
|
|
|
|
|
* All persistent state lives in the database
|
|
|
|
|
*/
|
|
|
|
|
interface ActiveGenerationStream {
|
|
|
|
|
userId: string;
|
|
|
|
|
topicId: string;
|
|
|
|
|
clients: Set<ReadableStreamDefaultController<Uint8Array>>;
|
|
|
|
|
isGenerating: boolean;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
const activeGenerationStreams = new Map<string, ActiveGenerationStream>();
|
|
|
|
|
const db = useDrizzle();
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Add a client connection to an active generation stream
|
|
|
|
|
*/
|
|
|
|
|
export const addClientToGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): boolean => {
|
|
|
|
|
const stream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (!stream) {
|
|
|
|
|
return false;
|
|
|
|
|
}
|
|
|
|
|
stream.clients.add(controller);
|
|
|
|
|
return true;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Remove a client connection from an active generation stream
|
|
|
|
|
*/
|
|
|
|
|
export const removeClientFromGeneration = (generationId: string, controller: ReadableStreamDefaultController<Uint8Array>): void => {
|
|
|
|
|
const stream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (stream) {
|
|
|
|
|
stream.clients.delete(controller);
|
|
|
|
|
// Clean up if no clients left and generation is complete
|
|
|
|
|
if (stream.clients.size === 0 && !stream.isGenerating) {
|
|
|
|
|
activeGenerationStreams.delete(generationId);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Send an event to all connected clients for a generation
|
|
|
|
|
*/
|
|
|
|
|
export const sendToClients = (generationId: string, event: GenerationStreamEvent): void => {
|
|
|
|
|
const stream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (!stream) return;
|
|
|
|
|
|
|
|
|
|
for (const client of stream.clients) {
|
|
|
|
|
sendToClient(client, event);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Send an event to a single client
|
|
|
|
|
*/
|
|
|
|
|
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);
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Build a prompt from chat messages
|
|
|
|
|
*/
|
|
|
|
|
const buildPrompt = (messages: ChatMessage[]): string => {
|
|
|
|
|
return messages
|
|
|
|
|
.map((msg: ChatMessage) => {
|
|
|
|
|
const roleMap: Record<typeof msg.type, string> = {
|
|
|
|
|
system: 'System',
|
|
|
|
|
user: 'User',
|
|
|
|
|
agent: 'Assistant'
|
|
|
|
|
};
|
|
|
|
|
return `${roleMap[msg.type]}: ${msg.message}`;
|
|
|
|
|
})
|
|
|
|
|
.join('\n\n');
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Create a new pending generation in the database
|
|
|
|
|
* Returns the generation ID
|
|
|
|
|
*/
|
|
|
|
|
export const createPendingGeneration = async (userId: string, topicId: string, messages: ChatMessage[], regeneratesFrom?: string): Promise<string> => {
|
|
|
|
|
const generationValues: any = {
|
|
|
|
|
userId,
|
|
|
|
|
topicId,
|
|
|
|
|
status: 'pending' as GenerationStatus,
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
if (regeneratesFrom) {
|
|
|
|
|
generationValues.regeneratesFrom = regeneratesFrom;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
const [generation] = await db
|
|
|
|
|
.insert(generations)
|
|
|
|
|
.values(generationValues)
|
|
|
|
|
.returning();
|
|
|
|
|
|
|
|
|
|
return generation.id;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Get current generation status and content from database
|
|
|
|
|
*/
|
|
|
|
|
export const getGenerationStatus = async (generationId: string) => {
|
|
|
|
|
const [generation] = await db
|
|
|
|
|
.select()
|
|
|
|
|
.from(generations)
|
|
|
|
|
.where(eq(generations.id, generationId));
|
|
|
|
|
|
|
|
|
|
return generation || null;
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Start a generation: update status to active and begin streaming
|
|
|
|
|
* This is called when a client connects to the stream
|
|
|
|
|
*/
|
|
|
|
|
export const startGeneration = async (
|
|
|
|
|
generationId: string,
|
|
|
|
|
userId: string,
|
|
|
|
|
topicId: string,
|
|
|
|
|
messages: ChatMessage[],
|
|
|
|
|
controller: ReadableStreamDefaultController<Uint8Array>
|
|
|
|
|
): Promise<void> => {
|
|
|
|
|
try {
|
|
|
|
|
// Get current generation from database
|
|
|
|
|
const generation = await getGenerationStatus(generationId);
|
|
|
|
|
if (!generation) {
|
|
|
|
|
throw new Error('Generation not found');
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
// Create active stream tracking
|
|
|
|
|
activeGenerationStreams.set(generationId, {
|
|
|
|
|
userId,
|
|
|
|
|
topicId,
|
|
|
|
|
clients: new Set([controller]),
|
|
|
|
|
isGenerating: true,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Update status to active
|
|
|
|
|
await db
|
|
|
|
|
.update(generations)
|
|
|
|
|
.set({
|
|
|
|
|
status: 'active' as GenerationStatus,
|
|
|
|
|
startedAt: new Date(),
|
|
|
|
|
})
|
|
|
|
|
.where(eq(generations.id, generationId));
|
|
|
|
|
|
|
|
|
|
sendToClients(generationId, {
|
|
|
|
|
type: 'start',
|
|
|
|
|
data: null,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
const prompt = buildPrompt(messages);
|
|
|
|
|
const dummyResponse = generateDummyResponse(prompt, messages);
|
|
|
|
|
const tokens = dummyResponse.split(' ');
|
|
|
|
|
|
|
|
|
|
// Simulate token streaming
|
|
|
|
|
let accumulatedContent = '';
|
|
|
|
|
for (const token of tokens) {
|
|
|
|
|
await new Promise(resolve => setTimeout(resolve, 50));
|
|
|
|
|
|
|
|
|
|
accumulatedContent += token + ' ';
|
|
|
|
|
|
|
|
|
|
sendToClients(generationId, {
|
|
|
|
|
type: 'token',
|
|
|
|
|
data: token + ' ',
|
|
|
|
|
});
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
const finalContent = accumulatedContent.trim();
|
|
|
|
|
|
|
|
|
|
// Create message record for this generation
|
|
|
|
|
const messageValues: any = {
|
|
|
|
|
topicId,
|
|
|
|
|
userId,
|
|
|
|
|
content: finalContent,
|
|
|
|
|
isUser: false,
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
// If this is a regeneration, set the regeneratedFromId
|
|
|
|
|
if (generation.regeneratesFrom) {
|
|
|
|
|
messageValues.regeneratedFromId = generation.regeneratesFrom;
|
|
|
|
|
messageValues.isRegenerated = true;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
let [message] = await db
|
|
|
|
|
.insert(messages_drizzle)
|
|
|
|
|
.values(messageValues)
|
|
|
|
|
.returning();
|
|
|
|
|
|
|
|
|
|
// Update generation as completed
|
|
|
|
|
await db
|
|
|
|
|
.update(generations)
|
|
|
|
|
.set({
|
|
|
|
|
status: 'completed' as GenerationStatus,
|
|
|
|
|
completedAt: new Date(),
|
|
|
|
|
messageId: message.id,
|
|
|
|
|
})
|
|
|
|
|
.where(eq(generations.id, generationId));
|
|
|
|
|
|
|
|
|
|
let fmessage = await db.select().from(messages_drizzle).where(eq(messages_drizzle.userId, userId)).leftJoin(generations, eq(messages_drizzle.id, generations.messageId))
|
|
|
|
|
console.log(fmessage);
|
|
|
|
|
|
|
|
|
|
// Mark stream as no longer generating
|
|
|
|
|
const stream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (stream) {
|
|
|
|
|
stream.isGenerating = false;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
sendToClients(generationId, {
|
|
|
|
|
type: 'complete',
|
|
|
|
|
data: fmessage,
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Close all client connections
|
|
|
|
|
const finalStream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (finalStream) {
|
|
|
|
|
for (const client of finalStream.clients) {
|
|
|
|
|
try {
|
|
|
|
|
client.close();
|
|
|
|
|
} catch (error) {
|
|
|
|
|
// Client already closed
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
activeGenerationStreams.delete(generationId);
|
|
|
|
|
}
|
|
|
|
|
} catch (error) {
|
|
|
|
|
console.error('Generation failed:', error);
|
|
|
|
|
|
|
|
|
|
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
|
|
|
|
|
|
|
|
|
// Update generation as failed
|
|
|
|
|
await db
|
|
|
|
|
.update(generations)
|
|
|
|
|
.set({
|
|
|
|
|
status: 'failed' as GenerationStatus,
|
|
|
|
|
error: errorMessage,
|
|
|
|
|
completedAt: new Date(),
|
|
|
|
|
})
|
|
|
|
|
.where(eq(generations.id, generationId));
|
|
|
|
|
|
|
|
|
|
sendToClients(generationId, {
|
|
|
|
|
type: 'error',
|
|
|
|
|
data: { error: errorMessage },
|
|
|
|
|
});
|
|
|
|
|
|
|
|
|
|
// Close all client connections
|
|
|
|
|
const stream = activeGenerationStreams.get(generationId);
|
|
|
|
|
if (stream) {
|
|
|
|
|
for (const client of stream.clients) {
|
|
|
|
|
try {
|
|
|
|
|
client.close();
|
|
|
|
|
} catch (error) {
|
|
|
|
|
// Client already closed
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
activeGenerationStreams.delete(generationId);
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Check if a generation is currently being streamed
|
|
|
|
|
*/
|
|
|
|
|
export const isGenerationStreaming = (generationId: string): boolean => {
|
|
|
|
|
return activeGenerationStreams.has(generationId);
|
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* Generate a dummy response for testing
|
|
|
|
|
*/
|
|
|
|
|
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)];
|
|
|
|
|
};
|
