design.ts107 lines · main
1import { streamText, tool } from 'ai'
2import { source } from 'common-tags'
3import { NextApiRequest, NextApiResponse } from 'next'
4import { z } from 'zod'
5
6import { getModel } from '@/lib/ai/model'
7import { DEFAULT_COMPLETION_MODEL } from '@/lib/ai/model.utils'
8import apiWrapper from '@/lib/api/apiWrapper'
9
10export const maxDuration = 60
11
12const ServiceSchema = z.object({
13 name: z.enum(['Auth', 'Storage', 'Database', 'Edge Function', 'Cron', 'Queues', 'Vector']),
14 reason: z.string().describe("The reason why this service is needed for the user's use case"),
15})
16
17const getTools = () => {
18 return {
19 executeSql: tool({
20 description: 'Save the generated database schema definition',
21 inputSchema: z.object({
22 sql: z.string().describe('The SQL schema definition'),
23 }),
24 }),
25
26 reset: tool({
27 description: 'Reset the database, services and start over',
28 inputSchema: z.object({}),
29 }),
30
31 setServices: tool({
32 description:
33 'Set the entire list of Briven services needed for the project. Always include the full list',
34 inputSchema: z.object({
35 services: z
36 .array(ServiceSchema)
37 .describe('Array of services with reasons why they are needed'),
38 }),
39 }),
40
41 setTitle: tool({
42 description: "Set the project title based on the user's description",
43 inputSchema: z.object({
44 title: z.string().describe('The project title'),
45 }),
46 }),
47 }
48}
49
50async function handler(req: NextApiRequest, res: NextApiResponse) {
51 const { method } = req
52
53 switch (method) {
54 case 'POST':
55 return handlePost(req, res)
56 default:
57 res.setHeader('Allow', ['POST'])
58 res.status(405).json({ data: null, error: { message: `Method ${method} Not Allowed` } })
59 }
60}
61
62const wrapper = (req: NextApiRequest, res: NextApiResponse) =>
63 apiWrapper(req, res, handler, { withAuth: true })
64
65export default wrapper
66
67async function handlePost(req: NextApiRequest, res: NextApiResponse) {
68 const { modelParams, error: modelError } = await getModel({
69 provider: 'openai',
70 modelEntry: DEFAULT_COMPLETION_MODEL,
71 })
72
73 if (modelError) {
74 return res.status(500).json({ error: modelError.message })
75 }
76
77 const { messages } = req.body
78
79 const result = streamText({
80 ...modelParams,
81 system: source`
82 You are a Briven expert who helps people set up their Briven project. You specializes in database schema design. You are to help the user design a database schema for their application but also suggest Briven services they should use.
83
84 When designing database schemas, follow these rules:
85 - Generate the entire schema
86 - For primary keys, always use "id bigint primary key generated always as identity" (not serial)
87 - Prefer creating foreign key references in the create statement
88 - Prefer 'text' over 'varchar'
89 - Prefer 'timestamp with time zone' over 'date'
90 - In Briven, the auth schema already has a users table which is used to store users
91 - Create a profiles table in the public schema where the primary id is uuid and references the auth.users schema instead of creating a users table
92 - Always include appropriate indexes and foreign key constraints.
93
94 Follow these rules:
95 1. Generate a database schema that meets the user's requirements by calling the executeSql tool. Make your best guess without needing to ask for more.
96 2. Set the services required for the user's use case by calling the setServices tool.
97 3. Set the project title by calling the setTitle tool.
98 4. Always respond with a short single paragraph of less than 80 words of what you changed and the current state of the schema.
99
100 If user requests to reset the database, call the reset tool.
101 `,
102 messages,
103 tools: getTools(),
104 })
105
106 result.pipeUIMessageStreamToResponse(res, { headers: { 'Content-Encoding': 'none' } })
107}