konnosDashboard
rate.ts185 lines · main
1import { generateText, Output } from 'ai'
2import { currentLogger } from 'braintrust'
3import { IS_PLATFORM } from 'common'
4import { NextApiRequest, NextApiResponse } from 'next'
5import { z } from 'zod'
6
7import { rateMessageResponseSchema } from '@/components/ui/AIAssistantPanel/Message.utils'
8import type { AiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi'
9import { getOrgAIDetails, getProjectAIDetails } from '@/lib/ai/ai-details'
10import { IS_TRACING_ENABLED, isTracingAllowed } from '@/lib/ai/braintrust-logger'
11import { getModel } from '@/lib/ai/model'
12import { DEFAULT_COMPLETION_MODEL } from '@/lib/ai/model.utils'
13import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer'
14import apiWrapper from '@/lib/api/apiWrapper'
15
16export const maxDuration = 30
17
18async function handler(req: NextApiRequest, res: NextApiResponse) {
19 const { method } = req
20
21 switch (method) {
22 case 'POST':
23 return handlePost(req, res)
24 default:
25 res.setHeader('Allow', ['POST'])
26 res.status(405).json({ data: null, error: { message: `Method ${method} Not Allowed` } })
27 }
28}
29
30const requestBodySchema = z.object({
31 rating: z.enum(['positive', 'negative']),
32 messages: z.array(z.any()),
33 messageId: z.string(),
34 projectRef: z.string(),
35 orgSlug: z.string().optional(),
36 reason: z.string().optional(),
37 spanId: z.string().optional(),
38})
39
40export async function handlePost(req: NextApiRequest, res: NextApiResponse) {
41 const authorization = req.headers.authorization
42 const accessToken = authorization?.replace('Bearer ', '')
43
44 if (IS_PLATFORM && !accessToken) {
45 return res.status(401).json({ error: 'Authorization token is required' })
46 }
47
48 const body = typeof req.body === 'string' ? JSON.parse(req.body) : req.body
49 const { data, error: parseError } = requestBodySchema.safeParse(body)
50
51 if (parseError) {
52 return res.status(400).json({ error: 'Invalid request body', issues: parseError.issues })
53 }
54
55 const { rating, messages: rawMessages, projectRef, orgSlug, reason, spanId } = data
56
57 let aiOptInLevel: AiOptInLevel = 'disabled'
58 let orgHasHipaaAddon: boolean | undefined
59 let projectIsSensitive: boolean | undefined
60 let projectRegion: string | undefined
61
62 if (!IS_PLATFORM) {
63 aiOptInLevel = 'schema'
64 }
65
66 if (IS_PLATFORM && orgSlug && authorization && projectRef) {
67 try {
68 const [orgDetails, projectDetails] = await Promise.all([
69 getOrgAIDetails({ orgSlug, authorization }),
70 getProjectAIDetails({ projectRef, authorization }),
71 ])
72
73 aiOptInLevel = orgDetails.aiOptInLevel
74 orgHasHipaaAddon = orgDetails.hasHipaaAddon
75 projectIsSensitive = projectDetails.isSensitive
76 projectRegion = projectDetails.region
77 } catch (error) {
78 return res.status(400).json({
79 error: 'There was an error fetching your organization details',
80 })
81 }
82 }
83
84 // Only returns last 7 messages
85 // Filters out tool outputs based on opt-in level using sanitizeMessagePart
86 const messages = (rawMessages || []).slice(-7).map((msg: any) => {
87 if (msg && msg.role === 'assistant' && 'results' in msg) {
88 const cleanedMsg = { ...msg }
89 delete cleanedMsg.results
90 return cleanedMsg
91 }
92 if (msg && msg.role === 'assistant' && msg.parts) {
93 const cleanedParts = msg.parts.map((part: any) => {
94 return sanitizeMessagePart(part, aiOptInLevel)
95 })
96 return { ...msg, parts: cleanedParts }
97 }
98 return msg
99 })
100
101 try {
102 const { modelParams, error: modelError } = await getModel({
103 provider: 'openai',
104 modelEntry: DEFAULT_COMPLETION_MODEL,
105 })
106
107 if (modelError) {
108 return res.status(500).json({ error: modelError.message })
109 }
110
111 const { output } = await generateText({
112 ...modelParams,
113 output: Output.object({ schema: rateMessageResponseSchema }),
114 prompt: `
115Your job is to look at a Briven Assistant conversation, which the user has given feedback on, and classify it.
116
117The user gave this feedback: ${rating === 'positive' ? 'THUMBS UP (positive)' : 'THUMBS DOWN (negative)'}
118${reason ? `\nUser's reason: ${reason}` : ''}
119
120Raw conversation:
121${JSON.stringify(messages)}
122
123Instructions:
1241. Classify the conversation into ONE of these categories:
125 - sql_generation: Generating SQL queries, DML statements
126 - schema_design: Creating tables, columns, relationships
127 - rls_policies: Row Level Security policies
128 - edge_functions: Edge Functions or serverless functions
129 - database_optimization: Performance, indexes, optimization
130 - debugging: Helping debug errors or issues
131 - general_help: General questions about Briven features
132 - other: Anything else
133`,
134 })
135
136 // Log feedback to Braintrust if tracing is enabled and span ID is available
137 if (
138 IS_TRACING_ENABLED &&
139 isTracingAllowed({ orgHasHipaaAddon, projectIsSensitive, projectRegion }) &&
140 spanId
141 ) {
142 try {
143 const logger = currentLogger()
144 logger?.logFeedback({
145 id: spanId,
146 scores: { 'User Rating': rating === 'positive' ? 1 : 0 },
147 comment: reason,
148 source: 'external',
149 })
150 logger?.updateSpan({
151 id: spanId,
152 metadata: { feedbackCategory: output.category },
153 })
154 } catch (error) {
155 console.error('Failed to log feedback to Braintrust:', error)
156 }
157 }
158
159 return res.json({
160 category: output.category,
161 })
162 } catch (error) {
163 if (error instanceof Error) {
164 console.error(`Classifying feedback failed:`, error)
165
166 // Check for context length error
167 if (error.message.includes('context_length') || error.message.includes('too long')) {
168 return res.status(400).json({
169 error: 'The conversation is too large to analyze',
170 })
171 }
172 } else {
173 console.error(`Unknown error: ${error}`)
174 }
175
176 return res.status(500).json({
177 error: 'There was an unknown error analyzing the feedback.',
178 })
179 }
180}
181
182const wrapper = (req: NextApiRequest, res: NextApiResponse) =>
183 apiWrapper(req, res, handler, { withAuth: true })
184
185export default wrapper