@@ -10,21 +10,21 @@ import { isServiceError } from "@/lib/utils";
1010import { prisma } from "@/prisma" ;
1111import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock' ;
1212import { AnthropicProviderOptions , createAnthropic } from '@ai-sdk/anthropic' ;
13+ import { createAzure } from '@ai-sdk/azure' ;
14+ import { createDeepSeek } from '@ai-sdk/deepseek' ;
1315import { createGoogleGenerativeAI } from '@ai-sdk/google' ;
1416import { createVertex } from '@ai-sdk/google-vertex' ;
1517import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic' ;
16- import { createOpenAI , OpenAIResponsesProviderOptions } from "@ai-sdk/openai" ;
1718import { createMistral } from '@ai-sdk/mistral' ;
18- import { createXai } from ' @ai-sdk/xai' ;
19+ import { createOpenAI , OpenAIResponsesProviderOptions } from " @ai-sdk/openai" ;
1920import { LanguageModelV2 as AISDKLanguageModelV2 } from "@ai-sdk/provider" ;
21+ import { createXai } from '@ai-sdk/xai' ;
22+ import { createOpenRouter } from '@openrouter/ai-sdk-provider' ;
2023import * as Sentry from "@sentry/nextjs" ;
2124import { getTokenFromConfig } from "@sourcebot/crypto" ;
2225import { OrgRole } from "@sourcebot/db" ;
2326import { createLogger } from "@sourcebot/logger" ;
2427import { LanguageModel } from "@sourcebot/schemas/v3/index.type" ;
25- import { createAzure } from '@ai-sdk/azure' ;
26- import { createDeepSeek } from '@ai-sdk/deepseek' ;
27- import { createOpenRouter } from '@openrouter/ai-sdk-provider' ;
2828import {
2929 createUIMessageStream ,
3030 createUIMessageStreamResponse ,
@@ -139,7 +139,6 @@ const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandl
139139
140140 const { model, providerOptions, headers } = await getAISDKLanguageModelAndOptions ( languageModelConfig , org . id ) ;
141141
142- // @todo : refactor this
143142 if (
144143 messages . length === 1 &&
145144 messages [ 0 ] . role === "user" &&
@@ -149,15 +148,10 @@ const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandl
149148 const content = messages [ 0 ] . parts [ 0 ] . text ;
150149
151150 const title = await generateChatTitle ( content , model ) ;
152- if ( title ) {
153- updateChatName ( {
154- chatId : id ,
155- name : title ,
156- } , domain ) ;
157- }
158- else {
159- logger . error ( "Failed to generate chat title." ) ;
160- }
151+ await updateChatName ( {
152+ chatId : id ,
153+ name : title ,
154+ } , domain ) ;
161155 }
162156
163157 const traceId = randomUUID ( ) ;
@@ -184,86 +178,73 @@ const chatHandler = ({ messages, id, selectedRepos, languageModelId }: ChatHandl
184178 }
185179 } ) . filter ( message => message !== undefined ) ;
186180
187- try {
188- const stream = createUIMessageStream < SBChatMessage > ( {
189- execute : async ( { writer } ) => {
190- writer . write ( {
191- type : 'start' ,
192- } ) ;
193-
194- const startTime = new Date ( ) ;
195-
196- const researchStream = await createAgentStream ( {
197- model,
198- providerOptions,
199- headers,
200- inputMessages : messageHistory ,
201- inputSources : sources ,
202- selectedRepos,
203- onWriteSource : ( source ) => {
204- writer . write ( {
205- type : 'data-source' ,
206- data : source ,
207- } ) ;
208- } ,
181+ const stream = createUIMessageStream < SBChatMessage > ( {
182+ execute : async ( { writer } ) => {
183+ writer . write ( {
184+ type : 'start' ,
185+ } ) ;
186+
187+ const startTime = new Date ( ) ;
188+
189+ const researchStream = await createAgentStream ( {
190+ model,
191+ providerOptions,
192+ headers,
193+ inputMessages : messageHistory ,
194+ inputSources : sources ,
195+ selectedRepos,
196+ onWriteSource : ( source ) => {
197+ writer . write ( {
198+ type : 'data-source' ,
199+ data : source ,
200+ } ) ;
201+ } ,
202+ traceId,
203+ } ) ;
204+
205+ await mergeStreamAsync ( researchStream , writer , {
206+ sendReasoning : true ,
207+ sendStart : false ,
208+ sendFinish : false ,
209+ } ) ;
210+
211+ const totalUsage = await researchStream . totalUsage ;
212+
213+ writer . write ( {
214+ type : 'message-metadata' ,
215+ messageMetadata : {
216+ totalTokens : totalUsage . totalTokens ,
217+ totalInputTokens : totalUsage . inputTokens ,
218+ totalOutputTokens : totalUsage . outputTokens ,
219+ totalResponseTimeMs : new Date ( ) . getTime ( ) - startTime . getTime ( ) ,
220+ modelName : languageModelConfig . displayName ?? languageModelConfig . model ,
209221 traceId,
210- } ) ;
211-
212- await mergeStreamAsync ( researchStream , writer , {
213- sendReasoning : true ,
214- sendStart : false ,
215- sendFinish : false ,
216- } ) ;
217-
218- const totalUsage = await researchStream . totalUsage ;
219-
220- writer . write ( {
221- type : 'message-metadata' ,
222- messageMetadata : {
223- totalTokens : totalUsage . totalTokens ,
224- totalInputTokens : totalUsage . inputTokens ,
225- totalOutputTokens : totalUsage . outputTokens ,
226- totalResponseTimeMs : new Date ( ) . getTime ( ) - startTime . getTime ( ) ,
227- modelName : languageModelConfig . displayName ?? languageModelConfig . model ,
228- traceId,
229- }
230- } )
231-
232-
233- writer . write ( {
234- type : 'finish' ,
235- } ) ;
236- } ,
237- onError : errorHandler ,
238- originalMessages : messages ,
239- onFinish : async ( { messages } ) => {
240- await updateChatMessages ( {
241- chatId : id ,
242- messages
243- } , domain ) ;
244- } ,
245- } ) ;
222+ }
223+ } )
246224
247- return createUIMessageStreamResponse ( {
248- stream,
249- } ) ;
250- } catch ( error ) {
251- logger . error ( error )
252- logger . error ( "Error stack:" , error instanceof Error ? error . stack : "No stack trace" )
253- Sentry . captureException ( error ) ;
254225
255- return serviceErrorResponse ( {
256- statusCode : StatusCodes . INTERNAL_SERVER_ERROR ,
257- errorCode : ErrorCode . UNEXPECTED_ERROR ,
258- message : error instanceof Error ? error . message : "Unknown error" ,
259- } ) ;
260- }
226+ writer . write ( {
227+ type : 'finish' ,
228+ } ) ;
229+ } ,
230+ onError : errorHandler ,
231+ originalMessages : messages ,
232+ onFinish : async ( { messages } ) => {
233+ await updateChatMessages ( {
234+ chatId : id ,
235+ messages
236+ } , domain ) ;
237+ } ,
238+ } ) ;
239+
240+ return createUIMessageStreamResponse ( {
241+ stream,
242+ } ) ;
261243 } , /* minRequiredRole = */ OrgRole . GUEST ) , /* allowSingleTenantUnauthedAccess = */ true
262244 ) ) ;
263245
264246const generateChatTitle = async ( message : string , model : AISDKLanguageModelV2 ) => {
265- try {
266- const prompt = `Convert this question into a short topic title (max 50 characters).
247+ const prompt = `Convert this question into a short topic title (max 50 characters).
267248
268249Rules:
269250- Do NOT include question words (what, where, how, why, when, which)
@@ -279,17 +260,12 @@ Examples:
279260
280261User question: ${ message } ` ;
281262
282- const result = await generateText ( {
283- model,
284- prompt,
285- maxOutputTokens : 20 ,
286- } ) ;
263+ const result = await generateText ( {
264+ model,
265+ prompt,
266+ } ) ;
287267
288- return result . text ;
289- } catch ( error ) {
290- logger . error ( "Error generating summary:" , error )
291- return undefined ;
292- }
268+ return result . text ;
293269}
294270
295271const getAISDKLanguageModelAndOptions = async ( config : LanguageModel , orgId : number ) : Promise < {
0 commit comments