Skip to content

Instantly share code, notes, and snippets.

@kirtirajsinh
Created January 12, 2024 15:42
Show Gist options
  • Select an option

  • Save kirtirajsinh/798432e5f2cc6162fb5f0394cbc0e482 to your computer and use it in GitHub Desktop.

Select an option

Save kirtirajsinh/798432e5f2cc6162fb5f0394cbc0e482 to your computer and use it in GitHub Desktop.
import { ReplicateStream, streamToResponse } from "ai";
import { experimental_buildLlama2Prompt } from "ai/prompts";
import { prisma } from "@/server/db";
import { replicate } from "@/utils/replicate";
import { getServerSession } from "next-auth";
import { authOptions } from "@/server/auth";
export default async function handler(req, res) {
const session = await getServerSession(req, res, authOptions);
if (!session) {
res.status(401).json({ message: "Unauthorized" });
return;
}
console.log(session, "session from chat.js");
const body = await req.body;
console.log(body, "body from chat.js");
if (session.user.remainingChat <= 0) {
res.status(401).json({ message: "You have no more chats left" });
return;
} else {
const messages = [
{
role: "system",
content: body?.prompt,
},
...body?.messages,
];
try {
const response = await reduceRemainingChat(session);
console.log(response, "response from chat.js");
} catch (e) {
console.log(e);
}
try {
const response = await replicate.predictions.create({
// You must enable streaming.
stream: true,
max_tokens: 100,
// The model must support streaming. See https://replicate.com/docs/streaming
// This is the model ID for Llama 2 70b Chat
version:
"f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
// Format the message list into the format expected by Llama 2
// @see https://github.com/vercel/ai/blob/99cf16edf0a09405d15d3867f997c96a8da869c6/packages/core/prompts/huggingface.ts#L53C1-L78C2
input: {
prompt: experimental_buildLlama2Prompt(messages),
},
});
const stream = await ReplicateStream(response);
return streamToResponse(stream, res);
} catch (e) {
console.log(e);
res.status(500).json({ error: "An error occurred" });
return;
}
}
}
async function reduceRemainingChat(session) {
console.log("I am in a reduce chat function");
const update = prisma.User.update({
where: { email: session.user.email },
data: {
remainingChat: session.user.remainingChat - 1,
},
});
return update;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment