From e6562f0b46dd914f5e341239583f87b4c4f9f3c4 Mon Sep 17 00:00:00 2001
From: Sam Brenner <sam.brenner@datadoghq.com>
Date: Thu, 24 Oct 2024 11:13:38 -0400
Subject: [PATCH] add some tests for not auto-annotating in certain cases

---
 .../dd-trace/test/llmobs/sdk/index.spec.js    | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/packages/dd-trace/test/llmobs/sdk/index.spec.js b/packages/dd-trace/test/llmobs/sdk/index.spec.js
index dccb06f6f4..5c48e7bb88 100644
--- a/packages/dd-trace/test/llmobs/sdk/index.spec.js
+++ b/packages/dd-trace/test/llmobs/sdk/index.spec.js
@@ -378,6 +378,62 @@ describe('sdk', () => {
 
           expect(span.finish).to.have.been.called
         })
+
+        it('does not auto-annotate llm spans', () => {
+          let span
+          function myLLM (input) {
+            span = llmobs._active()
+            return ''
+          }
+
+          const wrappedMyLLM = llmobs.wrap({ kind: 'llm' }, myLLM)
+
+          wrappedMyLLM('input')
+
+          expect(LLMObsTagger.tagMap.get(span)).to.deep.equal({
+            '_ml_obs.meta.span.kind': 'llm',
+            '_ml_obs.meta.ml_app': 'mlApp',
+            '_ml_obs.llmobs_parent_id': 'undefined'
+          })
+        })
+
+        it('does not auto-annotate embedding spans input', () => {
+          let span
+          function myEmbedding (input) {
+            span = llmobs._active()
+            return 'output'
+          }
+
+          const wrappedMyEmbedding = llmobs.wrap({ kind: 'embedding' }, myEmbedding)
+
+          wrappedMyEmbedding('input')
+
+          expect(LLMObsTagger.tagMap.get(span)).to.deep.equal({
+            '_ml_obs.meta.span.kind': 'embedding',
+            '_ml_obs.meta.ml_app': 'mlApp',
+            '_ml_obs.llmobs_parent_id': 'undefined',
+            '_ml_obs.meta.output.value': 'output'
+          })
+        })
+
+        it('does not auto-annotate retrieval spans output', () => {
+          let span
+          function myRetrieval (input) {
+            span = llmobs._active()
+            return 'output'
+          }
+
+          const wrappedMyRetrieval = llmobs.wrap({ kind: 'retrieval' }, myRetrieval)
+
+          wrappedMyRetrieval('input')
+
+          expect(LLMObsTagger.tagMap.get(span)).to.deep.equal({
+            '_ml_obs.meta.span.kind': 'retrieval',
+            '_ml_obs.meta.ml_app': 'mlApp',
+            '_ml_obs.llmobs_parent_id': 'undefined',
+            '_ml_obs.meta.input.value': 'input'
+          })
+        })
       })
 
       describe('parentage', () => {