test: Add unit tests for LLM API methods (#1878)

newrelic · Nov 17, 2023 · 3d8b987 · 3d8b987
1 parent 039705f
commit 3d8b987
Showing 1 changed file with 74 additions and 0 deletions.
diff --git a/test/unit/api/api-llm.test.js b/test/unit/api/api-llm.test.js
@@ -29,6 +29,7 @@ tap.test('Agent API LLM methods', (t) => {
     loggerMock.warn.reset()
     const agent = helper.loadMockedAgent()
     t.context.api = new API(agent)
+    t.context.api.agent.config.ai_monitoring.enabled = true
   })
 
   t.afterEach((t) => {
@@ -69,6 +70,26 @@ tap.test('Agent API LLM methods', (t) => {
     t.equal(loggerMock.warn.args[0][0], 'getLlmMessageIds invoked but ai_monitoring is disabled.')
   })
 
+  t.test('geLlmMessageIds is no-op when no transaction is available', async (t) => {
+    const { api } = t.context
+    const trackedIds = api.getLlmMessageIds({ responseId: 'test' })
+    t.equal(trackedIds, undefined)
+    t.equal(loggerMock.warn.callCount, 1)
+    t.equal(
+      loggerMock.warn.args[0][0],
+      'getLlmMessageIds must be called within the scope of a transaction.'
+    )
+  })
+
+  t.test('getLlmMessageIds returns undefined for unrecognized id', async (t) => {
+    const { api } = t.context
+    helper.runInTransaction(api.agent, () => {
+      const trackedIds = api.getLlmMessageIds({ responseId: 'test' })
+      t.equal(trackedIds, undefined)
+      t.equal(loggerMock.warn.callCount, 0)
+    })
+  })
+
   t.test('recordLlmFeedbackEvent is no-op when ai_monitoring is disabled', async (t) => {
     const { api } = t.context
     api.agent.config.ai_monitoring.enabled = false
@@ -85,4 +106,57 @@ tap.test('Agent API LLM methods', (t) => {
       'recordLlmFeedbackEvent invoked but ai_monitoring is disabled.'
     )
   })
+
+  t.test('recordLlmFeedbackEvent is no-op when no transaction is available', async (t) => {
+    const { api } = t.context
+
+    const result = api.recordLlmFeedbackEvent({
+      messageId: 'test',
+      category: 'test',
+      rating: 'test'
+    })
+    t.equal(result, undefined)
+    t.equal(loggerMock.warn.callCount, 1)
+    t.equal(
+      loggerMock.warn.args[0][0],
+      'No message feedback events will be recorded. recordLlmFeedbackEvent must be called within the scope of a transaction.'
+    )
+  })
+
+  t.test('recordLlmFeedbackEvent returns undefined on success', async (t) => {
+    const { api } = t.context
+
+    const rce = api.recordCustomEvent
+    let event
+    api.recordCustomEvent = (name, data) => {
+      event = { name, data }
+      return rce.call(api, name, data)
+    }
+    t.teardown(() => {
+      api.recordCustomEvent = rce
+    })
+
+    helper.runInTransaction(api.agent, () => {
+      const result = api.recordLlmFeedbackEvent({
+        messageId: 'test',
+        category: 'test-cat',
+        rating: '5 star',
+        metadata: { foo: 'foo' }
+      })
+      t.equal(result, undefined)
+      t.equal(loggerMock.warn.callCount, 0)
+      t.equal(event.name, 'LlmFeedbackMessage')
+      t.match(event.data, {
+        id: /[\w\d]{32}/,
+        conversation_id: '',
+        request_id: '',
+        message_id: 'test',
+        category: 'test-cat',
+        rating: '5 star',
+        message: '',
+        foo: 'foo',
+        ingest_source: 'Node'
+      })
+    })
+  })
 })