add save to prisma

This commit is contained in:
Shu Ding 2023-05-22 13:10:52 +02:00
parent 562dbe7442
commit 609194517e
No known key found for this signature in database
GPG key ID: B84C6E25F5FEA16B
3 changed files with 97 additions and 76 deletions

View file

@ -1,7 +1,6 @@
import { OpenAIStream, openai } from "@/lib/openai"; import { OpenAIStream, openai } from "@/lib/openai";
import { getServerSession } from "@/lib/session/get-server-session"; import { getServerSession } from "@/lib/session/get-server-session";
import { prisma } from "@/lib/prisma";
export const runtime = "edge";
export async function POST(req: Request) { export async function POST(req: Request) {
const json = await req.json(); const json = await req.json();
@ -20,7 +19,25 @@ export async function POST(req: Request) {
const stream = await OpenAIStream(res); const stream = await OpenAIStream(res);
return new Response(stream, { let fullResponse = "";
const decoder = new TextDecoder();
const saveToPrisma = new TransformStream({
transform: async (chunk, controller) => {
controller.enqueue(chunk);
fullResponse += decoder.decode(chunk);
},
flush: async () => {
await prisma.chat.upsert({
where: {
id: json.id,
},
create: json,
update: json,
});
},
});
return new Response(stream.pipeThrough(saveToPrisma), {
status: 200, status: 200,
headers: { "Content-Type": "text/event-stream" }, headers: { "Content-Type": "text/event-stream" },
}); });

View file

@ -13,7 +13,7 @@ export interface ChatProps {
} }
export function Chat({ export function Chat({
id: _id, id,
// create, // create,
messages, messages,
}: ChatProps) { }: ChatProps) {
@ -22,7 +22,7 @@ export function Chat({
const { isLoading, messageList, appendUserMessage, reloadLastMessage } = const { isLoading, messageList, appendUserMessage, reloadLastMessage } =
usePrompt({ usePrompt({
messages, messages,
_id, id,
// onCreate: (id: string) => { // onCreate: (id: string) => {
// router.push(`/chat/${id}`); // router.push(`/chat/${id}`);
// }, // },

View file

@ -4,10 +4,10 @@ import { nanoid } from "@/lib/utils";
export function usePrompt({ export function usePrompt({
messages = [], messages = [],
_id, id,
}: { }: {
messages?: Message[]; messages?: Message[];
_id: string | undefined | null; id: string | undefined;
}) { }) {
const [isLoading, setIsLoading] = useState(false); const [isLoading, setIsLoading] = useState(false);
const [messageList, setMessageList] = useState(messages); const [messageList, setMessageList] = useState(messages);
@ -23,83 +23,87 @@ export function usePrompt({
messageListRef.current = messageList; messageListRef.current = messageList;
}, [messageList]); }, [messageList]);
const appendUserMessage = useCallback(async (content: string | Message) => { const appendUserMessage = useCallback(
// Prevent multiple requests at once async (content: string | Message) => {
if (isLoadingRef.current) return; // Prevent multiple requests at once
if (isLoadingRef.current) return;
const userMsg = const userMsg =
typeof content === "string" typeof content === "string"
? ({ id: nanoid(10), role: "user", content } as Message) ? ({ id: nanoid(10), role: "user", content } as Message)
: content; : content;
const assMsg = { const assMsg = {
id: nanoid(10), id: nanoid(10),
role: "assistant", role: "assistant",
content: "", content: "",
} as Message; } as Message;
const messageListSnapshot = messageListRef.current; const messageListSnapshot = messageListRef.current;
// Reset output // Reset output
setIsLoading(true); setIsLoading(true);
try { try {
// Set user input immediately // Set user input immediately
setMessageList([...messageListSnapshot, userMsg]); setMessageList([...messageListSnapshot, userMsg]);
// If streaming, we need to use fetchEventSource directly // If streaming, we need to use fetchEventSource directly
const response = await fetch(`/api/generate`, { const response = await fetch(`/api/generate`, {
method: "POST", method: "POST",
body: JSON.stringify({ body: JSON.stringify({
messages: [...messageListSnapshot, userMsg].map((m) => ({ id: id || nanoid(10),
role: m.role, messages: [...messageListSnapshot, userMsg].map((m) => ({
content: m.content, role: m.role,
})), content: m.content,
}), })),
headers: { "Content-Type": "application/json" }, }),
}); headers: { "Content-Type": "application/json" },
// This data is a ReadableStream });
const data = response.body; // This data is a ReadableStream
if (!data) { const data = response.body;
return; if (!data) {
} return;
const reader = data.getReader();
const decoder = new TextDecoder();
let done = false;
let accumulatedValue = ""; // Variable to accumulate chunks
while (!done) {
const { value, done: doneReading } = await reader.read();
done = doneReading;
const chunkValue = decoder.decode(value);
accumulatedValue += chunkValue; // Accumulate the chunk value
// Check if the accumulated value contains the delimiter
const delimiter = "\n";
const chunks = accumulatedValue.split(delimiter);
// Process all chunks except the last one (which may be incomplete)
while (chunks.length > 1) {
const chunkToDispatch = chunks.shift(); // Get the first chunk
if (chunkToDispatch && chunkToDispatch.length > 0) {
const chunk = JSON.parse(chunkToDispatch);
assMsg.content += chunk;
setMessageList([...messageListSnapshot, userMsg, assMsg]);
}
} }
// The last chunk may be incomplete, so keep it in the accumulated value const reader = data.getReader();
accumulatedValue = chunks[0]; const decoder = new TextDecoder();
} let done = false;
let accumulatedValue = ""; // Variable to accumulate chunks
// Process any remaining accumulated value after the loop is done while (!done) {
if (accumulatedValue.length > 0) { const { value, done: doneReading } = await reader.read();
assMsg.content += accumulatedValue; done = doneReading;
setMessageList([...messageListSnapshot, userMsg, assMsg]); const chunkValue = decoder.decode(value);
accumulatedValue += chunkValue; // Accumulate the chunk value
// Check if the accumulated value contains the delimiter
const delimiter = "\n";
const chunks = accumulatedValue.split(delimiter);
// Process all chunks except the last one (which may be incomplete)
while (chunks.length > 1) {
const chunkToDispatch = chunks.shift(); // Get the first chunk
if (chunkToDispatch && chunkToDispatch.length > 0) {
const chunk = JSON.parse(chunkToDispatch);
assMsg.content += chunk;
setMessageList([...messageListSnapshot, userMsg, assMsg]);
}
}
// The last chunk may be incomplete, so keep it in the accumulated value
accumulatedValue = chunks[0];
}
// Process any remaining accumulated value after the loop is done
if (accumulatedValue.length > 0) {
assMsg.content += accumulatedValue;
setMessageList([...messageListSnapshot, userMsg, assMsg]);
}
} finally {
setIsLoading(false);
} }
} finally { },
setIsLoading(false); [id]
} );
}, []);
const reloadLastMessage = useCallback(async () => { const reloadLastMessage = useCallback(async () => {
// Prevent multiple requests at once // Prevent multiple requests at once