fix

Author: amychisholm03

test/versioned/langchain-aws/runnables-streaming.test.js

@@ -14,9 +14,7 @@ const { match } = require('../../lib/custom-assertions')
 const {
   runStreamingEnabledTests,
   runStreamingDisabledTest,
-  runAiMonitoringDisabledTests,
-  runBothDisabledTest,
-  runStreamingEdgeCasesTests
+  runAiMonitoringDisabledTests
 } = require('../langchain/runnables-streaming')
 const { FAKE_CREDENTIALS, getAiResponseServer } = require('../../lib/aws-server-stubs')
 const helper = require('../../lib/agent_helper')
@@ -112,22 +110,3 @@ test('ai_monitoring disabled', async (t) => {
     expectedContent: () => 'This is a test.'
   })(t)
 })
-
-test('both ai_monitoring and streaming disabled', async (t) => {
-  t.beforeEach((ctx) => beforeEach({ enabled: false, ctx }))
-  t.afterEach((ctx) => afterEach(ctx))
-
-  await runBothDisabledTest({
-    inputData: { topic: 'streamed' },
-    expectedContent: () => 'This is a test.'
-  })(t)
-})
-
-test('streaming enabled - edge cases', async (t) => {
-  t.beforeEach((ctx) => beforeEach({ enabled: true, ctx }))
-  t.afterEach((ctx) => afterEach(ctx))
-
-  await runStreamingEdgeCasesTests({
-    inputData: { topic: 'streamed' }
-  })(t)
-})

test/versioned/langchain-openai/runnables-streaming.test.js

@@ -13,9 +13,7 @@ const { match } = require('../../lib/custom-assertions')
 const {
   runStreamingEnabledTests,
   runStreamingDisabledTest,
-  runAiMonitoringDisabledTests,
-  runBothDisabledTest,
-  runStreamingEdgeCasesTests
+  runAiMonitoringDisabledTests
 } = require('../langchain/runnables-streaming')
 const createOpenAIMockServer = require('../openai/mock-server')
 const mockResponses = require('../openai/mock-chat-api-responses')
@@ -114,22 +112,3 @@ test('ai_monitoring disabled', async (t) => {
     expectedContent: () => mockResponses.get('Streamed response').streamData
   })(t)
 })
-
-test('both ai_monitoring and streaming disabled', async (t) => {
-  t.beforeEach((ctx) => beforeEach({ enabled: false, ctx }))
-  t.afterEach((ctx) => afterEach(ctx))
-
-  await runBothDisabledTest({
-    inputData: { topic: 'Streamed' },
-    expectedContent: () => mockResponses.get('Streamed response').streamData
-  })(t)
-})
-
-test('streaming enabled - edge cases', async (t) => {
-  t.beforeEach((ctx) => beforeEach({ enabled: true, ctx }))
-  t.afterEach((ctx) => afterEach(ctx))
-
-  await runStreamingEdgeCasesTests({
-    inputData: { topic: 'Streamed' }
-  })(t)
-})

test/versioned/langchain/runnables-streaming.js

@@ -541,6 +541,75 @@ function runStreamingEnabledTests(config) {
         end()
       })
     })
+
+    await t.test(
+      'should handle metadata properly during stream processing',
+      (t, end) => {
+        const { agent, prompt, model, outputParser } = t.nr
+
+        helper.runInTransaction(agent, async (tx) => {
+          const input = inputData
+          const options = {
+            metadata: { streamKey: 'streamValue', anotherKey: 'anotherValue' },
+            tags: ['stream-tag1', 'stream-tag2']
+          }
+
+          const chain = prompt.pipe(model).pipe(outputParser)
+          const stream = await chain.stream(input, options)
+          for await (const chunk of stream) {
+            consumeStreamChunk(chunk)
+          }
+
+          const events = agent.customEventAggregator.events.toArray()
+          const langchainEvents = filterLangchainEvents(events)
+          const langChainSummaryEvents = filterLangchainEventsByType(
+            langchainEvents,
+            'LlmChatCompletionSummary'
+          )
+
+          const [[, summary]] = langChainSummaryEvents
+          assert.equal(summary['metadata.streamKey'], 'streamValue')
+          assert.equal(summary['metadata.anotherKey'], 'anotherValue')
+
+          const tags = summary.tags.split(',')
+          assert.ok(tags.includes('stream-tag1'))
+          assert.ok(tags.includes('stream-tag2'))
+
+          tx.end()
+          end()
+        })
+      }
+    )
+
+    await t.test(
+      'should properly extend segment duration on each stream iteration',
+      (t, end) => {
+        const { agent, prompt, model, outputParser } = t.nr
+
+        helper.runInTransaction(agent, async (tx) => {
+          const input = inputData
+
+          const chain = prompt.pipe(model).pipe(outputParser)
+          const stream = await chain.stream(input)
+
+          const [segment] = tx.trace.getChildren(tx.trace.root.id)
+          assert.equal(segment.name, 'Llm/chain/Langchain/stream', 'should find the Langchain stream segment')
+
+          let chunkCount = 0
+          for await (const chunk of stream) {
+            consumeStreamChunk(chunk)
+            chunkCount++
+          }
+
+          // Segment should have been touched multiple times during streaming
+          assert.ok(chunkCount > 1, 'should have received multiple chunks')
+          assert.ok(segment.timer.hrDuration)
+
+          tx.end()
+          end()
+        })
+      }
+    )
   }
 }
 
@@ -601,55 +670,9 @@ function runStreamingDisabledTest(config) {
         })
       }
     )
-  }
-}
-
-/**
- * Runs the ai_monitoring disabled tests
- * @param {object} config Configuration for the test suite
- * @param {object} config.inputData The input data to pass to stream calls
- * @param {Function} [config.expectedContent] Function to get expected content
- */
-function runAiMonitoringDisabledTests(config) {
-  const { inputData, expectedContent } = config
-
-  function consumeStreamChunk() {
-    // A no-op function used to consume chunks of a stream.
-  }
 
-  return async (t) => {
     await t.test(
-      'should not create llm events when `ai_monitoring.enabled` is false',
-      (t, end) => {
-        const { agent, prompt, outputParser, model } = t.nr
-        agent.config.ai_monitoring.enabled = false
-
-        helper.runInTransaction(agent, async (tx) => {
-          const input = inputData
-
-          const chain = prompt.pipe(model).pipe(outputParser)
-          const stream = await chain.stream(input)
-          let content = ''
-          for await (const chunk of stream) {
-            content += chunk
-          }
-
-          if (expectedContent) {
-            const expected = expectedContent()
-            assert.equal(content, expected)
-          }
-
-          const events = agent.customEventAggregator.events.toArray()
-          assert.equal(events.length, 0, 'should not create llm events when ai_monitoring is disabled')
-
-          tx.end()
-          end()
-        })
-      }
-    )
-
-    await t.test(
-      'should not create segment when `ai_monitoring.enabled` is false',
+      'should not create segment when `ai_monitoring.streaming.enabled` is false',
       (t, end) => {
         const { agent, prompt, outputParser, model } = t.nr
         agent.config.ai_monitoring.enabled = false
@@ -675,17 +698,21 @@ function runAiMonitoringDisabledTests(config) {
 }
 
 /**
- * Runs the test for both ai_monitoring and streaming disabled
+ * Runs the ai_monitoring disabled tests
  * @param {object} config Configuration for the test suite
  * @param {object} config.inputData The input data to pass to stream calls
  * @param {Function} [config.expectedContent] Function to get expected content
  */
-function runBothDisabledTest(config) {
+function runAiMonitoringDisabledTests(config) {
   const { inputData, expectedContent } = config
 
+  function consumeStreamChunk() {
+    // A no-op function used to consume chunks of a stream.
+  }
+
   return async (t) => {
     await t.test(
-      'should not create llm events when both `ai_monitoring.enabled` and `ai_monitoring.streaming.enabled` are false',
+      'should not create llm events when `ai_monitoring.enabled` is false',
       (t, end) => {
         const { agent, prompt, outputParser, model } = t.nr
         agent.config.ai_monitoring.enabled = false
@@ -706,61 +733,7 @@ function runBothDisabledTest(config) {
           }
 
           const events = agent.customEventAggregator.events.toArray()
-          assert.equal(events.length, 0, 'should not create llm events when both configs are disabled')
-
-          tx.end()
-          end()
-        })
-      }
-    )
-  }
-}
-
-/**
- * Runs the streaming enabled edge cases tests
- * @param {object} config Configuration for the test suite
- * @param {object} config.inputData The input data to pass to stream calls
- */
-function runStreamingEdgeCasesTests(config) {
-  const { inputData } = config
-
-  function consumeStreamChunk() {
-    // A no-op function used to consume chunks of a stream.
-  }
-
-  return async (t) => {
-    await t.test(
-      'should handle metadata properly during stream processing',
-      (t, end) => {
-        const { agent, prompt, model, outputParser } = t.nr
-
-        helper.runInTransaction(agent, async (tx) => {
-          const input = inputData
-          const options = {
-            metadata: { streamKey: 'streamValue', anotherKey: 'anotherValue' },
-            tags: ['stream-tag1', 'stream-tag2']
-          }
-
-          const chain = prompt.pipe(model).pipe(outputParser)
-          const stream = await chain.stream(input, options)
-          for await (const chunk of stream) {
-            consumeStreamChunk(chunk)
-          }
-
-          const events = agent.customEventAggregator.events.toArray()
-          const langchainEvents = filterLangchainEvents(events)
-          const langChainSummaryEvents = filterLangchainEventsByType(
-            langchainEvents,
-            'LlmChatCompletionSummary'
-          )
-
-          const [[, summary]] = langChainSummaryEvents
-          assert.equal(summary['metadata.streamKey'], 'streamValue')
-          assert.equal(summary['metadata.anotherKey'], 'anotherValue')
-
-          const tags = summary.tags.split(',')
-          assert.ok(tags.includes('stream-tag1'))
-          assert.ok(tags.includes('stream-tag2'))
+          assert.equal(events.length, 0, 'should not create llm events when ai_monitoring is disabled')
 
           tx.end()
           end()
@@ -769,28 +742,22 @@ function runStreamingEdgeCasesTests(config) {
     )
 
     await t.test(
-      'should properly extend segment duration on each stream iteration',
+      'should not create segment when `ai_monitoring.enabled` is false',
       (t, end) => {
-        const { agent, prompt, model, outputParser } = t.nr
+        const { agent, prompt, outputParser, model } = t.nr
+        agent.config.ai_monitoring.enabled = false
 
         helper.runInTransaction(agent, async (tx) => {
           const input = inputData
 
           const chain = prompt.pipe(model).pipe(outputParser)
           const stream = await chain.stream(input)
-
-          const [segment] = tx.trace.getChildren(tx.trace.root.id)
-          assert.equal(segment.name, 'Llm/chain/Langchain/stream', 'should find the Langchain stream segment')
-
-          let chunkCount = 0
           for await (const chunk of stream) {
             consumeStreamChunk(chunk)
-            chunkCount++
           }
 
-          // Segment should have been touched multiple times during streaming
-          assert.ok(chunkCount > 1, 'should have received multiple chunks')
-          assert.ok(segment.timer.hrDuration)
+          const segment = findSegment(tx.trace, tx.trace.root, 'Llm/chain/Langchain/stream')
+          assert.equal(segment, undefined, 'should not create Llm/chain/Langchain/stream segment when ai_monitoring is disabled')
 
           tx.end()
           end()
@@ -803,7 +770,5 @@ function runStreamingEdgeCasesTests(config) {
 module.exports = {
   runStreamingEnabledTests,
   runStreamingDisabledTest,
-  runAiMonitoringDisabledTests,
-  runBothDisabledTest,
-  runStreamingEdgeCasesTests
+  runAiMonitoringDisabledTests
 }