11import { sew } from "@/actions" ;
22import { _getConfiguredLanguageModelsFull , _getAISDKLanguageModelAndOptions , _updateChatMessages , _isOwnerOfChat } from "@/features/chat/actions" ;
33import { createAgentStream } from "@/features/chat/agent" ;
4- import { additionalChatRequestParamsSchema , LanguageModelInfo , SBChatMessage , SearchScope } from "@/features/chat/types" ;
4+ import { additionalChatRequestParamsSchema , LanguageModelInfo , SBChatMessage , SBChatMessageMetadata , SearchScope } from "@/features/chat/types" ;
55import { getAnswerPartFromAssistantMessage , getLanguageModelKey } from "@/features/chat/utils" ;
66import { apiHandler } from "@/lib/apiHandler" ;
77import { ErrorCode } from "@/lib/errorCodes" ;
@@ -10,8 +10,7 @@ import { isServiceError } from "@/lib/utils";
1010import { withOptionalAuthV2 } from "@/withAuthV2" ;
1111import { LanguageModelV2 as AISDKLanguageModelV2 } from "@ai-sdk/provider" ;
1212import * as Sentry from "@sentry/nextjs" ;
13- import { PrismaClient } from "@sourcebot/db" ;
14- import { createLogger } from "@sourcebot/shared" ;
13+ import { createLogger , env } from "@sourcebot/shared" ;
1514import { captureEvent } from "@/lib/posthog" ;
1615import {
1716 createUIMessageStream ,
@@ -89,20 +88,34 @@ export const POST = apiHandler(async (req: NextRequest) => {
8988
9089 const { model, providerOptions } = await _getAISDKLanguageModelAndOptions ( languageModelConfig ) ;
9190
91+ const expandedRepos = ( await Promise . all ( selectedSearchScopes . map ( async ( scope ) => {
92+ if ( scope . type === 'repo' ) return [ scope . value ] ;
93+ if ( scope . type === 'reposet' ) {
94+ const reposet = await prisma . searchContext . findFirst ( {
95+ where : { orgId : org . id , name : scope . value } ,
96+ include : { repos : true }
97+ } ) ;
98+ return reposet ? reposet . repos . map ( r => r . name ) : [ ] ;
99+ }
100+ return [ ] ;
101+ } ) ) ) . flat ( ) ;
102+
92103 await captureEvent ( 'wa_chat_message_sent' , {
93104 chatId : id ,
94105 messageCount : messages . length ,
95- } ) ;
106+ ...( env . EXPERIMENT_ASK_GH_ENABLED === 'true' ? { selectedRepos : expandedRepos } : { } ) ,
107+ } ) ;
96108
97109 const stream = await createMessageStream ( {
98110 chatId : id ,
99111 messages,
100- selectedSearchScopes,
112+ metadata : {
113+ selectedSearchScopes,
114+ } ,
115+ selectedRepos : expandedRepos ,
101116 model,
102117 modelName : languageModelConfig . displayName ?? languageModelConfig . model ,
103118 modelProviderOptions : providerOptions ,
104- orgId : org . id ,
105- prisma,
106119 onFinish : async ( { messages } ) => {
107120 await _updateChatMessages ( { chatId : id , messages, prisma } ) ;
108121 } ,
@@ -152,25 +165,23 @@ const mergeStreamAsync = async (stream: StreamTextResult<any, any>, writer: UIMe
152165interface CreateMessageStreamResponseProps {
153166 chatId : string ;
154167 messages : SBChatMessage [ ] ;
155- selectedSearchScopes : SearchScope [ ] ;
168+ selectedRepos : string [ ] ;
156169 model : AISDKLanguageModelV2 ;
157170 modelName : string ;
158- modelProviderOptions ?: Record < string , Record < string , JSONValue > > ;
159- orgId : number ;
160- prisma : PrismaClient ;
161171 onFinish : UIMessageStreamOnFinishCallback < SBChatMessage > ;
162172 onError : ( error : unknown ) => string ;
173+ modelProviderOptions ?: Record < string , Record < string , JSONValue > > ;
174+ metadata ?: Partial < SBChatMessageMetadata > ;
163175}
164176
165177export const createMessageStream = async ( {
166178 chatId,
167179 messages,
168- selectedSearchScopes,
180+ metadata,
181+ selectedRepos,
169182 model,
170183 modelName,
171184 modelProviderOptions,
172- orgId,
173- prisma,
174185 onFinish,
175186 onError,
176187} : CreateMessageStreamResponseProps ) => {
@@ -211,36 +222,12 @@ export const createMessageStream = async ({
211222
212223 const startTime = new Date ( ) ;
213224
214- const expandedRepos = ( await Promise . all ( selectedSearchScopes . map ( async ( scope ) => {
215- if ( scope . type === 'repo' ) {
216- return [ scope . value ] ;
217- }
218-
219- if ( scope . type === 'reposet' ) {
220- const reposet = await prisma . searchContext . findFirst ( {
221- where : {
222- orgId,
223- name : scope . value
224- } ,
225- include : {
226- repos : true
227- }
228- } ) ;
229-
230- if ( reposet ) {
231- return reposet . repos . map ( repo => repo . name ) ;
232- }
233- }
234-
235- return [ ] ;
236- } ) ) ) . flat ( )
237-
238225 const researchStream = await createAgentStream ( {
239226 model,
240227 providerOptions : modelProviderOptions ,
241228 inputMessages : messageHistory ,
242229 inputSources : sources ,
243- selectedRepos : expandedRepos ,
230+ selectedRepos,
244231 onWriteSource : ( source ) => {
245232 writer . write ( {
246233 type : 'data-source' ,
@@ -267,8 +254,8 @@ export const createMessageStream = async ({
267254 totalOutputTokens : totalUsage . outputTokens ,
268255 totalResponseTimeMs : new Date ( ) . getTime ( ) - startTime . getTime ( ) ,
269256 modelName,
270- selectedSearchScopes,
271257 traceId,
258+ ...metadata ,
272259 }
273260 } ) ;
274261
0 commit comments