design.ts107 lines · main
| 1 | import { streamText, tool } from 'ai' |
| 2 | import { source } from 'common-tags' |
| 3 | import { NextApiRequest, NextApiResponse } from 'next' |
| 4 | import { z } from 'zod' |
| 5 | |
| 6 | import { getModel } from '@/lib/ai/model' |
| 7 | import { DEFAULT_COMPLETION_MODEL } from '@/lib/ai/model.utils' |
| 8 | import apiWrapper from '@/lib/api/apiWrapper' |
| 9 | |
| 10 | export const maxDuration = 60 |
| 11 | |
| 12 | const 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 | |
| 17 | const 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 | |
| 50 | async 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 | |
| 62 | const wrapper = (req: NextApiRequest, res: NextApiResponse) => |
| 63 | apiWrapper(req, res, handler, { withAuth: true }) |
| 64 | |
| 65 | export default wrapper |
| 66 | |
| 67 | async 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 | } |