@@ -92,13 +92,14 @@ test('should set error to true', (t, end) => {
9292 }
9393
9494 const api = helper . getAgentApi ( )
95- helper . runInTransaction ( agent , ( ) => {
95+ helper . runInTransaction ( agent , ( tx ) => {
9696 api . startSegment ( 'fakeSegment' , false , ( ) => {
9797 const segment = agent . tracer . getSegment ( )
9898 segment . end ( )
9999 const embeddingEvent = new LlmEmbedding ( {
100100 agent,
101101 segment,
102+ transaction : tx ,
102103 request : req ,
103104 response : res ,
104105 error : true
@@ -123,11 +124,12 @@ test('respects record_content by not recording content when set to false', (t, e
123124 const api = helper . getAgentApi ( )
124125 api . setLlmTokenCountCallback ( cb )
125126
126- helper . runInTransaction ( agent , ( ) => {
127+ helper . runInTransaction ( agent , ( tx ) => {
127128 const segment = agent . tracer . getSegment ( )
128129 const embeddingEvent = new LlmEmbedding ( {
129130 agent,
130131 segment,
132+ transaction : tx ,
131133 request : req ,
132134 response : res
133135 } )
@@ -150,11 +152,12 @@ test('respects record_content by recording content when true', (t, end) => {
150152
151153 const api = helper . getAgentApi ( )
152154 api . setLlmTokenCountCallback ( cb )
153- helper . runInTransaction ( agent , ( ) => {
155+ helper . runInTransaction ( agent , ( tx ) => {
154156 const segment = agent . tracer . getSegment ( )
155157 const embeddingEvent = new LlmEmbedding ( {
156158 agent,
157159 segment,
160+ transaction : tx ,
158161 request : req ,
159162 response : res
160163 } )
@@ -176,11 +179,12 @@ test('does not calculate tokens when no content exists', (t, end) => {
176179
177180 const api = helper . getAgentApi ( )
178181 api . setLlmTokenCountCallback ( cb )
179- helper . runInTransaction ( agent , ( ) => {
182+ helper . runInTransaction ( agent , ( tx ) => {
180183 const segment = agent . tracer . getSegment ( )
181184 const embeddingEvent = new LlmEmbedding ( {
182185 agent,
183186 segment,
187+ transaction : tx ,
184188 request : req ,
185189 response : res
186190 } )
@@ -197,13 +201,14 @@ test('assigns total_tokens from response', (t, end) => {
197201 model : 'gpt-3.5-turbo-0613'
198202 }
199203
200- helper . runInTransaction ( agent , ( ) => {
204+ helper . runInTransaction ( agent , ( tx ) => {
201205 const segment = agent . tracer . getSegment ( )
202206 const embeddingEvent = new LlmEmbedding ( {
203207 agent,
204208 segment,
205209 request : req ,
206- response : res
210+ response : res ,
211+ transaction : tx
207212 } )
208213 assert . equal ( embeddingEvent [ 'response.usage.total_tokens' ] , 30 )
209214 end ( )
0 commit comments